<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strings"
	"syscall"
	"time"
	//"user"

	"pty"
)

type cgroupname string

func (c cgroupname) apply(s string, f func(s string)) {
	// range of strings.Split("",",") is 1.
	// not exactly what we might expect.
	if s == "" {
		return
	}
	for _, g := range strings.Split(s, ",") {
		p := path.Join(g)
		f(p)
	}
}

func (c cgroupname) Validate(s string) {
	c.apply(s, func(s string) {
		if st, err := os.Stat(path.Join(string(c), s)); err != nil {
			log.Fatal("%v", err)
		} else if !st.IsDir() {
			log.Fatal("%s: not a directory", s)
		}
	})
}

func (c cgroupname) Create(s, name string) {
	if err := os.MkdirAll(path.Join(string(c), s, name), 0755); err != nil {
		log.Fatal(err)
	}
}

func (c cgroupname) Attach(s, name string, pid int) {
	t := path.Join(string(c), s, name, "tasks")
	b := []byte(fmt.Sprintf("%v", pid))
	if err := ioutil.WriteFile(t, b, 0600); err != nil {
		log.Fatal(err)
	}
}

func (c cgroupname) Destroy(s, n string) {
	if err := os.RemoveAll(path.Join(string(c), s, n)); err != nil {
		log.Fatal(err)
	}
}

func (c cgroupname) Do(groups string, pid int) {
	cgn := fmt.Sprintf("pflask.%d", pid)
	c.apply(groups, func(s string) {
		c.Create(s, cgn)
		c.Attach(s, cgn, pid)
	})
}

type mount struct {
	src, dst, mtype, opts string
	flags                 uintptr
	dir                   bool
}

// Add adds  to teh mouuntlist. Don't know if we need it, but we might for additional volumes?
func Add(src, dst, mtype, opts string, flags uintptr, dir bool) {
	mounts = append(mounts, mount{src: src, dst: dst, mtype: mtype, flags: flags, opts: opts, dir: dir})

}

// One mounts one mountpoint, using base as a prefix for the destination.
// If anything goes wrong, we just bail out; we've privatized the namespace
// so there is nothing more to do.
func (m *mount) One(base string) {
	dst := path.Join(base, m.dst)
	if m.dir {
		if err := os.MkdirAll(dst, 0755); err != nil {
			log.Fatalf("One: mkdirall %v: %v", m.dst, err)
		}
	}
	if err := syscall.Mount(m.src, dst, m.mtype, m.flags, m.opts); err != nil {
		log.Fatalf("Mount :%s: on :%s: type :%s: flags %x: opts :%v: %v\n",
			m.src, m.dst, m.mtype, m.flags, m.opts, err)
	}
}
func MountAll(base string) {
	root.One("")
	for _, m := range mounts {
		m.One(base)
	}
}

func make_console(base, console string) {
	if err := os.Chmod(console, 0600); err != nil {
		log.Printf("%v", err)
	}
	if err := os.Chown(console, 0, 0); err != nil {
		log.Printf("%v", err)
	}

	st, err := os.Stat(console)
	if err != nil {
		log.Printf("%v", err)
	}

	nn := path.Join(base, "/dev/console")
	if err := syscall.Mknod(nn, uint32(st.Mode()), int(st.Sys().(*syscall.Stat_t).Dev)); err != nil {
		log.Printf("%v", err)
	}

	// if any previous steps failed, this one will too, so we can bail here.
	if err := syscall.Mount(console, nn, "", syscall.MS_BIND, ""); err != nil {
		log.Fatalf("Mount :%s: on :%s: flags %v: %v",
			console, nn, syscall.MS_BIND, err)
	}

}
func copy_nodes(base string) {
	nodes := []string{
		"/dev/tty",
		"/dev/full",
		"/dev/null",
		"/dev/zero",
		"/dev/random",
		"/dev/urandom"}

	for _, n := range nodes {
		st, err := os.Stat(n)
		if err != nil {
			log.Printf("%v", err)
		}
		nn := path.Join(base, n)
		if err := syscall.Mknod(nn, uint32(st.Mode()), int(st.Sys().(*syscall.Stat_t).Dev)); err != nil {
			log.Printf("%v", err)
		}

	}
}

func make_ptmx(base string) {
	dst := path.Join(base, "/dev/ptmx")

	if err := os.Symlink("/dev/ptmx", dst); err != nil {
		log.Printf("%v", err)
	}
}

func make_symlinks(base string) {
	linkit := []struct {
		src, dst string
	}{
		{"/dev/pts/ptmx", "/dev/ptmx"},
		{"/proc/kcore", "/dev/core"},
		{"/proc/self/fd", "/dev/fd"},
		{"/proc/self/fd/0", "/dev/stdin"},
		{"/proc/self/fd/1", "/dev/stdout"},
		{"/proc/self/fd/2", "/dev/stderr"},
	}

	for i := range linkit {
		dst := path.Join(base, linkit[i].dst)

		if err := os.Symlink(linkit[i].src, dst); err != nil {
			log.Printf("%v", err)
		}
	}
}

var (
	cgpath  = flag.String("cgpath", "/sys/fs/cgroup", "set the cgroups")
	cgroup  = flag.String("cgroup", "", "set the cgroups")
	mnt     = flag.String("mount", "", "define mounts")
	chroot  = flag.String("chroot", "", "where to chroot to")
	chdir   = flag.String("chdir", "", "where to chrdir to in the chroot")
	console = flag.String("console", "/dev/console", "where the root is")
	keepenv = flag.Bool("keepenv", false, "Keep the environment")
	env     = flag.String("env", "", "other environment variables")
	user    = flag.String("user", "root" /*user.User.Username*/, "User name")
	root = &mount {"", "/", "", "", syscall.MS_SLAVE|syscall.MS_REC, false}
	mounts = []mount {
		{"proc", "/proc", "proc", "", syscall.MS_NOSUID|syscall.MS_NOEXEC|syscall.MS_NODEV, true},
		{"/proc/sys", "/proc/sys", "", "", syscall.MS_BIND, true},
		{"", "/proc/sys", "", "", syscall.MS_BIND|syscall.MS_RDONLY|syscall.MS_REMOUNT, true},
		{"sysfs", "/sys", "sysfs", "", syscall.MS_NOSUID|syscall.MS_NOEXEC|syscall.MS_NODEV|syscall.MS_RDONLY, true},
		{"tmpfs", "/dev", "tmpfs", "mode=755", syscall.MS_NOSUID|syscall.MS_STRICTATIME, true},
		{"devpts", "/dev/pts", "devpts", "newinstance,ptmxmode=000,mode=620,gid=5", syscall.MS_NOSUID|syscall.MS_NOEXEC, true},
		{"tmpfs", "/dev/shm", "tmpfs", "mode=1777", syscall.MS_NOSUID|syscall.MS_STRICTATIME|syscall.MS_NODEV, true},
		{"tmpfs", "/run", "tmpfs", "mode=755", syscall.MS_NOSUID|syscall.MS_NODEV|syscall.MS_STRICTATIME, true},
	}

)

func main() {
	flag.Parse()

	if len(flag.Args()) < 1 {
		os.Exit(1)
	}

	// note the unshare system call worketh not for Go.
	// So do it ourselves.
	a := os.Args
	if a[len(a)-1] != "#" {

		a = append(a, "#")
		// spawn ourselves with the right unsharing settings.
		c := exec.Command(a[0], a[1:]...)
		c.SysProcAttr = &syscall.SysProcAttr{Cloneflags: syscall.CLONE_NEWNS | syscall.CLONE_NEWUTS | syscall.CLONE_NEWIPC | syscall.CLONE_NEWPID}
		//		c.SysProcAttr.Cloneflags |= syscall.CLONE_NEWNET

		c.Stdin = os.Stdin
		c.Stdout = os.Stdout
		c.Stderr = os.Stderr
		t, err := getTermios(1)
		if err != nil {
			log.Fatalf("Can't get termios on fd 1: %v", err)
		}
		if err := c.Run(); err != nil {
			log.Printf(err.Error())
		}
		if err := t.set(1); err != nil {
			log.Printf("Can't reset termios on fd1: %v", err)
		}
		os.Exit(1)
	}

	if *chroot == "" {
		log.Fatalf("you are required to set the chroot via --chroot")
	}

	a = flag.Args()
	//log.Printf("greetings %v\n", a)
	a = a[:len(a)-1]

	ptm, pts, sname, err := pty.Open()
	if err != nil {
		log.Fatalf(err.Error())
	}

	// child code. Not really. What really happens here is we set
	// ourselves into the container, and spawn the child. It's a bit odd
	// but we're the master, but we'll run in the container? I don't know
	// how else to do it. This may require we set some things up first,
	// esp. the network. But, it's all fun and games until someone loses
	// an eye.
	MountAll(*chroot)

	copy_nodes(*chroot)

	make_ptmx(*chroot)

	make_symlinks(*chroot)

	make_console(*chroot, sname)

	//umask(0022);

	/* TODO: drop capabilities */

	//do_user(user);

	e := make(map[string]string)
	if *keepenv {
		for _, v := range os.Environ() {
			k := strings.SplitN(v, "=", 2)
			e[k[0]] = k[1]
		}
	}

	term := os.Getenv("TERM")
	e["TERM"] = term
	e["PATH"] = "/usr/sbin:/usr/bin:/sbin:/bin"
	e["USER"] = *user
	e["LOGNAME"] = *user
	e["HOME"] = "/root"

	if *env != "" {
		for _, c := range strings.Split(*env, ",") {
			k := strings.SplitN(c, "=", 2)
			if len(k) != 2 {
				log.Printf("Bogus environment string %v", c)
				continue
			}
			e[k[0]] = k[1]
		}
	}
	e["container"] = "pflask"

	if len(a) == 0 {
		a = []string{"/bin/bash", "bash"}
	}
	c := exec.Command(a[0]) // , a[1:]...)
	c.Env = nil
	for k, v := range e {
		c.Env = append(c.Env, k+"="+v)
	}

	c.SysProcAttr = &syscall.SysProcAttr{
		Chroot:  *chroot,
		Setctty: true,
		Setsid:  true,
	}
	c.Stdout = pts
	c.Stdin = pts
	c.Stderr = c.Stdout
	c.SysProcAttr.Setctty = true
	c.SysProcAttr.Setsid = true
	c.SysProcAttr.Cloneflags = 0
	c.SysProcAttr.Ptrace = true
	c.Dir = "/"
	err = c.Start()
	if err != nil {
		panic(err)
	}
	kid := c.Process.Pid
	fmt.Printf("Started %d\n", kid)

	// set up the containers, then resume the process.
	// Its children will get the containers as it clones.

	cg := cgroupname(*cgpath)
	cg.Do(*cgroup, kid)

	// sometimes the detach fails. Looks like a race condition: we're
	// sending the detach before the child has hit the TRACE_ME point.
	// Experimentally, when it fails, even one seconds it too short to
	// sleep. Sleep for 5 seconds.

	for {
		if err = syscall.PtraceDetach(kid); err != nil {
			log.Printf("Could not detach %v, sleeping five seconds", kid)
			time.Sleep(5*time.Second)
			continue
		}
		break
	}

	// we don't set raw until the very last, so if they see an issue they can hit ^C
	t, err := getTermios(1)
	if err != nil {
		log.Fatalf(err.Error())
	}
	if err = t.setRaw(1); err != nil {
		log.Fatalf(err.Error())
	}

	go func() {
		io.Copy(os.Stdout, ptm)
		os.Exit(1)
	}()
	io.Copy(ptm, os.Stdin)
}
<commit_msg>Set the chdir for the child process.<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strings"
	"syscall"
	"time"
	//"user"

	"pty"
)

type cgroupname string

func (c cgroupname) apply(s string, f func(s string)) {
	// range of strings.Split("",",") is 1.
	// not exactly what we might expect.
	if s == "" {
		return
	}
	for _, g := range strings.Split(s, ",") {
		p := path.Join(g)
		f(p)
	}
}

func (c cgroupname) Validate(s string) {
	c.apply(s, func(s string) {
		if st, err := os.Stat(path.Join(string(c), s)); err != nil {
			log.Fatal("%v", err)
		} else if !st.IsDir() {
			log.Fatal("%s: not a directory", s)
		}
	})
}

func (c cgroupname) Create(s, name string) {
	if err := os.MkdirAll(path.Join(string(c), s, name), 0755); err != nil {
		log.Fatal(err)
	}
}

func (c cgroupname) Attach(s, name string, pid int) {
	t := path.Join(string(c), s, name, "tasks")
	b := []byte(fmt.Sprintf("%v", pid))
	if err := ioutil.WriteFile(t, b, 0600); err != nil {
		log.Fatal(err)
	}
}

func (c cgroupname) Destroy(s, n string) {
	if err := os.RemoveAll(path.Join(string(c), s, n)); err != nil {
		log.Fatal(err)
	}
}

func (c cgroupname) Do(groups string, pid int) {
	cgn := fmt.Sprintf("pflask.%d", pid)
	c.apply(groups, func(s string) {
		c.Create(s, cgn)
		c.Attach(s, cgn, pid)
	})
}

type mount struct {
	src, dst, mtype, opts string
	flags                 uintptr
	dir                   bool
}

// Add adds  to teh mouuntlist. Don't know if we need it, but we might for additional volumes?
func Add(src, dst, mtype, opts string, flags uintptr, dir bool) {
	mounts = append(mounts, mount{src: src, dst: dst, mtype: mtype, flags: flags, opts: opts, dir: dir})

}

// One mounts one mountpoint, using base as a prefix for the destination.
// If anything goes wrong, we just bail out; we've privatized the namespace
// so there is nothing more to do.
func (m *mount) One(base string) {
	dst := path.Join(base, m.dst)
	if m.dir {
		if err := os.MkdirAll(dst, 0755); err != nil {
			log.Fatalf("One: mkdirall %v: %v", m.dst, err)
		}
	}
	if err := syscall.Mount(m.src, dst, m.mtype, m.flags, m.opts); err != nil {
		log.Fatalf("Mount :%s: on :%s: type :%s: flags %x: opts :%v: %v\n",
			m.src, m.dst, m.mtype, m.flags, m.opts, err)
	}
}
func MountAll(base string) {
	root.One("")
	for _, m := range mounts {
		m.One(base)
	}
}

func make_console(base, console string) {
	if err := os.Chmod(console, 0600); err != nil {
		log.Printf("%v", err)
	}
	if err := os.Chown(console, 0, 0); err != nil {
		log.Printf("%v", err)
	}

	st, err := os.Stat(console)
	if err != nil {
		log.Printf("%v", err)
	}

	nn := path.Join(base, "/dev/console")
	if err := syscall.Mknod(nn, uint32(st.Mode()), int(st.Sys().(*syscall.Stat_t).Dev)); err != nil {
		log.Printf("%v", err)
	}

	// if any previous steps failed, this one will too, so we can bail here.
	if err := syscall.Mount(console, nn, "", syscall.MS_BIND, ""); err != nil {
		log.Fatalf("Mount :%s: on :%s: flags %v: %v",
			console, nn, syscall.MS_BIND, err)
	}

}
func copy_nodes(base string) {
	nodes := []string{
		"/dev/tty",
		"/dev/full",
		"/dev/null",
		"/dev/zero",
		"/dev/random",
		"/dev/urandom"}

	for _, n := range nodes {
		st, err := os.Stat(n)
		if err != nil {
			log.Printf("%v", err)
		}
		nn := path.Join(base, n)
		if err := syscall.Mknod(nn, uint32(st.Mode()), int(st.Sys().(*syscall.Stat_t).Dev)); err != nil {
			log.Printf("%v", err)
		}

	}
}

func make_ptmx(base string) {
	dst := path.Join(base, "/dev/ptmx")

	if err := os.Symlink("/dev/ptmx", dst); err != nil {
		log.Printf("%v", err)
	}
}

func make_symlinks(base string) {
	linkit := []struct {
		src, dst string
	}{
		{"/dev/pts/ptmx", "/dev/ptmx"},
		{"/proc/kcore", "/dev/core"},
		{"/proc/self/fd", "/dev/fd"},
		{"/proc/self/fd/0", "/dev/stdin"},
		{"/proc/self/fd/1", "/dev/stdout"},
		{"/proc/self/fd/2", "/dev/stderr"},
	}

	for i := range linkit {
		dst := path.Join(base, linkit[i].dst)

		if err := os.Symlink(linkit[i].src, dst); err != nil {
			log.Printf("%v", err)
		}
	}
}

var (
	cgpath  = flag.String("cgpath", "/sys/fs/cgroup", "set the cgroups")
	cgroup  = flag.String("cgroup", "", "set the cgroups")
	mnt     = flag.String("mount", "", "define mounts")
	chroot  = flag.String("chroot", "", "where to chroot to")
	chdir   = flag.String("chdir", "/", "where to chrdir to in the chroot")
	console = flag.String("console", "/dev/console", "where the root is")
	keepenv = flag.Bool("keepenv", false, "Keep the environment")
	env     = flag.String("env", "", "other environment variables")
	user    = flag.String("user", "root" /*user.User.Username*/, "User name")
	root = &mount {"", "/", "", "", syscall.MS_SLAVE|syscall.MS_REC, false}
	mounts = []mount {
		{"proc", "/proc", "proc", "", syscall.MS_NOSUID|syscall.MS_NOEXEC|syscall.MS_NODEV, true},
		{"/proc/sys", "/proc/sys", "", "", syscall.MS_BIND, true},
		{"", "/proc/sys", "", "", syscall.MS_BIND|syscall.MS_RDONLY|syscall.MS_REMOUNT, true},
		{"sysfs", "/sys", "sysfs", "", syscall.MS_NOSUID|syscall.MS_NOEXEC|syscall.MS_NODEV|syscall.MS_RDONLY, true},
		{"tmpfs", "/dev", "tmpfs", "mode=755", syscall.MS_NOSUID|syscall.MS_STRICTATIME, true},
		{"devpts", "/dev/pts", "devpts", "newinstance,ptmxmode=000,mode=620,gid=5", syscall.MS_NOSUID|syscall.MS_NOEXEC, true},
		{"tmpfs", "/dev/shm", "tmpfs", "mode=1777", syscall.MS_NOSUID|syscall.MS_STRICTATIME|syscall.MS_NODEV, true},
		{"tmpfs", "/run", "tmpfs", "mode=755", syscall.MS_NOSUID|syscall.MS_NODEV|syscall.MS_STRICTATIME, true},
	}

)

func main() {
	flag.Parse()

	if len(flag.Args()) < 1 {
		os.Exit(1)
	}

	// note the unshare system call worketh not for Go.
	// So do it ourselves.
	a := os.Args
	if a[len(a)-1] != "#" {

		a = append(a, "#")
		// spawn ourselves with the right unsharing settings.
		c := exec.Command(a[0], a[1:]...)
		c.SysProcAttr = &syscall.SysProcAttr{Cloneflags: syscall.CLONE_NEWNS | syscall.CLONE_NEWUTS | syscall.CLONE_NEWIPC | syscall.CLONE_NEWPID}
		//		c.SysProcAttr.Cloneflags |= syscall.CLONE_NEWNET

		c.Stdin = os.Stdin
		c.Stdout = os.Stdout
		c.Stderr = os.Stderr
		t, err := getTermios(1)
		if err != nil {
			log.Fatalf("Can't get termios on fd 1: %v", err)
		}
		if err := c.Run(); err != nil {
			log.Printf(err.Error())
		}
		if err := t.set(1); err != nil {
			log.Printf("Can't reset termios on fd1: %v", err)
		}
		os.Exit(1)
	}

	if *chroot == "" {
		log.Fatalf("you are required to set the chroot via --chroot")
	}

	a = flag.Args()
	//log.Printf("greetings %v\n", a)
	a = a[:len(a)-1]

	ptm, pts, sname, err := pty.Open()
	if err != nil {
		log.Fatalf(err.Error())
	}

	// child code. Not really. What really happens here is we set
	// ourselves into the container, and spawn the child. It's a bit odd
	// but we're the master, but we'll run in the container? I don't know
	// how else to do it. This may require we set some things up first,
	// esp. the network. But, it's all fun and games until someone loses
	// an eye.
	MountAll(*chroot)

	copy_nodes(*chroot)

	make_ptmx(*chroot)

	make_symlinks(*chroot)

	make_console(*chroot, sname)

	//umask(0022);

	/* TODO: drop capabilities */

	//do_user(user);

	e := make(map[string]string)
	if *keepenv {
		for _, v := range os.Environ() {
			k := strings.SplitN(v, "=", 2)
			e[k[0]] = k[1]
		}
	}

	term := os.Getenv("TERM")
	e["TERM"] = term
	e["PATH"] = "/usr/sbin:/usr/bin:/sbin:/bin"
	e["USER"] = *user
	e["LOGNAME"] = *user
	e["HOME"] = "/root"

	if *env != "" {
		for _, c := range strings.Split(*env, ",") {
			k := strings.SplitN(c, "=", 2)
			if len(k) != 2 {
				log.Printf("Bogus environment string %v", c)
				continue
			}
			e[k[0]] = k[1]
		}
	}
	e["container"] = "pflask"

	if len(a) == 0 {
		a = []string{"/bin/bash", "bash"}
	}
	c := exec.Command(a[0]) // , a[1:]...)
	c.Env = nil
	for k, v := range e {
		c.Env = append(c.Env, k+"="+v)
	}

	c.SysProcAttr = &syscall.SysProcAttr{
		Chroot:  *chroot,
		Setctty: true,
		Setsid:  true,
	}
	c.Stdout = pts
	c.Stdin = pts
	c.Stderr = c.Stdout
	c.SysProcAttr.Setctty = true
	c.SysProcAttr.Setsid = true
	c.SysProcAttr.Cloneflags = 0
	c.SysProcAttr.Ptrace = true
	c.Dir = *chdir
	err = c.Start()
	if err != nil {
		panic(err)
	}
	kid := c.Process.Pid
	fmt.Printf("Started %d\n", kid)

	// set up the containers, then resume the process.
	// Its children will get the containers as it clones.

	cg := cgroupname(*cgpath)
	cg.Do(*cgroup, kid)

	// sometimes the detach fails. Looks like a race condition: we're
	// sending the detach before the child has hit the TRACE_ME point.
	// Experimentally, when it fails, even one seconds it too short to
	// sleep. Sleep for 5 seconds.

	for {
		if err = syscall.PtraceDetach(kid); err != nil {
			log.Printf("Could not detach %v, sleeping five seconds", kid)
			time.Sleep(5*time.Second)
			continue
		}
		break
	}

	// we don't set raw until the very last, so if they see an issue they can hit ^C
	t, err := getTermios(1)
	if err != nil {
		log.Fatalf(err.Error())
	}
	if err = t.setRaw(1); err != nil {
		log.Fatalf(err.Error())
	}

	go func() {
		io.Copy(os.Stdout, ptm)
		os.Exit(1)
	}()
	io.Copy(ptm, os.Stdin)
}
<|endoftext|>


<commit_before>package framework

import (
	"net/http"
)

type Server interface {
	SetHandler(http.Handler)
	SetLogger(LogWriter)
	SetAddr(string)
	ListenAndServe() error
	ListenAndServeTLS(string, string) error
}

type Error interface {
	error
	Code() int
	Type() string
	Name() string
}

type Config interface {
	Get(string) Value
}

type Valuer interface {
	Values() Values
	SetValue(interface{}, interface{})
	GetValue(interface{}) Value
	Value(interface{}) interface{}
}

type ParamValuer interface {
	Param(string) ParamValue
	ParamOk(string) (ParamValue, bool)
	Params(...string) ParamValues
}

type Context interface {
	LogWriter
	Request
	Response
	Valuer
	Operation() Operation
	Locale() Locale
	Id() int64
	User() User
}

type Request interface {
	Sessioner
	ParamValuer
	Path() string
	Method() string
}

type Response interface {
	http.ResponseWriter
	IsIndented() bool
}

type Operations interface {
	Append(...Operation)
	Set(...Operation)
	Get() []Operation
	GetByIndex(int) Operation
	GetByName(string) Operation
	Len() int
}

type ModelStore interface {
	Get()
	Create()
	Edit()
	Remove()
}

type Operation interface {
	BuildPather
	Route() Route
	GetName() string
	GetMethods() []string
	GetPath() string
	GetFullPath() string
	// Param(string) Param
	// Params(...string) Params
	// BodyParam(string) Param
	// BodyParams(...string) Params
	// HeaderParam(string) Param
	// HeaderParams(...string) Params
	// PathParam(string) Param
	// PathParams(...string) Params
	// QueryParam(string) Param
	// QueryParams(...string) Params
}

type Resource interface {
	Create(Context) interface{}
	Retrieve(Context) interface{}
	Update(Context) interface{}
	Delete(Context) interface{}
}

type Values map[interface{}]interface{}

func (vs *Values) Set(k interface{}, v interface{}) {
	(*vs)[k] = v
}

func (vs *Values) Get(k interface{}) Value {
	return &value{(*vs)[k]}
}

type Value interface {
	Value() interface{}
	Int() int
	Int64() int64
	Float() float32
	Float64() float64
	String() string
	Bool() bool
}

type Params []Param

type Param interface {
}

type ParamValues map[string]ParamValue

type ParamValue interface {
	Value
	Name() string
	As() ValueType
	RawString() string
	IntE() (int, error)
	Int64E() (int64, error)
	FloatE() (float32, error)
	Float64E() (float64, error)
	BoolE() (bool, error)
}

type BuildPather interface {
	BuildPath(...interface{}) string
}

type Route interface {
	BuildPather
	Path() string
	FullPath() string
	SetPath(string)
}
<commit_msg>Edited framework interfaces<commit_after>package framework

import (
	"net/http"
)

type Server interface {
	SetHandler(http.Handler)
	SetLogger(LogWriter)
	SetAddr(string)
	ListenAndServe() error
	ListenAndServeTLS(string, string) error
}

type Error interface {
	error
	Code() int
	Type() string
	Name() string
}

type Config interface {
	Get(string) Value
}

type Valuer interface {
	Values() Values
	SetValue(interface{}, interface{})
	GetValue(interface{}) Value
	Value(interface{}) interface{}
}

type ParamValuer interface {
	Param(string) ParamValue
	ParamOk(string) (ParamValue, bool)
	Params(...string) ParamValues
}

type Context interface {
	LogWriter
	Request
	Response
	Valuer
	Operation() Operation
	Locale() Locale
	Id() int64
	User() User
}

type Request interface {
	Sessioner
	ParamValuer
	Path() string
	Method() string
}

type Response interface {
	http.ResponseWriter
	IsIndented() bool
}

type Operations interface {
	Append(...Operation)
	Set(...Operation)
	Get() []Operation
	GetByIndex(int) Operation
	GetByName(string) Operation
	Len() int
}

type ModelStore interface {
	Get()
	Create()
	Edit()
	Remove()
}

type Operation interface {
	BuildPather
	Route() Route
	GetName() string
	GetMethods() []string
	GetPath() string
	GetFullPath() string
	GetSchemes() []string
	GetAllSchemes() []string
	GetAccepts() []string
	GetAllAccepts() []string
	GetReturns() []string
	GetAllReturns() []string
	// Param(string) Param
	// Params(...string) Params
	// BodyParam(string) Param
	// BodyParams(...string) Params
	// HeaderParam(string) Param
	// HeaderParams(...string) Params
	// PathParam(string) Param
	// PathParams(...string) Params
	// QueryParam(string) Param
	// QueryParams(...string) Params
}

type Resource interface {
	Create(Context) interface{}
	Retrieve(Context) interface{}
	Update(Context) interface{}
	Delete(Context) interface{}
}

type Values map[interface{}]interface{}

func (vs *Values) Set(k interface{}, v interface{}) {
	(*vs)[k] = v
}

func (vs *Values) Get(k interface{}) Value {
	return &value{(*vs)[k]}
}

type Value interface {
	Value() interface{}
	Int() int
	Int64() int64
	Float() float32
	Float64() float64
	String() string
	Bool() bool
}

type Params interface {
	ContainsFiles() bool
	ContainsBodyParams() bool
}

type Param interface {
}

type ParamValues map[string]ParamValue

type ParamValue interface {
	Value
	Name() string
	As() ValueType
	RawString() string
	IntE() (int, error)
	Int64E() (int64, error)
	FloatE() (float32, error)
	Float64E() (float64, error)
	BoolE() (bool, error)
}

type BuildPather interface {
	BuildPath(...interface{}) string
}

type Route interface {
	BuildPather
	GetPath() string
	GetFullPath() string
	GetSchemes() []string
	GetAllSchemes() []string
	GetAccepts() []string
	GetAllAccepts() []string
	GetReturns() []string
	GetAllReturns() []string
}
<|endoftext|>


<commit_before>// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see <http://www.gnu.org/licenses/>.

package twitter

import (
	"fmt"
	"github.com/ChimeraCoder/anaconda"
	"github.com/nmeum/marvin/irc"
	"github.com/nmeum/marvin/modules"
	"html"
	"net/url"
	"strconv"
	"strings"
)

type Module struct {
	api               *anaconda.TwitterApi
	ReadOnly          bool   `json:"read_only"`
	ConsumerKey       string `json:"consumer_key"`
	ConsumerSecret    string `json:"consumer_secret"`
	AccessToken       string `json:"access_token"`
	AccessTokenSecret string `json:"access_token_secret"`
}

func Init(moduleSet *modules.ModuleSet) {
	moduleSet.Register(new(Module))
}

func (m *Module) Name() string {
	return "twitter"
}

func (m *Module) Help() string {
	return "USAGE: !tweet TEXT || !reply ID TEXT || !retweet ID || !favorite ID"
}

func (m *Module) Defaults() {
	m.ReadOnly = false
}

func (m *Module) Load(client *irc.Client) error {
	anaconda.SetConsumerKey(m.ConsumerKey)
	anaconda.SetConsumerSecret(m.ConsumerSecret)
	m.api = anaconda.NewTwitterApi(m.AccessToken, m.AccessTokenSecret)

	if !m.ReadOnly {
		client.CmdHook("privmsg", m.tweetCmd)
		client.CmdHook("privmsg", m.replyCmd)
		client.CmdHook("privmsg", m.retweetCmd)
		client.CmdHook("privmsg", m.favoriteCmd)
	}

	values := url.Values{}
	values.Add("replies", "all")
	values.Add("with", "user")

	go func(c *irc.Client, v url.Values) {
		for {
			m.streamHandler(c, v)
		}
	}(client, values)

	return nil
}

func (m *Module) tweetCmd(client *irc.Client, msg irc.Message) error {
	splited := strings.Fields(msg.Data)
	if len(splited) < 2 || splited[0] != "!tweet" || !client.Connected(msg.Receiver) {
		return nil
	}

	status := strings.Join(splited[1:], " ")
	if _, err := m.api.PostTweet(status, url.Values{}); err != nil {
		return client.Write("NOTICE %s :ERROR: %s",
			msg.Receiver, err.Error())
	}

	return nil
}

func (m *Module) replyCmd(client *irc.Client, msg irc.Message) error {
	splited := strings.Fields(msg.Data)
	if len(splited) < 3 || splited[0] != "!reply" || !client.Connected(msg.Receiver) {
		return nil
	}

	values := url.Values{}
	values.Add("in_reply_to_status_id", splited[1])

	status := strings.Join(splited[2:], " ")
	if !strings.Contains(status, "@") {
		return client.Write("NOTICE %s :ERROR: %s",
			msg.Receiver, "A reply must contain a @mention")
	}

	if _, err := m.api.PostTweet(status, values); err != nil {
		return client.Write("NOTICE %s :ERROR: %s",
			msg.Receiver, err.Error())
	}

	return nil
}

func (m *Module) retweetCmd(client *irc.Client, msg irc.Message) error {
	splited := strings.Fields(msg.Data)
	if len(splited) < 2 || splited[0] != "!retweet" || !client.Connected(msg.Receiver) {
		return nil
	}

	id, err := strconv.Atoi(splited[1])
	if err != nil {
		return err
	}

	if _, err := m.api.Retweet(int64(id), false); err != nil {
		return client.Write("NOTICE %s :ERROR: %s",
			msg.Receiver, err.Error())
	}

	return nil
}

func (m *Module) favoriteCmd(client *irc.Client, msg irc.Message) error {
	splited := strings.Fields(msg.Data)
	if len(splited) < 2 || splited[0] != "!favorite" || !client.Connected(msg.Receiver) {
		return nil
	}

	id, err := strconv.Atoi(splited[1])
	if err != nil {
		return err
	}

	if _, err := m.api.Favorite(int64(id)); err != nil {
		return client.Write("NOTICE %s :ERROR: %s",
			msg.Receiver, err.Error())
	}

	return nil
}

func (m *Module) streamHandler(client *irc.Client, values url.Values) {
	stream := m.api.UserStream(values)
	for {
		select {
		case event, ok := <-stream.C:
			if !ok {
				break
			}

			if t := m.formatEvent(event); len(t) > 0 {
				m.notify(client, t)
			}
		}
	}

	stream.Stop()
}

func (m *Module) formatEvent(event interface{}) string {
	var msg string
	switch t := event.(type) {
	case anaconda.ApiError:
		msg = fmt.Sprintf("Twitter API error %d: %s", t.StatusCode, t.Decoded.Error())
	case anaconda.StatusDeletionNotice:
		msg = fmt.Sprintf("Tweet %d has been deleted", t.Id)
	case anaconda.DirectMessage:
		msg = fmt.Sprintf("Direct message %d by %s: %s", t.Id,
			t.SenderScreenName, t.Text)
	case anaconda.Tweet:
		msg = fmt.Sprintf("Tweet %d by %s: %s", t.Id, t.User.ScreenName,
			html.UnescapeString(t.Text))
	case anaconda.EventTweet:
		if t.Event.Event != "favorite" {
			break
		}

		text := html.UnescapeString(t.TargetObject.Text)
		msg = fmt.Sprintf("%s favorited tweet %d by %s: %s",
			t.Source.ScreenName, t.TargetObject.Id, t.Target.ScreenName, text)
	}

	return msg
}

func (m *Module) notify(client *irc.Client, text string) {
	for _, ch := range client.Channels {
		client.Write("NOTICE %s :%s", ch, html.UnescapeString(text))
	}
}
<commit_msg>twitter: Better feedback when message is too long<commit_after>// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see <http://www.gnu.org/licenses/>.

package twitter

import (
	"fmt"
	"github.com/ChimeraCoder/anaconda"
	"github.com/nmeum/marvin/irc"
	"github.com/nmeum/marvin/modules"
	"html"
	"net/url"
	"strconv"
	"strings"
)

// Maximum amount of characters allowed in a tweet.
const maxChars = 140

type Module struct {
	api               *anaconda.TwitterApi
	ReadOnly          bool   `json:"read_only"`
	ConsumerKey       string `json:"consumer_key"`
	ConsumerSecret    string `json:"consumer_secret"`
	AccessToken       string `json:"access_token"`
	AccessTokenSecret string `json:"access_token_secret"`
}

func Init(moduleSet *modules.ModuleSet) {
	moduleSet.Register(new(Module))
}

func (m *Module) Name() string {
	return "twitter"
}

func (m *Module) Help() string {
	return "USAGE: !tweet TEXT || !reply ID TEXT || !retweet ID || !favorite ID"
}

func (m *Module) Defaults() {
	m.ReadOnly = false
}

func (m *Module) Load(client *irc.Client) error {
	anaconda.SetConsumerKey(m.ConsumerKey)
	anaconda.SetConsumerSecret(m.ConsumerSecret)
	m.api = anaconda.NewTwitterApi(m.AccessToken, m.AccessTokenSecret)

	if !m.ReadOnly {
		client.CmdHook("privmsg", m.tweetCmd)
		client.CmdHook("privmsg", m.replyCmd)
		client.CmdHook("privmsg", m.retweetCmd)
		client.CmdHook("privmsg", m.favoriteCmd)
	}

	values := url.Values{}
	values.Add("replies", "all")
	values.Add("with", "user")

	go func(c *irc.Client, v url.Values) {
		for {
			m.streamHandler(c, v)
		}
	}(client, values)

	return nil
}

func (m *Module) tweet(t string, v url.Values, c *irc.Client, p irc.Message) error {
	tLen := len(t)
	if tLen > maxChars {
		return c.Write("NOTICE %s :ERROR: Tweet is too long, remove %d characters",
			p.Receiver, tLen-maxChars)
	}

	if _, err := m.api.PostTweet(t, v); err != nil {
		return c.Write("NOTICE %s :ERROR: %s", p.Receiver, err.Error())
	}

	return nil
}

func (m *Module) tweetCmd(client *irc.Client, msg irc.Message) error {
	splited := strings.Fields(msg.Data)
	if len(splited) < 2 || splited[0] != "!tweet" || !client.Connected(msg.Receiver) {
		return nil
	}

	status := strings.Join(splited[1:], " ")
	return m.tweet(status, url.Values{}, client, msg)
}

func (m *Module) replyCmd(client *irc.Client, msg irc.Message) error {
	splited := strings.Fields(msg.Data)
	if len(splited) < 3 || splited[0] != "!reply" || !client.Connected(msg.Receiver) {
		return nil
	}

	status := strings.Join(splited[2:], " ")
	if !strings.Contains(status, "@") {
		return client.Write("NOTICE %s :ERROR: %s",
			msg.Receiver, "A reply must contain an @mention")
	}

	values := url.Values{}
	values.Add("in_reply_to_status_id", splited[1])

	return m.tweet(status, values, client, msg)
}

func (m *Module) retweetCmd(client *irc.Client, msg irc.Message) error {
	splited := strings.Fields(msg.Data)
	if len(splited) < 2 || splited[0] != "!retweet" || !client.Connected(msg.Receiver) {
		return nil
	}

	id, err := strconv.Atoi(splited[1])
	if err != nil {
		return err
	}

	if _, err := m.api.Retweet(int64(id), false); err != nil {
		return client.Write("NOTICE %s :ERROR: %s",
			msg.Receiver, err.Error())
	}

	return nil
}

func (m *Module) favoriteCmd(client *irc.Client, msg irc.Message) error {
	splited := strings.Fields(msg.Data)
	if len(splited) < 2 || splited[0] != "!favorite" || !client.Connected(msg.Receiver) {
		return nil
	}

	id, err := strconv.Atoi(splited[1])
	if err != nil {
		return err
	}

	if _, err := m.api.Favorite(int64(id)); err != nil {
		return client.Write("NOTICE %s :ERROR: %s",
			msg.Receiver, err.Error())
	}

	return nil
}

func (m *Module) streamHandler(client *irc.Client, values url.Values) {
	stream := m.api.UserStream(values)
	for {
		select {
		case event, ok := <-stream.C:
			if !ok {
				break
			}

			if t := m.formatEvent(event); len(t) > 0 {
				m.notify(client, t)
			}
		}
	}

	stream.Stop()
}

func (m *Module) formatEvent(event interface{}) string {
	var msg string
	switch t := event.(type) {
	case anaconda.ApiError:
		msg = fmt.Sprintf("Twitter API error %d: %s", t.StatusCode, t.Decoded.Error())
	case anaconda.StatusDeletionNotice:
		msg = fmt.Sprintf("Tweet %d has been deleted", t.Id)
	case anaconda.DirectMessage:
		msg = fmt.Sprintf("Direct message %d by %s: %s", t.Id,
			t.SenderScreenName, t.Text)
	case anaconda.Tweet:
		msg = fmt.Sprintf("Tweet %d by %s: %s", t.Id, t.User.ScreenName,
			html.UnescapeString(t.Text))
	case anaconda.EventTweet:
		if t.Event.Event != "favorite" {
			break
		}

		text := html.UnescapeString(t.TargetObject.Text)
		msg = fmt.Sprintf("%s favorited tweet %d by %s: %s",
			t.Source.ScreenName, t.TargetObject.Id, t.Target.ScreenName, text)
	}

	return msg
}

func (m *Module) notify(client *irc.Client, text string) {
	for _, ch := range client.Channels {
		client.Write("NOTICE %s :%s", ch, html.UnescapeString(text))
	}
}
<|endoftext|>


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

package tabletserver

import (
	"fmt"
	"math/rand"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/mysqlconn/fakesqldb"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/tabletserver/tabletenv"
	"github.com/youtube/vitess/go/vt/vterrors"

	vtrpcpb "github.com/youtube/vitess/go/vt/proto/vtrpc"
)

func TestTxPoolExecuteRollback(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddQuery("rollback", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	defer txPool.Rollback(ctx, transactionID)
	txConn.RecordQuery(sql)
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
}

func TestTxPoolRollbackNonBusy(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddQuery("rollback", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	txid1, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	_, err = txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	conn1, err := txPool.Get(txid1, "for query")
	if err != nil {
		t.Fatal(err)
	}
	// This should rollback only txid2.
	txPool.RollbackNonBusy(ctx)
	if sz := txPool.activePool.Size(); sz != 1 {
		t.Errorf("txPool.activePool.Size(): %d, want 1", sz)
	}
	conn1.Recycle()
	// This should rollback txid1.
	txPool.RollbackNonBusy(ctx)
	if sz := txPool.activePool.Size(); sz != 0 {
		t.Errorf("txPool.activePool.Size(): %d, want 0", sz)
	}
}

func TestTxPoolTransactionKiller(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	// make sure transaction killer will run frequent enough
	txPool.SetTimeout(1 * time.Millisecond)
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	killCount := tabletenv.KillStats.Counts()["Transactions"]
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	txConn.RecordQuery(sql)
	txConn.Recycle()
	// transaction killer should kill the query
	txPool.WaitForEmpty()
	killCountDiff := tabletenv.KillStats.Counts()["Transactions"] - killCount
	if killCountDiff != 1 {
		t.Fatalf("query: %s should be killed by transaction killer", sql)
	}
}

func TestTxPoolBeginAfterConnPoolClosed(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	txPool := newTxPool()
	txPool.SetTimeout(time.Duration(10))
	txPool.Open(db.ConnParams(), db.ConnParams())

	txPool.Close()

	_, err := txPool.Begin(context.Background())
	if err == nil {
		t.Fatalf("expect to get an error")
	}
	terr, ok := err.(*tabletenv.TabletError)
	if !ok || terr != tabletenv.ErrConnPoolClosed {
		t.Fatalf("get error: %v, but expect: %v", terr, tabletenv.ErrConnPoolClosed)
	}
}

func TestTxPoolBeginWithPoolConnectionError(t *testing.T) {
	db := fakesqldb.New(t)
	db.Close()
	//	db.EnableConnFail()
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	want := "errno 2003"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolBeginWithError(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddRejectedQuery("begin", errRejected)
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	want := "error: rejected"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
	if got, want := vterrors.RecoverVtErrorCode(err), vtrpcpb.ErrorCode_UNKNOWN_ERROR; got != want {
		t.Errorf("wrong error code for Begin error: got = %v, want = %v", got, want)
	}
}

func TestTxPoolRollbackFail(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddRejectedQuery("rollback", errRejected)

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	txConn.RecordQuery(sql)
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	err = txPool.Rollback(ctx, transactionID)
	want := "error: rejected"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolGetConnFail(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	_, err := txPool.Get(12345, "for query")
	want := "not_in_tx: Transaction 12345: not found"
	if err == nil || err.Error() != want {
		t.Errorf("Get: %v, want %s", err, want)
	}
}

func TestTxPoolExecFailDueToConnFail(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	sql := "alter table test_table add test_column int"

	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	db.EnableConnFail()
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err == nil {
		t.Fatalf("exec should fail because of a conn error")
	}
}

func TestTxPoolCloseKillsStrayTransactions(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())

	// Start stray transaction.
	_, err := txPool.Begin(context.Background())
	if err != nil {
		t.Fatal(err)
	}

	// Close kills stray transaction.
	txPool.Close()
	if got, want := tabletenv.InternalErrors.Counts()["StrayTransactions"], int64(1); got != want {
		t.Fatalf("internal error count for stray transactions not increased: got = %v, want = %v", got, want)
	}
	if got, want := txPool.conns.Capacity(), int64(0); got != want {
		t.Fatalf("resource pool was not closed. capacity: got = %v, want = %v", got, want)
	}
}

func newTxPool() *TxPool {
	randID := rand.Int63()
	poolName := fmt.Sprintf("TestTransactionPool-%d", randID)
	transactionCap := 300
	transactionTimeout := time.Duration(30 * time.Second)
	idleTimeout := time.Duration(30 * time.Second)
	return NewTxPool(
		poolName,
		transactionCap,
		transactionTimeout,
		idleTimeout,
		DummyChecker,
	)
}
<commit_msg>tabletserver: Rename test to better express the intent.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletserver

import (
	"fmt"
	"math/rand"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/mysqlconn/fakesqldb"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/tabletserver/tabletenv"
	"github.com/youtube/vitess/go/vt/vterrors"

	vtrpcpb "github.com/youtube/vitess/go/vt/proto/vtrpc"
)

func TestTxPoolExecuteRollback(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddQuery("rollback", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	defer txPool.Rollback(ctx, transactionID)
	txConn.RecordQuery(sql)
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
}

func TestTxPoolRollbackNonBusy(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddQuery("rollback", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	txid1, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	_, err = txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	conn1, err := txPool.Get(txid1, "for query")
	if err != nil {
		t.Fatal(err)
	}
	// This should rollback only txid2.
	txPool.RollbackNonBusy(ctx)
	if sz := txPool.activePool.Size(); sz != 1 {
		t.Errorf("txPool.activePool.Size(): %d, want 1", sz)
	}
	conn1.Recycle()
	// This should rollback txid1.
	txPool.RollbackNonBusy(ctx)
	if sz := txPool.activePool.Size(); sz != 0 {
		t.Errorf("txPool.activePool.Size(): %d, want 0", sz)
	}
}

func TestTxPoolTransactionKiller(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	// make sure transaction killer will run frequent enough
	txPool.SetTimeout(1 * time.Millisecond)
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	killCount := tabletenv.KillStats.Counts()["Transactions"]
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	txConn.RecordQuery(sql)
	txConn.Recycle()
	// transaction killer should kill the query
	txPool.WaitForEmpty()
	killCountDiff := tabletenv.KillStats.Counts()["Transactions"] - killCount
	if killCountDiff != 1 {
		t.Fatalf("query: %s should be killed by transaction killer", sql)
	}
}

func TestTxPoolBeginAfterConnPoolClosed(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	txPool := newTxPool()
	txPool.SetTimeout(time.Duration(10))
	txPool.Open(db.ConnParams(), db.ConnParams())

	txPool.Close()

	_, err := txPool.Begin(context.Background())
	if err == nil {
		t.Fatalf("expect to get an error")
	}
	terr, ok := err.(*tabletenv.TabletError)
	if !ok || terr != tabletenv.ErrConnPoolClosed {
		t.Fatalf("get error: %v, but expect: %v", terr, tabletenv.ErrConnPoolClosed)
	}
}

func TestTxPoolBeginWithPoolConnectionError(t *testing.T) {
	db := fakesqldb.New(t)
	db.Close()
	//	db.EnableConnFail()
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	want := "errno 2003"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolBeginWithError(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddRejectedQuery("begin", errRejected)
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	_, err := txPool.Begin(ctx)
	want := "error: rejected"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
	if got, want := vterrors.RecoverVtErrorCode(err), vtrpcpb.ErrorCode_UNKNOWN_ERROR; got != want {
		t.Errorf("wrong error code for Begin error: got = %v, want = %v", got, want)
	}
}

func TestTxPoolRollbackFail(t *testing.T) {
	sql := "alter table test_table add test_column int"
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery(sql, &sqltypes.Result{})
	db.AddQuery("begin", &sqltypes.Result{})
	db.AddRejectedQuery("rollback", errRejected)

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	txConn.RecordQuery(sql)
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err != nil {
		t.Fatalf("got error: %v", err)
	}
	err = txPool.Rollback(ctx, transactionID)
	want := "error: rejected"
	if err == nil || !strings.Contains(err.Error(), want) {
		t.Errorf("Begin: %v, want %s", err, want)
	}
}

func TestTxPoolGetConnNonExistentTransaction(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	_, err := txPool.Get(12345, "for query")
	want := "not_in_tx: Transaction 12345: not found"
	if err == nil || err.Error() != want {
		t.Errorf("Get: %v, want %s", err, want)
	}
}

func TestTxPoolExecFailDueToConnFail(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())
	defer txPool.Close()
	ctx := context.Background()
	sql := "alter table test_table add test_column int"

	transactionID, err := txPool.Begin(ctx)
	if err != nil {
		t.Fatal(err)
	}
	txConn, err := txPool.Get(transactionID, "for query")
	if err != nil {
		t.Fatal(err)
	}
	db.EnableConnFail()
	_, err = txConn.Exec(ctx, sql, 1, true)
	txConn.Recycle()
	if err == nil {
		t.Fatalf("exec should fail because of a conn error")
	}
}

func TestTxPoolCloseKillsStrayTransactions(t *testing.T) {
	db := fakesqldb.New(t)
	defer db.Close()
	db.AddQuery("begin", &sqltypes.Result{})

	txPool := newTxPool()
	txPool.Open(db.ConnParams(), db.ConnParams())

	// Start stray transaction.
	_, err := txPool.Begin(context.Background())
	if err != nil {
		t.Fatal(err)
	}

	// Close kills stray transaction.
	txPool.Close()
	if got, want := tabletenv.InternalErrors.Counts()["StrayTransactions"], int64(1); got != want {
		t.Fatalf("internal error count for stray transactions not increased: got = %v, want = %v", got, want)
	}
	if got, want := txPool.conns.Capacity(), int64(0); got != want {
		t.Fatalf("resource pool was not closed. capacity: got = %v, want = %v", got, want)
	}
}

func newTxPool() *TxPool {
	randID := rand.Int63()
	poolName := fmt.Sprintf("TestTransactionPool-%d", randID)
	transactionCap := 300
	transactionTimeout := time.Duration(30 * time.Second)
	idleTimeout := time.Duration(30 * time.Second)
	return NewTxPool(
		poolName,
		transactionCap,
		transactionTimeout,
		idleTimeout,
		DummyChecker,
	)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/cosn/firebase"
	"log"
	"time"
)

type Station struct {
	Code string `json:"CODE_LIGNE"`
	Name        string `json:"NOM"`
	Latitude   string `json:"LATITUDE"`
	Longitude  string `json:"LONGITUDE"`
	Owner      string `json:"owner"`
	OwnerMail  string `json:"OwnerMail"`
	When       int64  `json:"when"`
}

func main() {
	client := new(firebase.Client)
	client.Init("https://shining-inferno-9452.firebaseio.com", "a6O8UkBnRaeQ88BjUHv46VMNFBAhUDwKdFDfcfYz", nil)

	for i := 1; i < 6441; i++ {
		station := &Station{};
		stationChild := client.Child(fmt.Sprintf("%v", i), nil, station)
		if (station.Owner != "") {
			if (time.Now().UnixNano()/int64(time.Millisecond)-station.When > 5*3600*1000) {
				log.Printf("Free station : %v", i)
				stationChild.Set("", &Station{Code: station.Code, Name: station.Name, Latitude: station.Latitude, Longitude: station.Longitude }, nil)
			}
		}
	}
}
<commit_msg>Update script<commit_after>package main

import (
	"fmt"
	"github.com/cosn/firebase"
	"log"
	"time"
)

type Station struct {
	Code        string `json:"CODE_LIGNE"`
	Name        string `json:"NOM"`
	Latitude    string `json:"LATITUDE"`
	Longitude   string `json:"LONGITUDE"`
	Owner       string `json:"owner"`
	OwnerMail   string `json:"OwnerMail"`
	When        int64  `json:"when"`
}

func main() {
	client := new(firebase.Client)
	client.Init("https://shining-inferno-9452.firebaseio.com", "a6O8UkBnRaeQ88BjUHv46VMNFBAhUDwKdFDfcfYz", nil)

	for i := 1; i < 6441; i++ {
		station := &Station{};
		stationChild := client.Child("stations", nil, nil).Child(fmt.Sprintf("%v", i), nil, station)
		if (station.Owner != "") {
			if (time.Now().UnixNano()/int64(time.Millisecond)-station.When > 5*3600*1000) {
				log.Printf("Free station : %v", i)
				stationChild.Set("", &Station{Code: station.Code, Name: station.Name, Latitude: station.Latitude, Longitude: station.Longitude }, nil)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

type BlobOwnership struct {
	OID    string               `json:"oid"`
	Length int64                `json:"length"`
	Nodes  map[string]time.Time `json:"nodes"`
	Type   string               `json:"type"`
}

type internodeCommand uint8

const (
	removeObjectCmd = internodeCommand(iota)
	acquireObjectCmd
	fetchObjectCmd
)

type internodeTask struct {
	node     StorageNode
	cmd      internodeCommand
	oid      string
	prevNode string
}

var taskWorkers = flag.Int("taskWorkers", 4,
	"Number of blob move/removal workers.")

func (b BlobOwnership) ResolveNodes() NodeList {
	keys := make([]string, 0, len(b.Nodes))
	for k := range b.Nodes {
		keys = append(keys, "/"+k)
	}
	resps := couchbase.GetBulk(keys)

	rv := make(NodeList, 0, len(resps))

	for k, v := range resps {
		if v.Status == gomemcached.SUCCESS {
			a := StorageNode{}
			err := json.Unmarshal(v.Body, &a)
			if err == nil {
				a.name = k[1:]
				rv = append(rv, a)
			}
		}
	}

	sort.Sort(rv)

	return rv
}

// Get the most recent storer of a blob
func (b BlobOwnership) mostRecent() (string, time.Time) {
	rvnode := ""
	rvt := time.Time{}

	for node, t := range b.Nodes {
		if t.After(rvt) {
			rvnode = node
			rvt = t
		}
	}

	return rvnode, rvt
}

func (b BlobOwnership) ResolveRemoteNodes() NodeList {
	return b.ResolveNodes().minusLocal()
}

func getBlobOwnership(oid string) (BlobOwnership, error) {
	rv := BlobOwnership{}
	oidkey := "/" + oid
	err := couchbase.Get(oidkey, &rv)
	return rv, err
}

func copyBlob(w io.Writer, oid string) error {
	f, err := os.Open(hashFilename(*root, oid))
	if err == nil {
		// Doing it locally
		defer f.Close()
		_, err = io.Copy(w, f)
		return err
	} else {
		// Doing it remotely
		c := captureResponseWriter{w: w}
		getBlobFromRemote(&c, oid, http.Header{}, 0)
		if c.statusCode != 200 {
			return fmt.Errorf("Error grabbing remote object, got %v",
				c.statusCode)
		}
		return nil
	}
	panic("unreachable")
}

func recordBlobOwnership(h string, l int64, force bool) error {
	k := "/" + h
	err := couchbase.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.CAS(vb, k, func(in []byte) ([]byte, memcached.CasOp) {
			ownership := BlobOwnership{}
			err := json.Unmarshal(in, &ownership)
			if err == nil {
				if _, ok := ownership.Nodes[serverId]; ok && !force {
					// Skip it fast if it already knows us
					return nil, memcached.CASQuit
				}
				ownership.Nodes[serverId] = time.Now().UTC()
			} else {
				ownership.Nodes = map[string]time.Time{
					serverId: time.Now().UTC(),
				}
			}
			ownership.OID = h
			ownership.Length = l
			ownership.Type = "blob"
			return mustEncode(&ownership), memcached.CASStore
		}, 0)
		return err
	})
	if err == memcached.CASQuit {
		err = nil
	}
	return err
}

func recordBlobAccess(h string) {
	_, err := couchbase.Incr("/"+h+"/r", 1, 1, 0)
	if err != nil {
		log.Printf("Error incrementing counter for %v: %v", h, err)
	}

	_, err = couchbase.Incr("/"+serverId+"/r", 1, 1, 0)
	if err != nil {
		log.Printf("Error incrementing node identifier: %v", err)
	}
}

// Returns the number of known owners (-1 if it can't be determined)
func removeBlobOwnershipRecord(h, node string) int {
	log.Printf("Cleaning up %v from %v", h, node)
	numOwners := -1

	k := "/" + h
	err := couchbase.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.CAS(vb, k, func(in []byte) ([]byte, memcached.CasOp) {
			ownership := BlobOwnership{}

			if len(in) == 0 {
				return nil, memcached.CASQuit
			}

			err := json.Unmarshal(in, &ownership)
			if err == nil {
				delete(ownership.Nodes, node)
			} else {
				log.Printf("Error unmarhaling blob removal from %s: %v",
					in, err)
				return nil, memcached.CASQuit
			}

			var rv []byte
			op := memcached.CASStore

			numOwners = len(ownership.Nodes)

			if len(ownership.Nodes) == 0 && node == serverId {
				op = memcached.CASDelete
			} else {
				rv = mustEncode(&ownership)
			}

			return rv, op
		}, 0)
		return err
	})
	if err != nil && err != memcached.CASQuit {
		log.Printf("Error cleaning %v from %v: %v", node, h, err)
		numOwners = -1
	}
	if numOwners == 0 {
		couchbase.Delete(k + "/r")
	}

	return numOwners
}

func increaseReplicaCount(oid string, length int64, by int) error {
	nl, err := findAllNodes()
	if err != nil {
		return err
	}
	onto := nl.candidatesFor(oid, NodeList{})
	if len(onto) > by {
		onto = onto[:by]
	}
	for _, n := range onto {
		log.Printf("Asking %v to acquire %v", n, oid)
		queueBlobAcquire(n, oid, "")
	}
	return nil
}

func ensureMinimumReplicaCount() error {
	return runMarkedTask("ensureMinReplCount", "garbageCollectBlobs",
		ensureMinimumReplicaCountTask)
}

func ensureMinimumReplicaCountTask() error {
	// Don't let this run concurrently with the garbage collector.
	// They don't get along.
	for taskRunning("garbageCollectBlobs") {
		log.Printf("Waiting for gc to finish for ensureMinReplCount")
		time.Sleep(5 * time.Second)
	}

	nl, err := findAllNodes()
	if err != nil {
		return err
	}

	viewRes := struct {
		Rows []struct {
			Id string
		}
	}{}

	// Don't bother trying to replicate to more nodes than exist.
	endKey := globalConfig.MinReplicas - 1
	if globalConfig.MinReplicas > len(nl) {
		endKey = len(nl) - 1
	}

	// Find some less replicated docs to suck in.
	err = couchbase.ViewCustom("cbfs", "repcounts",
		map[string]interface{}{
			"reduce":   false,
			"limit":    globalConfig.ReplicationCheckLimit,
			"startkey": 1,
			"endkey":   endKey,
			"stale":    false,
		},
		&viewRes)

	if err != nil {
		return err
	}

	log.Printf("Increasing replica count of %v items",
		len(viewRes.Rows))

	for _, r := range viewRes.Rows {
		salvageBlob(r.Id[1:], "", nl)
	}
	return nil
}

func pruneBlob(oid string, nodemap map[string]string, nl NodeList) {
	if len(nodemap) <= globalConfig.MaxReplicas {
		log.Printf("Asked to prune a blob that has too few replicas: %v",
			oid)
	}

	log.Printf("Pruning blob %v down from %v repls to %v",
		oid, len(nodemap), globalConfig.MaxReplicas)

	nm := map[string]StorageNode{}
	for _, n := range nl {
		nm[n.name] = n
	}

	remaining := len(nodemap)
	for n := range nodemap {
		if remaining <= globalConfig.MaxReplicas {
			break
		}
		remaining--
		if sn, ok := nm[n]; ok {
			queueBlobRemoval(sn, oid)
		}
	}

}

func pruneExcessiveReplicas() error {
	nl, err := findAllNodes()
	if err != nil {
		return err
	}

	viewRes := struct {
		Rows []struct {
			Id  string
			Doc struct {
				Json struct {
					Nodes map[string]string
				}
			}
		}
	}{}

	// Find some less replicated docs to suck in.
	err = couchbase.ViewCustom("cbfs", "repcounts",
		map[string]interface{}{
			"descending":   true,
			"reduce":       false,
			"include_docs": true,
			"limit":        globalConfig.ReplicationCheckLimit,
			"endkey":       globalConfig.MaxReplicas + 1,
			"stale":        false,
		},
		&viewRes)

	if err != nil {
		return err
	}

	log.Printf("Decreasing replica count of %v items",
		len(viewRes.Rows))

	// Short-circuit when there's nothing to clean
	if len(viewRes.Rows) == 0 {
		return nil
	}

	for _, r := range viewRes.Rows {
		pruneBlob(r.Id[1:], r.Doc.Json.Nodes, nl)
	}
	return nil
}

func hasBlob(oid string) bool {
	_, err := os.Stat(hashFilename(*root, oid))
	return err == nil
}

func performFetch(oid, prev string) {
	c := captureResponseWriter{w: ioutil.Discard}

	// If we already have it, we don't need it more.
	f, err := os.Open(hashFilename(*root, oid))
	if err == nil {
		f.Close()
		return
	}

	getBlobFromRemote(&c, oid, http.Header{}, 100)

	if c.statusCode == 200 {
		if prev != "" {
			removeBlobOwnershipRecord(oid, prev)
			log.Printf("Removing ownership of %v from %v after takeover",
				oid, prev)
		}
	} else {
		log.Printf("Error grabbing remote object, got %v",
			c.statusCode)
	}
}

func salvageBlob(oid, deadNode string, nl NodeList) {
	candidates := nl.candidatesFor(oid,
		NodeList{nl.named(deadNode)})

	if len(candidates) == 0 {
		log.Printf("Couldn't find a candidate for blob!")
	} else {
		log.Printf("Recommending %v get a copy of %v",
			candidates[0], oid)
		queueBlobAcquire(candidates[0], oid, deadNode)
	}
}

var internodeTaskQueue = make(chan internodeTask, 1000)

func internodeTaskWorker() {
	for c := range internodeTaskQueue {
		switch c.cmd {
		case removeObjectCmd:
			if err := c.node.deleteBlob(c.oid); err != nil {
				log.Printf("Error deleting %v from %v: %v",
					c.oid, c.node, err)
			}
		case acquireObjectCmd:
			if err := c.node.acquireBlob(c.oid, c.prevNode); err != nil {
				log.Printf("Error acquiring %v from %v: %v",
					c.oid, c.node, err)
			}
		case fetchObjectCmd:
			performFetch(c.oid, c.prevNode)
		default:
			log.Fatalf("Unhandled worker task: %v", c)
		}
	}
}

func initTaskQueueWorkers() {
	for i := 0; i < *taskWorkers; i++ {
		go internodeTaskWorker()
	}
}

func queueBlobRemoval(n StorageNode, oid string) {
	internodeTaskQueue <- internodeTask{
		node: n,
		cmd:  removeObjectCmd,
		oid:  oid,
	}
}

// Ask a remote node to go get a blob
func queueBlobAcquire(n StorageNode, oid string, prev string) {
	internodeTaskQueue <- internodeTask{
		node:     n,
		cmd:      acquireObjectCmd,
		oid:      oid,
		prevNode: prev,
	}
}

// Ask this node to go get a blob
func queueBlobFetch(oid, prev string) {
	internodeTaskQueue <- internodeTask{
		cmd:      fetchObjectCmd,
		oid:      oid,
		prevNode: prev,
	}
}
<commit_msg>No reason to avoid locally caching while making a .zip<commit_after>package main

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

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

type BlobOwnership struct {
	OID    string               `json:"oid"`
	Length int64                `json:"length"`
	Nodes  map[string]time.Time `json:"nodes"`
	Type   string               `json:"type"`
}

type internodeCommand uint8

const (
	removeObjectCmd = internodeCommand(iota)
	acquireObjectCmd
	fetchObjectCmd
)

type internodeTask struct {
	node     StorageNode
	cmd      internodeCommand
	oid      string
	prevNode string
}

var taskWorkers = flag.Int("taskWorkers", 4,
	"Number of blob move/removal workers.")

func (b BlobOwnership) ResolveNodes() NodeList {
	keys := make([]string, 0, len(b.Nodes))
	for k := range b.Nodes {
		keys = append(keys, "/"+k)
	}
	resps := couchbase.GetBulk(keys)

	rv := make(NodeList, 0, len(resps))

	for k, v := range resps {
		if v.Status == gomemcached.SUCCESS {
			a := StorageNode{}
			err := json.Unmarshal(v.Body, &a)
			if err == nil {
				a.name = k[1:]
				rv = append(rv, a)
			}
		}
	}

	sort.Sort(rv)

	return rv
}

// Get the most recent storer of a blob
func (b BlobOwnership) mostRecent() (string, time.Time) {
	rvnode := ""
	rvt := time.Time{}

	for node, t := range b.Nodes {
		if t.After(rvt) {
			rvnode = node
			rvt = t
		}
	}

	return rvnode, rvt
}

func (b BlobOwnership) ResolveRemoteNodes() NodeList {
	return b.ResolveNodes().minusLocal()
}

func getBlobOwnership(oid string) (BlobOwnership, error) {
	rv := BlobOwnership{}
	oidkey := "/" + oid
	err := couchbase.Get(oidkey, &rv)
	return rv, err
}

func copyBlob(w io.Writer, oid string) error {
	f, err := os.Open(hashFilename(*root, oid))
	if err == nil {
		// Doing it locally
		defer f.Close()
		_, err = io.Copy(w, f)
		return err
	} else {
		// Doing it remotely
		c := captureResponseWriter{w: w}
		getBlobFromRemote(&c, oid, http.Header{}, *cachePercentage)
		if c.statusCode != 200 {
			return fmt.Errorf("Error grabbing remote object, got %v",
				c.statusCode)
		}
		return nil
	}
	panic("unreachable")
}

func recordBlobOwnership(h string, l int64, force bool) error {
	k := "/" + h
	err := couchbase.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.CAS(vb, k, func(in []byte) ([]byte, memcached.CasOp) {
			ownership := BlobOwnership{}
			err := json.Unmarshal(in, &ownership)
			if err == nil {
				if _, ok := ownership.Nodes[serverId]; ok && !force {
					// Skip it fast if it already knows us
					return nil, memcached.CASQuit
				}
				ownership.Nodes[serverId] = time.Now().UTC()
			} else {
				ownership.Nodes = map[string]time.Time{
					serverId: time.Now().UTC(),
				}
			}
			ownership.OID = h
			ownership.Length = l
			ownership.Type = "blob"
			return mustEncode(&ownership), memcached.CASStore
		}, 0)
		return err
	})
	if err == memcached.CASQuit {
		err = nil
	}
	return err
}

func recordBlobAccess(h string) {
	_, err := couchbase.Incr("/"+h+"/r", 1, 1, 0)
	if err != nil {
		log.Printf("Error incrementing counter for %v: %v", h, err)
	}

	_, err = couchbase.Incr("/"+serverId+"/r", 1, 1, 0)
	if err != nil {
		log.Printf("Error incrementing node identifier: %v", err)
	}
}

// Returns the number of known owners (-1 if it can't be determined)
func removeBlobOwnershipRecord(h, node string) int {
	log.Printf("Cleaning up %v from %v", h, node)
	numOwners := -1

	k := "/" + h
	err := couchbase.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.CAS(vb, k, func(in []byte) ([]byte, memcached.CasOp) {
			ownership := BlobOwnership{}

			if len(in) == 0 {
				return nil, memcached.CASQuit
			}

			err := json.Unmarshal(in, &ownership)
			if err == nil {
				delete(ownership.Nodes, node)
			} else {
				log.Printf("Error unmarhaling blob removal from %s: %v",
					in, err)
				return nil, memcached.CASQuit
			}

			var rv []byte
			op := memcached.CASStore

			numOwners = len(ownership.Nodes)

			if len(ownership.Nodes) == 0 && node == serverId {
				op = memcached.CASDelete
			} else {
				rv = mustEncode(&ownership)
			}

			return rv, op
		}, 0)
		return err
	})
	if err != nil && err != memcached.CASQuit {
		log.Printf("Error cleaning %v from %v: %v", node, h, err)
		numOwners = -1
	}
	if numOwners == 0 {
		couchbase.Delete(k + "/r")
	}

	return numOwners
}

func increaseReplicaCount(oid string, length int64, by int) error {
	nl, err := findAllNodes()
	if err != nil {
		return err
	}
	onto := nl.candidatesFor(oid, NodeList{})
	if len(onto) > by {
		onto = onto[:by]
	}
	for _, n := range onto {
		log.Printf("Asking %v to acquire %v", n, oid)
		queueBlobAcquire(n, oid, "")
	}
	return nil
}

func ensureMinimumReplicaCount() error {
	return runMarkedTask("ensureMinReplCount", "garbageCollectBlobs",
		ensureMinimumReplicaCountTask)
}

func ensureMinimumReplicaCountTask() error {
	// Don't let this run concurrently with the garbage collector.
	// They don't get along.
	for taskRunning("garbageCollectBlobs") {
		log.Printf("Waiting for gc to finish for ensureMinReplCount")
		time.Sleep(5 * time.Second)
	}

	nl, err := findAllNodes()
	if err != nil {
		return err
	}

	viewRes := struct {
		Rows []struct {
			Id string
		}
	}{}

	// Don't bother trying to replicate to more nodes than exist.
	endKey := globalConfig.MinReplicas - 1
	if globalConfig.MinReplicas > len(nl) {
		endKey = len(nl) - 1
	}

	// Find some less replicated docs to suck in.
	err = couchbase.ViewCustom("cbfs", "repcounts",
		map[string]interface{}{
			"reduce":   false,
			"limit":    globalConfig.ReplicationCheckLimit,
			"startkey": 1,
			"endkey":   endKey,
			"stale":    false,
		},
		&viewRes)

	if err != nil {
		return err
	}

	log.Printf("Increasing replica count of %v items",
		len(viewRes.Rows))

	for _, r := range viewRes.Rows {
		salvageBlob(r.Id[1:], "", nl)
	}
	return nil
}

func pruneBlob(oid string, nodemap map[string]string, nl NodeList) {
	if len(nodemap) <= globalConfig.MaxReplicas {
		log.Printf("Asked to prune a blob that has too few replicas: %v",
			oid)
	}

	log.Printf("Pruning blob %v down from %v repls to %v",
		oid, len(nodemap), globalConfig.MaxReplicas)

	nm := map[string]StorageNode{}
	for _, n := range nl {
		nm[n.name] = n
	}

	remaining := len(nodemap)
	for n := range nodemap {
		if remaining <= globalConfig.MaxReplicas {
			break
		}
		remaining--
		if sn, ok := nm[n]; ok {
			queueBlobRemoval(sn, oid)
		}
	}

}

func pruneExcessiveReplicas() error {
	nl, err := findAllNodes()
	if err != nil {
		return err
	}

	viewRes := struct {
		Rows []struct {
			Id  string
			Doc struct {
				Json struct {
					Nodes map[string]string
				}
			}
		}
	}{}

	// Find some less replicated docs to suck in.
	err = couchbase.ViewCustom("cbfs", "repcounts",
		map[string]interface{}{
			"descending":   true,
			"reduce":       false,
			"include_docs": true,
			"limit":        globalConfig.ReplicationCheckLimit,
			"endkey":       globalConfig.MaxReplicas + 1,
			"stale":        false,
		},
		&viewRes)

	if err != nil {
		return err
	}

	log.Printf("Decreasing replica count of %v items",
		len(viewRes.Rows))

	// Short-circuit when there's nothing to clean
	if len(viewRes.Rows) == 0 {
		return nil
	}

	for _, r := range viewRes.Rows {
		pruneBlob(r.Id[1:], r.Doc.Json.Nodes, nl)
	}
	return nil
}

func hasBlob(oid string) bool {
	_, err := os.Stat(hashFilename(*root, oid))
	return err == nil
}

func performFetch(oid, prev string) {
	c := captureResponseWriter{w: ioutil.Discard}

	// If we already have it, we don't need it more.
	f, err := os.Open(hashFilename(*root, oid))
	if err == nil {
		f.Close()
		return
	}

	getBlobFromRemote(&c, oid, http.Header{}, 100)

	if c.statusCode == 200 {
		if prev != "" {
			removeBlobOwnershipRecord(oid, prev)
			log.Printf("Removing ownership of %v from %v after takeover",
				oid, prev)
		}
	} else {
		log.Printf("Error grabbing remote object, got %v",
			c.statusCode)
	}
}

func salvageBlob(oid, deadNode string, nl NodeList) {
	candidates := nl.candidatesFor(oid,
		NodeList{nl.named(deadNode)})

	if len(candidates) == 0 {
		log.Printf("Couldn't find a candidate for blob!")
	} else {
		log.Printf("Recommending %v get a copy of %v",
			candidates[0], oid)
		queueBlobAcquire(candidates[0], oid, deadNode)
	}
}

var internodeTaskQueue = make(chan internodeTask, 1000)

func internodeTaskWorker() {
	for c := range internodeTaskQueue {
		switch c.cmd {
		case removeObjectCmd:
			if err := c.node.deleteBlob(c.oid); err != nil {
				log.Printf("Error deleting %v from %v: %v",
					c.oid, c.node, err)
			}
		case acquireObjectCmd:
			if err := c.node.acquireBlob(c.oid, c.prevNode); err != nil {
				log.Printf("Error acquiring %v from %v: %v",
					c.oid, c.node, err)
			}
		case fetchObjectCmd:
			performFetch(c.oid, c.prevNode)
		default:
			log.Fatalf("Unhandled worker task: %v", c)
		}
	}
}

func initTaskQueueWorkers() {
	for i := 0; i < *taskWorkers; i++ {
		go internodeTaskWorker()
	}
}

func queueBlobRemoval(n StorageNode, oid string) {
	internodeTaskQueue <- internodeTask{
		node: n,
		cmd:  removeObjectCmd,
		oid:  oid,
	}
}

// Ask a remote node to go get a blob
func queueBlobAcquire(n StorageNode, oid string, prev string) {
	internodeTaskQueue <- internodeTask{
		node:     n,
		cmd:      acquireObjectCmd,
		oid:      oid,
		prevNode: prev,
	}
}

// Ask this node to go get a blob
func queueBlobFetch(oid, prev string) {
	internodeTaskQueue <- internodeTask{
		cmd:      fetchObjectCmd,
		oid:      oid,
		prevNode: prev,
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 - Michael Woolnough <michael.woolnough@gmail.com>
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package minecraft

import (
	"fmt"
	"github.com/MJKWoolnough/equaler"
	"github.com/MJKWoolnough/minecraft/nbt"
)

type Tick struct {
	I, T, P int32
}

type Block struct {
	BlockId  uint16
	Data     uint8
	metadata []nbt.Tag
	ticks    []Tick
}

func (b Block) Equal(e equaler.Equaler) bool {
	c, ok := e.(*Block)
	if !ok {
		if d, ok := e.(Block); ok {
			c = &d
		}
	}
	if c != nil {
		if b.BlockId == c.BlockId && b.Data == c.Data && len(b.metadata) == len(c.metadata) && len(b.ticks) == len(c.ticks) {
			for _, bT := range b.ticks {
				found := false
				for _, cT := range c.ticks {
					if bT.I == cT.I && bT.T == cT.T && bT.P == cT.P {
						found = true
						break
					}
				}
				if !found {
					return false
				}
			}
			for _, v := range b.metadata {
				name := v.Name()
				found := false
				for _, w := range c.metadata {
					if w.Name() == name {
						if !v.Equal(w) {
							return false
						}
						found = true
						break
					}
				}
				if !found {
					return false
				}
			}
			return true
		}
	}
	return false
}

// Opacity returns how much light is blocked by this block.
func (b Block) Opacity() uint8 {
	if b.BlockId == 8 || b.BlockId == 9 {
		return 3
	}
	for i := 0; i < len(transparentBlocks); i++ {
		if transparentBlocks[i] == b.BlockId {
			return 1
		}
	}
	return 16
}

func (b Block) IsLiquid() bool {
	return b.BlockId == 8 || b.BlockId == 9 || b.BlockId == 10 || b.BlockId == 11
}

func (b Block) HasMetadata() bool {
	return len(b.metadata) > 0
}

func (b Block) GetMetadata() []nbt.Tag {
	if b.metadata == nil {
		return nil
	}
	a := make([]nbt.Tag, len(b.metadata))
	for i, j := range b.metadata {
		a[i] = j.Copy()
	}
	return a
}

func (b *Block) SetMetadata(data []nbt.Tag) {
	metadata := make([]nbt.Tag, 0)
	for i := 0; i < len(data); i++ {
		name := data[i].Name()
		if name == "x" || name == "y" || name == "z" {
			continue
		} else if data[i].TagId() == nbt.Tag_End {
			break
		}
		metadata = append(metadata, data[i].Copy())
	}
	if len(metadata) > 0 {
		b.metadata = metadata
	} else {
		b.metadata = nil
	}
}

func (b Block) HasTicks() bool {
	return len(b.ticks) > 0
}

func (b Block) GetTicks() []Tick {
	t := make([]Tick, len(b.ticks))
	copy(t, b.ticks)
	return t
}

func (b *Block) AddTicks(t ...Tick) {
	b.ticks = append(b.ticks, t...)
}

func (b *Block) SetTicks(t []Tick) {
	b.ticks = make([]Tick, len(t))
	copy(b.ticks, t)
}

func (b Block) String() string {
	toRet := fmt.Sprintf("Block ID: %d\nData: %d\n", b.BlockId, b.Data)
	if b.metadata != nil && len(b.metadata) != 0 {
		toRet += "Metadata:\n"
		for i := 0; i < len(b.metadata); i++ {
			toRet += "	" + b.metadata[i].String() + "\n"
		}
	}
	for n, tick := range b.ticks {
		toRet += fmt.Sprintf("	Tick: %d, i: %d, t: %d, p: %d\n", n+1, tick.I, tick.T, tick.P)
	}
	return toRet
}
<commit_msg>Added Block.Light() method to return block generated light<commit_after>// Copyright (c) 2013 - Michael Woolnough <michael.woolnough@gmail.com>
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package minecraft

import (
	"fmt"
	"github.com/MJKWoolnough/equaler"
	"github.com/MJKWoolnough/minecraft/nbt"
)

type Tick struct {
	I, T, P int32
}

type Block struct {
	BlockId  uint16
	Data     uint8
	metadata []nbt.Tag
	ticks    []Tick
}

func (b Block) Equal(e equaler.Equaler) bool {
	c, ok := e.(*Block)
	if !ok {
		if d, ok := e.(Block); ok {
			c = &d
		}
	}
	if c != nil {
		if b.BlockId == c.BlockId && b.Data == c.Data && len(b.metadata) == len(c.metadata) && len(b.ticks) == len(c.ticks) {
			for _, bT := range b.ticks {
				found := false
				for _, cT := range c.ticks {
					if bT.I == cT.I && bT.T == cT.T && bT.P == cT.P {
						found = true
						break
					}
				}
				if !found {
					return false
				}
			}
			for _, v := range b.metadata {
				name := v.Name()
				found := false
				for _, w := range c.metadata {
					if w.Name() == name {
						if !v.Equal(w) {
							return false
						}
						found = true
						break
					}
				}
				if !found {
					return false
				}
			}
			return true
		}
	}
	return false
}

// Opacity returns how much light is blocked by this block.
func (b Block) Opacity() uint8 {
	if b.BlockId == 8 || b.BlockId == 9 {
		return 3
	}
	for i := 0; i < len(transparentBlocks); i++ {
		if transparentBlocks[i] == b.BlockId {
			return 1
		}
	}
	return 16
}

// Light returns how much light is generated by this block.
func (b Block) Light() uint8 {
	if l, ok := lightBlocks[b.BlockId]; ok {
		return l
	}
	return 0
}

func (b Block) IsLiquid() bool {
	return b.BlockId == 8 || b.BlockId == 9 || b.BlockId == 10 || b.BlockId == 11
}

func (b Block) HasMetadata() bool {
	return len(b.metadata) > 0
}

func (b Block) GetMetadata() []nbt.Tag {
	if b.metadata == nil {
		return nil
	}
	a := make([]nbt.Tag, len(b.metadata))
	for i, j := range b.metadata {
		a[i] = j.Copy()
	}
	return a
}

func (b *Block) SetMetadata(data []nbt.Tag) {
	metadata := make([]nbt.Tag, 0)
	for i := 0; i < len(data); i++ {
		name := data[i].Name()
		if name == "x" || name == "y" || name == "z" {
			continue
		} else if data[i].TagId() == nbt.Tag_End {
			break
		}
		metadata = append(metadata, data[i].Copy())
	}
	if len(metadata) > 0 {
		b.metadata = metadata
	} else {
		b.metadata = nil
	}
}

func (b Block) HasTicks() bool {
	return len(b.ticks) > 0
}

func (b Block) GetTicks() []Tick {
	t := make([]Tick, len(b.ticks))
	copy(t, b.ticks)
	return t
}

func (b *Block) AddTicks(t ...Tick) {
	b.ticks = append(b.ticks, t...)
}

func (b *Block) SetTicks(t []Tick) {
	b.ticks = make([]Tick, len(t))
	copy(b.ticks, t)
}

func (b Block) String() string {
	toRet := fmt.Sprintf("Block ID: %d\nData: %d\n", b.BlockId, b.Data)
	if b.metadata != nil && len(b.metadata) != 0 {
		toRet += "Metadata:\n"
		for i := 0; i < len(b.metadata); i++ {
			toRet += "	" + b.metadata[i].String() + "\n"
		}
	}
	for n, tick := range b.ticks {
		toRet += fmt.Sprintf("	Tick: %d, i: %d, t: %d, p: %d\n", n+1, tick.I, tick.T, tick.P)
	}
	return toRet
}
<|endoftext|>


<commit_before>// Package blstr (blaster? blast string?) is an ill-named library
// for easily setting up best-effort broadcast communication between
// goroutines.
//
// Synchronization and consistency is traded for speed and
// predictability of send operations.
//
// The best use case is where perfect consistency may not be
// necessary and some message loss on slow receivers is more tolerable
// than choking upstream senders.
package blstr

import (
	"errors"
	"sync"
)

// ByteHub implements a multiuser communication channel.
type ByteHub struct {
	subscribers map[int]chan<- []byte
	sync.RWMutex
}

// New returns a new ByteHub
func New() *ByteHub {
	return &ByteHub{subscribers: make(map[int]chan<- []byte)}
}

// Subscribe to be notified when broadcast events occur on the hub.
//
// Messages can be dropped if receiver stops listening, or if messages
// arrive faster than they can be consumed.
// Channels can be buffered to mitigate this to some extent.
func (bh *ByteHub) Subscribe(id int, ch chan<- []byte) error {
	bh.Lock()
	defer bh.Unlock()

	if _, ok := bh.subscribers[id]; ok {
		return errors.New("subscriber already exists")
	}

	bh.subscribers[id] = ch

	return nil
}

// Unsubscribe tries to remove the subscription by ID if it exists.
func (bh *ByteHub) Unsubscribe(id int) {
	bh.Lock()
	defer bh.Unlock()

	delete(bh.subscribers, id)
}

// Send provides unicast communication over the hub.
//
// An attempt is made to send a message to the specified subscriber.
// Errors are returned if the subscriber could not be found or
// is unable to receive the message.
func (bh *ByteHub) Send(to int, msg []byte) error {
	bh.RLock()
	defer bh.RUnlock()

	ch, ok := bh.subscribers[to]
	if !ok {
		return errors.New("subscriber does not exist")
	}

	select {
	case ch <- msg:
		return nil
	default:
		return errors.New("subscriber not listening")
	}
}

// Flood provides broadcast communication over the hub.
//
// All listening subscribers will receive the message provided.
// The from field is usually set with the sender ID to be excluded
// from the broadcast. This allows subscribers to send without
// having to filter out their own messages.
//
// A count of skipped subscribers (excluding the sender) that are
// unable to receive the message is returned.
func (bh *ByteHub) Flood(from int, msg []byte) (skipped int) {
	bh.RLock()
	defer bh.RUnlock()

	for id, ch := range bh.subscribers {
		if id != from {
			select {
			case ch <- msg:
			default:
				skipped++
			}
		}
	}

	return
}

// Count returns the current number of subscriptions to the hub.
func (bh *ByteHub) Count() int {
	bh.RLock()
	defer bh.RUnlock()

	return len(bh.subscribers)
}

// Reset removes all subscriptions from the hub and allows it to be reused.
func (bh *ByteHub) Reset() {
	bh.Lock()
	defer bh.Unlock()

	bh.subscribers = make(map[int]chan<- []byte)
}
<commit_msg>manually unlock mutex without using defer<commit_after>// Package blstr (blaster? blast string?) is an ill-named library
// for easily setting up best-effort broadcast communication between
// goroutines.
//
// Synchronization and consistency is traded for speed and
// predictability of send operations.
//
// The best use case is where perfect consistency may not be
// necessary and some message loss on slow receivers is more tolerable
// than choking upstream senders.
package blstr

import (
	"errors"
	"sync"
)

// ByteHub implements a multiuser communication channel.
type ByteHub struct {
	subscribers map[int]chan<- []byte
	sync.RWMutex
}

// New returns a new ByteHub
func New() *ByteHub {
	return &ByteHub{subscribers: make(map[int]chan<- []byte)}
}

// Subscribe to be notified when broadcast events occur on the hub.
//
// Messages can be dropped if receiver stops listening, or if messages
// arrive faster than they can be consumed.
// Channels can be buffered to mitigate this to some extent.
func (bh *ByteHub) Subscribe(id int, ch chan<- []byte) error {
	bh.Lock()

	if _, ok := bh.subscribers[id]; ok {
		bh.Unlock()
		return errors.New("subscriber already exists")
	}

	bh.subscribers[id] = ch

	bh.Unlock()
	return nil
}

// Unsubscribe tries to remove the subscription by ID if it exists.
func (bh *ByteHub) Unsubscribe(id int) {
	bh.Lock()
	delete(bh.subscribers, id)
	bh.Unlock()
}

// Send provides unicast communication over the hub.
//
// An attempt is made to send a message to the specified subscriber.
// Errors are returned if the subscriber could not be found or
// is unable to receive the message.
func (bh *ByteHub) Send(to int, msg []byte) (err error) {
	bh.RLock()

	ch, ok := bh.subscribers[to]
	if !ok {
		bh.RUnlock()
		return errors.New("subscriber does not exist")
	}

	select {
	case ch <- msg:
	default:
		err = errors.New("subscriber not listening")
	}

	bh.RUnlock()
	return
}

// Flood provides broadcast communication over the hub.
//
// All listening subscribers will receive the message provided.
// The from field is usually set with the sender ID to be excluded
// from the broadcast. This allows subscribers to send without
// having to filter out their own messages.
//
// A count of skipped subscribers (excluding the sender) that are
// unable to receive the message is returned.
func (bh *ByteHub) Flood(from int, msg []byte) (skipped int) {
	bh.RLock()

	for id, ch := range bh.subscribers {
		if id != from {
			select {
			case ch <- msg:
			default:
				skipped++
			}
		}
	}

	bh.RUnlock()
	return
}

// Count returns the current number of subscriptions to the hub.
func (bh *ByteHub) Count() int {
	bh.RLock()
	l := len(bh.subscribers)
	bh.RUnlock()
	return l
}

// Reset removes all subscriptions from the hub and allows it to be reused.
func (bh *ByteHub) Reset() {
	bh.Lock()
	bh.subscribers = make(map[int]chan<- []byte)
	bh.Unlock()
}
<|endoftext|>


<commit_before>// Copyright 2015 Square 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 forecast

import (
	"fmt"
	"math"

	"github.com/square/metrics/api"
	"github.com/square/metrics/function"
)

// HoltWintersModel computes a Holt-Winters model for the given time series.
var ModelHoltWinters = function.MetricFunction{
	Name:         "forecast.model_holt_winters",
	MinArguments: 4, // series, period, time offset to train, length of training period
	MaxArguments: 4,
	Compute: func(context *function.EvaluationContext, arguments []function.Expression, groups function.Groups) (function.Value, error) {
		period, err := function.EvaluateToDuration(arguments[1], context)
		if err != nil {
			return nil, err
		}
		if period <= 0 {
			return nil, fmt.Errorf("forecast.holt_winters expected period to be positive") // TODO: use a structured error
		}
		when, err := function.EvaluateToDuration(arguments[3], context)
		if err != nil {
			return nil, err
		}
		length, err := function.EvaluateToDuration(arguments[3], context)
		if err != nil {
			return nil, err
		}
		// We need to perform a fetch of length 'length' offset 'when' for all this data.
		// Then we apply the Holt-Winters model to each of the resulting series.
		newContext := *context
		newContext.Timerange = newContext.Timerange.Shift(when).SelectLength(length)

		original, err := function.EvaluateToSeriesList(arguments[0], &newContext)
		if err != nil {
			return nil, err
		}

		result := api.SeriesList{
			Series:    make([]api.Timeseries, len(original.Series)),
			Timerange: context.Timerange,
			Name:      original.Name,
			Query:     fmt.Sprintf("forecast.holt_winters(%s, %s, %s, %s)", original.Query, period.String(), when.String(), length.String()),
		}
		slotTrainingStart := int(when / context.Timerange.Resolution())
		slotQueryStart := int(context.Timerange.Start() / context.Timerange.ResolutionMillis())
		for s := range result.Series {
			training := original.Series[s].Values
			model, err := HoltWintersMultiplicativeEstimate(training, int(period/context.Timerange.Resolution()))
			if err != nil {
				return nil, err // TODO: determine if there's a more graceful way to indicate the error - probably not
			}
			result.Series[s] = api.Timeseries{
				TagSet: original.Series[s].TagSet,
				Raw:    original.Series[s].Raw,
				Values: model.EstimateRange(slotQueryStart-slotTrainingStart, context.Timerange.Slots()),
			}
		}
		return result, nil
	},
}

type weighted struct {
	value  float64
	weight float64
	rate   float64
}

func (w *weighted) get() float64 {
	return w.value
}
func (w *weighted) observe(y float64) {
	if math.IsNaN(y) {
		w.skip()
		return
	}
	w.weight *= 1 - w.rate
	w.value = (w.value*w.weight + y*w.rate) / (w.weight + w.rate)
	w.weight += w.rate
}
func (w *weighted) boostAdd(dy float64) {
	if math.IsNaN(dy) {
		return
	}
	w.value += dy
}
func (w *weighted) skip() {
	w.weight *= 1 - w.rate
}
func newWeighted(rate float64) *weighted {
	return &weighted{
		value:  0,
		weight: 0,
		rate:   rate,
	}
}

type cycle struct {
	season []*weighted
}

func (c cycle) index(i int) int {
	return (i%len(c.season) + len(c.season)) % len(c.season)
}
func (c cycle) get(i int) float64 {
	return c.season[c.index(i)].get()
}
func (c cycle) observe(i int, v float64) {
	c.season[c.index(i)].observe(v)
}
func (c cycle) skip(i int) {
	c.season[c.index(i)].skip()
}
func newCycle(rate float64, n int) cycle {
	c := make([]*weighted, n)
	for i := range c {
		c[i] = newWeighted(rate)
	}
	return cycle{
		season: c,
	}
}

var SmoothHoltWinters = function.MetricFunction{
	Name:         "forecast.smooth_holt_winters",
	MinArguments: 5, // Series, period, level learning rate,  trend learning rate, seasonal learning rate,
	MaxArguments: 5,
	Compute: func(context *function.EvaluationContext, arguments []function.Expression, groups function.Groups) (function.Value, error) {
		period, err := function.EvaluateToDuration(arguments[1], context)
		if err != nil {
			return nil, err
		}
		levelLearningRate, err := function.EvaluateToScalar(arguments[2], context)
		if err != nil {
			return nil, err
		}
		trendLearningRate, err := function.EvaluateToScalar(arguments[3], context)
		if err != nil {
			return nil, err
		}
		seasonalLearningRate, err := function.EvaluateToScalar(arguments[4], context)
		if err != nil {
			return nil, err
		}

		samples := int(period / context.Timerange.Resolution())
		if samples <= 0 {
			return nil, fmt.Errorf("forecast.holt_winters_adaptive expects the period parameter to mean at least one slot") // TODO: use a structured error
		}

		seriesList, err := function.EvaluateToSeriesList(arguments[0], context)
		if err != nil {
			return nil, err
		}

		result := api.SeriesList{
			Series:    make([]api.Timeseries, len(seriesList.Series)),
			Timerange: context.Timerange,
			Name:      seriesList.Name,
			Query:     fmt.Sprintf("forecast.holt_winters_adaptive(%s, %s, %f, %f)", seriesList.Query, period.String(), seasonalLearningRate, trendLearningRate),
		}

		for seriesIndex := range result.Series {
			// This will be the result.
			estimate := make([]float64, len(seriesList.Series[seriesIndex].Values))

			level := newWeighted(levelLearningRate)
			trend := newWeighted(trendLearningRate)
			season := newCycle(seasonalLearningRate, samples)

			for i, y := range seriesList.Series[seriesIndex].Values {
				// Remember the old values.
				oldLevel := level.get()
				oldTrend := trend.get()
				oldSeason := season.get(i)

				// Update the level, by increasing it by the estimate slope
				level.boostAdd(oldTrend)
				// Then observing the new y [if y is NaN, this skips, as desired]
				level.observe(y / oldSeason) // observe the y/s non-seasonal value

				// Next, observe the trend- difference between this level and last.
				// If y is NaN, we want to skip instead of updating.
				if math.IsNaN(y) {
					trend.skip()
				} else {
					// Compare the new level against the old.
					trend.observe(level.get() - oldLevel)
				}

				// Lastly, the seasonal value is just y / (l+b) the non-seasonal component.
				// If y is NaN, this will be NaN too, causing it to skip (as desired).
				season.observe(i, y/(oldLevel+oldTrend))

				// Our estimate is the level times the seasonal component.
				estimate[i] = level.get() * season.get(i)
			}

			result.Series[seriesIndex] = api.Timeseries{
				TagSet: seriesList.Series[seriesIndex].TagSet,
				Raw:    seriesList.Series[seriesIndex].Raw,
				Values: estimate,
			}
		}

		return result, nil
	},
}
<commit_msg>update documentary strings<commit_after>// Copyright 2015 Square 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 forecast

import (
	"fmt"
	"math"

	"github.com/square/metrics/api"
	"github.com/square/metrics/function"
)

// HoltWintersModel computes a Holt-Winters model for the given time series.
var ModelHoltWinters = function.MetricFunction{
	Name:         "forecast.model_holt_winters",
	MinArguments: 4, // series, period, time offset to train, length of training period
	MaxArguments: 4,
	Compute: func(context *function.EvaluationContext, arguments []function.Expression, groups function.Groups) (function.Value, error) {
		period, err := function.EvaluateToDuration(arguments[1], context)
		if err != nil {
			return nil, err
		}
		if period <= 0 {
			return nil, fmt.Errorf("forecast.holt_winters expected period to be positive") // TODO: use a structured error
		}
		when, err := function.EvaluateToDuration(arguments[3], context)
		if err != nil {
			return nil, err
		}
		length, err := function.EvaluateToDuration(arguments[3], context)
		if err != nil {
			return nil, err
		}
		// We need to perform a fetch of length 'length' offset 'when' for all this data.
		// Then we apply the Holt-Winters model to each of the resulting series.
		newContext := context.Copy()
		newContext.Timerange = newContext.Timerange.Shift(when).SelectLength(length)

		original, err := function.EvaluateToSeriesList(arguments[0], &newContext)
		if err != nil {
			return nil, err
		}

		result := api.SeriesList{
			Series:    make([]api.Timeseries, len(original.Series)),
			Timerange: context.Timerange,
			Name:      original.Name,
			Query:     fmt.Sprintf("forecast.model_holt_winters(%s, %s, %s, %s)", original.Query, period.String(), when.String(), length.String()),
		}
		slotTrainingStart := int(when / context.Timerange.Resolution())
		slotQueryStart := int(context.Timerange.Start() / context.Timerange.ResolutionMillis())
		for s := range result.Series {
			training := original.Series[s].Values
			model, err := HoltWintersMultiplicativeEstimate(training, int(period/context.Timerange.Resolution()))
			if err != nil {
				return nil, err // TODO: determine if there's a more graceful way to indicate the error - probably not
			}
			result.Series[s] = api.Timeseries{
				TagSet: original.Series[s].TagSet,
				Raw:    original.Series[s].Raw,
				Values: model.EstimateRange(slotQueryStart-slotTrainingStart, context.Timerange.Slots()),
			}
		}
		return result, nil
	},
}

type weighted struct {
	value  float64
	weight float64
	rate   float64
}

func (w *weighted) get() float64 {
	return w.value
}
func (w *weighted) observe(y float64) {
	if math.IsNaN(y) {
		w.skip()
		return
	}
	w.weight *= 1 - w.rate
	w.value = (w.value*w.weight + y*w.rate) / (w.weight + w.rate)
	w.weight += w.rate
}
func (w *weighted) boostAdd(dy float64) {
	if math.IsNaN(dy) {
		return
	}
	w.value += dy
}
func (w *weighted) skip() {
	w.weight *= 1 - w.rate
}
func newWeighted(rate float64) *weighted {
	return &weighted{
		value:  0,
		weight: 0,
		rate:   rate,
	}
}

type cycle struct {
	season []*weighted
}

func (c cycle) index(i int) int {
	return (i%len(c.season) + len(c.season)) % len(c.season)
}
func (c cycle) get(i int) float64 {
	return c.season[c.index(i)].get()
}
func (c cycle) observe(i int, v float64) {
	c.season[c.index(i)].observe(v)
}
func (c cycle) skip(i int) {
	c.season[c.index(i)].skip()
}
func newCycle(rate float64, n int) cycle {
	c := make([]*weighted, n)
	for i := range c {
		c[i] = newWeighted(rate)
	}
	return cycle{
		season: c,
	}
}

var SmoothHoltWinters = function.MetricFunction{
	Name:         "forecast.smooth_holt_winters",
	MinArguments: 5, // Series, period, level learning rate,  trend learning rate, seasonal learning rate,
	MaxArguments: 5,
	Compute: func(context *function.EvaluationContext, arguments []function.Expression, groups function.Groups) (function.Value, error) {
		period, err := function.EvaluateToDuration(arguments[1], context)
		if err != nil {
			return nil, err
		}
		levelLearningRate, err := function.EvaluateToScalar(arguments[2], context)
		if err != nil {
			return nil, err
		}
		trendLearningRate, err := function.EvaluateToScalar(arguments[3], context)
		if err != nil {
			return nil, err
		}
		seasonalLearningRate, err := function.EvaluateToScalar(arguments[4], context)
		if err != nil {
			return nil, err
		}

		samples := int(period / context.Timerange.Resolution())
		if samples <= 0 {
			return nil, fmt.Errorf("forecast.smooth_holt_winters expects the period parameter to mean at least one slot") // TODO: use a structured error
		}

		seriesList, err := function.EvaluateToSeriesList(arguments[0], context)
		if err != nil {
			return nil, err
		}

		result := api.SeriesList{
			Series:    make([]api.Timeseries, len(seriesList.Series)),
			Timerange: context.Timerange,
			Name:      seriesList.Name,
			Query:     fmt.Sprintf("forecast.smooth_holt_winters(%s, %s, %f, %f)", seriesList.Query, period.String(), seasonalLearningRate, trendLearningRate),
		}

		for seriesIndex := range result.Series {
			// This will be the result.
			estimate := make([]float64, len(seriesList.Series[seriesIndex].Values))

			level := newWeighted(levelLearningRate)
			trend := newWeighted(trendLearningRate)
			season := newCycle(seasonalLearningRate, samples)

			for i, y := range seriesList.Series[seriesIndex].Values {
				// Remember the old values.
				oldLevel := level.get()
				oldTrend := trend.get()
				oldSeason := season.get(i)

				// Update the level, by increasing it by the estimate slope
				level.boostAdd(oldTrend)
				// Then observing the new y [if y is NaN, this skips, as desired]
				level.observe(y / oldSeason) // observe the y/s non-seasonal value

				// Next, observe the trend- difference between this level and last.
				// If y is NaN, we want to skip instead of updating.
				if math.IsNaN(y) {
					trend.skip()
				} else {
					// Compare the new level against the old.
					trend.observe(level.get() - oldLevel)
				}

				// Lastly, the seasonal value is just y / (l+b) the non-seasonal component.
				// If y is NaN, this will be NaN too, causing it to skip (as desired).
				season.observe(i, y/(oldLevel+oldTrend))

				// Our estimate is the level times the seasonal component.
				estimate[i] = level.get() * season.get(i)
			}

			result.Series[seriesIndex] = api.Timeseries{
				TagSet: seriesList.Series[seriesIndex].TagSet,
				Raw:    seriesList.Series[seriesIndex].Raw,
				Values: estimate,
			}
		}

		return result, nil
	},
}
<|endoftext|>