<commit_before>package templates

func init() {
	registerTemplate("node.html", `
{{template "header.html" .}}
{{$g := .GraphOptions}}
{{$node := .Node}}
{{$numTopics := .TopicStats | len}}
{{$numMessages := .NumMessages | commafy}}
{{$numClients := .NumClients}}

Nodes /

{{$node}}

{{if $g.Enabled}}

GC Pressure

GC Pause Percentiles

GC Runs

Heap Objects In-Use

{{end}}

    {{if not .TopicStats}}

Notice
 No topics exist on this node.

    {{else}}
    {{range $t := .TopicStats}}

                {{$numTopics}} Topics

                {{$numMessages}} Messages
                {{$numClients}} Clients

                Topic
                Depth
                Memory + Disk

                Messages
                Channels

                ✘ {{.TopicName}}

                    {{if $g.Enabled}}{{end}}
                    {{.Depth | commafy}}
                {{.MemoryDepth | commafy}} + {{.BackendDepth | commafy}}

                    {{if $g.Enabled}}{{end}}
                    {{.MessageCount | commafy}}

                {{.ChannelCount}}

            {{if not .Channels}}

Notice
 No channels exist for this topic.

            {{else}}
            {{range $c := .Channels}}

                Channel
                Depth
                Memory + Disk
                In-Flight
                Deferred
                Requeued
                Timed Out
                Messages
                Connections

                    {{$c.ChannelName}}
                    {{if $c.Paused}}paused{{end}}

                    {{if $g.Enabled}}{{end}}
                    {{$c.Depth | commafy}}
                {{$c.MemoryDepth | commafy}} + {{$c.BackendDepth | commafy}}
                {{$c.InFlightCount | commafy}}
                {{$c.DeferredCount | commafy}}
                {{$c.RequeueCount | commafy}}
                {{$c.TimeoutCount | commafy}}
                {{$c.MessageCount | commafy}}

                    {{if $g.Enabled}}{{end}}
                    {{$c.ClientCount}}

            {{if not .Clients}}

Notice
No clients connected to this channel.

            {{else}}

                Client Host
                Protocol
                Attributes
                NSQd Host
                In-Flight
                Ready Count
                Finished
                Requeued
                Messages
                Connected

            {{range .Clients}}

                {{.ClientIdentifier}}
                {{.ClientVersion}} {{if .HasUserAgent}}({{.ClientUserAgent}}){{end}}

                  {{if .HasSampleRate}}
                  Sampled {{.SampleRate}}%
                  {{end}}
                  {{if .TLS}}
                  TLS
                  {{end}}
                  {{if .Deflate}}
                  Delfate
                  {{end}}
                  {{if .Snappy}}
                  Snappy
                  {{end}}

                {{.HostAddress}}
                {{.InFlightCount | commafy}}
                {{.ReadyCount | commafy}}
                {{.FinishCount | commafy}}
                {{.RequeueCount | commafy}}
                {{.MessageCount | commafy}}
                {{.ConnectedDuration}}

            {{end}}
            {{end}}
        {{end}}
        {{end}}

    {{end}}
    {{end}}

{{template "js.html" .}}
{{template "footer.html" .}}
`)
}
<commit_msg>nsqadmin: adjust columns on node page<commit_after>package templates

func init() {
	registerTemplate("node.html", `
{{template "header.html" .}}
{{$g := .GraphOptions}}
{{$node := .Node}}
{{$numTopics := .TopicStats | len}}
{{$numMessages := .NumMessages | commafy}}
{{$numClients := .NumClients}}

Nodes /

{{$node}}

{{if $g.Enabled}}

GC Pressure

GC Pause Percentiles

GC Runs

Heap Objects In-Use

{{end}}

    {{if not .TopicStats}}

Notice
 No topics exist on this node.

    {{else}}
    {{range $t := .TopicStats}}

                {{$numTopics}} Topics

                {{$numMessages}} Messages
                {{$numClients}} Clients

                Topic
                Depth
                Memory + Disk

                Messages
                Channels

                ✘ {{.TopicName}}

                    {{if $g.Enabled}}{{end}}
                    {{.Depth | commafy}}
                {{.MemoryDepth | commafy}} + {{.BackendDepth | commafy}}

                    {{if $g.Enabled}}{{end}}
                    {{.MessageCount | commafy}}

                {{.ChannelCount}}

            {{if not .Channels}}

Notice
 No channels exist for this topic.

            {{else}}
            {{range $c := .Channels}}

                Channel
                Depth
                Memory + Disk
                In-Flight
                Deferred
                Requeued
                Timed Out
                Messages
                Connections

                    {{$c.ChannelName}}
                    {{if $c.Paused}}paused{{end}}

                    {{if $g.Enabled}}{{end}}
                    {{$c.Depth | commafy}}
                {{$c.MemoryDepth | commafy}} + {{$c.BackendDepth | commafy}}
                {{$c.InFlightCount | commafy}}
                {{$c.DeferredCount | commafy}}
                {{$c.RequeueCount | commafy}}
                {{$c.TimeoutCount | commafy}}
                {{$c.MessageCount | commafy}}

                    {{if $g.Enabled}}{{end}}
                    {{$c.ClientCount}}

            {{if not .Clients}}

Notice
No clients connected to this channel.

            {{else}}

                Client Host
                Protocol
                Attributes
                NSQd Host
                In-Flight
                Ready Count
                Finished
                Requeued
                Messages
                Connected

            {{range .Clients}}

                {{.ClientIdentifier}}
                {{.ClientVersion}} {{if .HasUserAgent}}({{.ClientUserAgent}}){{end}}

                  {{if .HasSampleRate}}
                  Sampled {{.SampleRate}}%
                  {{end}}
                  {{if .TLS}}
                  TLS
                  {{end}}
                  {{if .Deflate}}
                  Delfate
                  {{end}}
                  {{if .Snappy}}
                  Snappy
                  {{end}}

                {{.HostAddress}}
                {{.InFlightCount | commafy}}
                {{.ReadyCount | commafy}}
                {{.FinishCount | commafy}}
                {{.RequeueCount | commafy}}
                {{.MessageCount | commafy}}
                {{.ConnectedDuration}}

            {{end}}
            {{end}}
        {{end}}
        {{end}}

    {{end}}
    {{end}}

{{template "js.html" .}}
{{template "footer.html" .}}
`)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"runtime"

	"github.com/containernetworking/cni/pkg/ipam"
	"github.com/containernetworking/cni/pkg/ns"
	"github.com/containernetworking/cni/pkg/skel"
	"github.com/containernetworking/cni/pkg/types"
	"github.com/vishvananda/netlink"
)

type NetConf struct {
	types.NetConf
	Master string `json:"master"`
	MAC    string `json:"mac"`
	VF     int    `json:"vf"`
	Vlan   int    `json:"vlan"`
}

func init() {
	// this ensures that main runs only on main thread (thread group leader).
	// since namespace ops (unshare, setns) are done for a single thread, we
	// must ensure that the goroutine does not jump from OS thread to thread
	runtime.LockOSThread()
}

func loadConf(bytes []byte) (*NetConf, error) {
	n := &NetConf{}
	if err := json.Unmarshal(bytes, n); err != nil {
		return nil, fmt.Errorf("failed to load netconf: %v", err)
	}
	if n.Master == "" {
		return nil, fmt.Errorf(`"master" field is required. It specifies the host interface name to virtualize`)
	}
	return n, nil
}

func setupVF(conf *NetConf, ifName string, netns ns.NetNS) error {

	masterName := conf.Master
	vfIdx := conf.VF

	m, err := netlink.LinkByName(masterName)
	if err != nil {
		return fmt.Errorf("failed to lookup master %q: %v", conf.Master, err)
	}

	vfDir := fmt.Sprintf("/sys/class/net/%s/device/virtfn%d/net", masterName, vfIdx)
	if _, err := os.Lstat(vfDir); err != nil {
		return err
	}

	infos, err := ioutil.ReadDir(vfDir)
	if err != nil {
		return err
	}

	if len(infos) != 1 {
		return fmt.Errorf("Mutiple network devices in directory %s", vfDir)
	}

	// VF NIC name
	vfDevName := infos[0].Name()
	vfDev, err := netlink.LinkByName(vfDevName)
	if err != nil {
		return fmt.Errorf("failed to lookup vf device %q: %v", vfDevName, err)
	}

	// set hardware address
	if conf.MAC != "" {
		macAddr, err := net.ParseMAC(conf.MAC)
		if err != nil {
			return err
		}
		if err = netlink.LinkSetVfHardwareAddr(m, conf.VF, macAddr); err != nil {
			return fmt.Errorf("failed to set vf %d macaddress: %v", conf.VF, err)
		}
	}

	if conf.Vlan != 0 {
		if err = netlink.LinkSetVfVlan(m, conf.VF, conf.Vlan); err != nil {
			return fmt.Errorf("failed to set vf %d vlan: %v", conf.VF, err)
		}
	}

	if err = netlink.LinkSetUp(vfDev); err != nil {
		return fmt.Errorf("failed to setup vf %d device: %v", conf.VF, err)
	}

	// move VF device to ns
	if err = netlink.LinkSetNsFd(vfDev, int(netns.Fd())); err != nil {
		return fmt.Errorf("failed to move vf %d to netns: %v", conf.VF, err)
	}

	return netns.Do(func(_ ns.NetNS) error {
		err := renameLink(vfDevName, ifName)
		if err != nil {
			return fmt.Errorf("failed to rename vf %d device %q to %q: %v", conf.VF, vfDevName, ifName, err)
		}
		return nil
	})
}

func releaseVF(conf *NetConf, ifName string, netns ns.NetNS) error {
	initns, err := ns.GetCurrentNS()
	if err != nil {
		return fmt.Errorf("failed to get init netns: %v", err)
	}

	if err = netns.Set(); err != nil {
		return fmt.Errorf("failed to enter netns %q: %v", netns, err)
	}

	// get VF device
	vfDev, err := netlink.LinkByName(ifName)
	if err != nil {
		return fmt.Errorf("failed to lookup vf %d device %q: %v", conf.VF, ifName, err)
	}

	// device name in init netns
	index := vfDev.Attrs().Index
	devName := fmt.Sprintf("dev%d", index)

	// shutdown VF device
	if err = netlink.LinkSetDown(vfDev); err != nil {
		return fmt.Errorf("failed to down vf %d device: %v", conf.VF, err)
	}

	// rename VF device
	err = renameLink(ifName, devName)
	if err != nil {
		return fmt.Errorf("failed to rename vf %d evice %q to %q: %v", conf.VF, ifName, devName, err)
	}

	// move VF device to init netns
	if err = netlink.LinkSetNsFd(vfDev, int(initns.Fd())); err != nil {
		return fmt.Errorf("failed to move vf %d to init netns: %v", conf.VF, err)
	}

	return nil
}

func cmdAdd(args *skel.CmdArgs) error {
	n, err := loadConf(args.StdinData)
	if err != nil {
		return err
	}

	netns, err := ns.GetNS(args.Netns)
	if err != nil {
		return fmt.Errorf("failed to open netns %q: %v", args.Netns, err)
	}
	defer netns.Close()

	if err = setupVF(n, args.IfName, netns); err != nil {
		return err
	}

	// run the IPAM plugin and get back the config to apply
	result, err := ipam.ExecAdd(n.IPAM.Type, args.StdinData)
	if err != nil {
		return err
	}
	if result.IP4 == nil {
		return errors.New("IPAM plugin returned missing IPv4 config")
	}

	err = netns.Do(func(_ ns.NetNS) error {
		return ipam.ConfigureIface(args.IfName, result)
	})
	if err != nil {
		return err
	}

	result.DNS = n.DNS
	return result.Print()
}

func cmdDel(args *skel.CmdArgs) error {
	n, err := loadConf(args.StdinData)
	if err != nil {
		return err
	}

	netns, err := ns.GetNS(args.Netns)
	if err != nil {
		return fmt.Errorf("failed to open netns %q: %v", netns, err)
	}
	defer netns.Close()

	if err = releaseVF(n, args.IfName, netns); err != nil {
		return err
	}

	err = ipam.ExecDel(n.IPAM.Type, args.StdinData)
	if err != nil {
		return err
	}

	return nil
}

func renameLink(curName, newName string) error {
	link, err := netlink.LinkByName(curName)
	if err != nil {
		return err
	}

	return netlink.LinkSetName(link, newName)
}

func main() {
	skel.PluginMain(cmdAdd, cmdDel)
}
<commit_msg>Change 'NetConf.VF' from 'int' to '*int'<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"runtime"

	"github.com/containernetworking/cni/pkg/ipam"
	"github.com/containernetworking/cni/pkg/ns"
	"github.com/containernetworking/cni/pkg/skel"
	"github.com/containernetworking/cni/pkg/types"
	"github.com/vishvananda/netlink"
)

type NetConf struct {
	types.NetConf
	Master string `json:"master"`
	MAC    string `json:"mac"`
	VF     *int   `json:"vf"`
	Vlan   int    `json:"vlan"`
}

func init() {
	// this ensures that main runs only on main thread (thread group leader).
	// since namespace ops (unshare, setns) are done for a single thread, we
	// must ensure that the goroutine does not jump from OS thread to thread
	runtime.LockOSThread()
}

func loadConf(bytes []byte) (*NetConf, error) {
	n := &NetConf{}
	if err := json.Unmarshal(bytes, n); err != nil {
		return nil, fmt.Errorf("failed to load netconf: %v", err)
	}
	if n.Master == "" {
		return nil, fmt.Errorf(`"master" field is required. It specifies the host interface name to virtualize`)
	}
	return n, nil
}

func setupVF(conf *NetConf, ifName string, netns ns.NetNS) error {
	vfIdx := 0
	masterName := conf.Master
	if conf.VF != nil {
		vfIdx = *conf.VF
	}
	// TODO: if conf.VF == nil, alloc vf randomly

	m, err := netlink.LinkByName(masterName)
	if err != nil {
		return fmt.Errorf("failed to lookup master %q: %v", conf.Master, err)
	}

	vfDir := fmt.Sprintf("/sys/class/net/%s/device/virtfn%d/net", masterName, vfIdx)
	if _, err := os.Lstat(vfDir); err != nil {
		return err
	}

	infos, err := ioutil.ReadDir(vfDir)
	if err != nil {
		return err
	}

	if len(infos) != 1 {
		return fmt.Errorf("no network devices in directory %s", vfDir)
	}

	// VF NIC name
	vfDevName := infos[0].Name()
	vfDev, err := netlink.LinkByName(vfDevName)
	if err != nil {
		return fmt.Errorf("failed to lookup vf device %q: %v", vfDevName, err)
	}

	// set hardware address
	if conf.MAC != "" {
		macAddr, err := net.ParseMAC(conf.MAC)
		if err != nil {
			return err
		}
		if err = netlink.LinkSetVfHardwareAddr(m, vfIdx, macAddr); err != nil {
			return fmt.Errorf("failed to set vf %d macaddress: %v", vfIdx, err)
		}
	}

	if conf.Vlan != 0 {
		if err = netlink.LinkSetVfVlan(m, vfIdx, conf.Vlan); err != nil {
			return fmt.Errorf("failed to set vf %d vlan: %v", vfIdx, err)
		}
	}

	if err = netlink.LinkSetUp(vfDev); err != nil {
		return fmt.Errorf("failed to setup vf %d device: %v", vfIdx, err)
	}

	// move VF device to ns
	if err = netlink.LinkSetNsFd(vfDev, int(netns.Fd())); err != nil {
		return fmt.Errorf("failed to move vf %d to netns: %v", vfIdx, err)
	}

	return netns.Do(func(_ ns.NetNS) error {
		err := renameLink(vfDevName, ifName)
		if err != nil {
			return fmt.Errorf("failed to rename vf %d device %q to %q: %v", vfIdx, vfDevName, ifName, err)
		}
		return nil
	})
}

func releaseVF(conf *NetConf, ifName string, netns ns.NetNS) error {
	vfIdx := 0
	if conf.VF != nil {
		vfIdx = *conf.VF
	}

	initns, err := ns.GetCurrentNS()
	if err != nil {
		return fmt.Errorf("failed to get init netns: %v", err)
	}

	if err = netns.Set(); err != nil {
		return fmt.Errorf("failed to enter netns %q: %v", netns, err)
	}

	// get VF device
	vfDev, err := netlink.LinkByName(ifName)
	if err != nil {
		return fmt.Errorf("failed to lookup vf %d device %q: %v", vfIdx, ifName, err)
	}

	// device name in init netns
	index := vfDev.Attrs().Index
	devName := fmt.Sprintf("dev%d", index)

	// shutdown VF device
	if err = netlink.LinkSetDown(vfDev); err != nil {
		return fmt.Errorf("failed to down vf %d device: %v", vfIdx, err)
	}

	// rename VF device
	err = renameLink(ifName, devName)
	if err != nil {
		return fmt.Errorf("failed to rename vf %d evice %q to %q: %v", vfIdx, ifName, devName, err)
	}

	// move VF device to init netns
	if err = netlink.LinkSetNsFd(vfDev, int(initns.Fd())); err != nil {
		return fmt.Errorf("failed to move vf %d to init netns: %v", vfIdx, err)
	}

	return nil
}

func cmdAdd(args *skel.CmdArgs) error {
	n, err := loadConf(args.StdinData)
	if err != nil {
		return err
	}

	netns, err := ns.GetNS(args.Netns)
	if err != nil {
		return fmt.Errorf("failed to open netns %q: %v", args.Netns, err)
	}
	defer netns.Close()

	if err = setupVF(n, args.IfName, netns); err != nil {
		return err
	}

	// run the IPAM plugin and get back the config to apply
	result, err := ipam.ExecAdd(n.IPAM.Type, args.StdinData)
	if err != nil {
		return err
	}
	if result.IP4 == nil {
		return errors.New("IPAM plugin returned missing IPv4 config")
	}

	err = netns.Do(func(_ ns.NetNS) error {
		return ipam.ConfigureIface(args.IfName, result)
	})
	if err != nil {
		return err
	}

	result.DNS = n.DNS
	return result.Print()
}

func cmdDel(args *skel.CmdArgs) error {
	n, err := loadConf(args.StdinData)
	if err != nil {
		return err
	}

	netns, err := ns.GetNS(args.Netns)
	if err != nil {
		return fmt.Errorf("failed to open netns %q: %v", netns, err)
	}
	defer netns.Close()

	if err = releaseVF(n, args.IfName, netns); err != nil {
		return err
	}

	err = ipam.ExecDel(n.IPAM.Type, args.StdinData)
	if err != nil {
		return err
	}

	return nil
}

func renameLink(curName, newName string) error {
	link, err := netlink.LinkByName(curName)
	if err != nil {
		return err
	}

	return netlink.LinkSetName(link, newName)
}

func main() {
	skel.PluginMain(cmdAdd, cmdDel)
}
<|endoftext|>


<commit_before>// The presence package is intended as a replacement for zookeeper ephemeral
// nodes; the primary difference is that node timeout is unrelated to session
// timeout, and this allows us to restart a presence-enabled process "silently"
// (from the perspective of the rest of the system) without dealing with the
// complication of session re-establishment.

package presence

import (
	"fmt"
	zk "launchpad.net/gozk/zookeeper"
	"launchpad.net/juju-core/log"
	"launchpad.net/tomb"
	"time"
)

// changeNode wraps a zookeeper node and can induce watches on that node to fire.
type changeNode struct {
	conn    *zk.Conn
	path    string
	content string
}

// change sets the zookeeper node's content (creating it if it doesn't exist) and
// returns the node's new MTime.
func (n *changeNode) change() (mtime time.Time, err error) {
	stat, err := n.conn.Set(n.path, n.content, -1)
	if zk.IsError(err, zk.ZNONODE) {
		_, err = n.conn.Create(n.path, n.content, 0, zk.WorldACL(zk.PERM_ALL))
		if err == nil || zk.IsError(err, zk.ZNODEEXISTS) {
			// *Someone* created the node anyway; just try again.
			return n.change()
		}
	}
	if err != nil {
		return
	}
	return stat.MTime(), nil
}

// Pinger periodically updates a node in zookeeper while it is running.
type Pinger struct {
	conn      *zk.Conn
	target    changeNode
	period    time.Duration
	tomb      tomb.Tomb
	stayAlive chan struct{}
}

// StartPinger returns a new Pinger that refreshes the content of path periodically.
func StartPinger(conn *zk.Conn, path string, period time.Duration) (*Pinger, error) {
	target := changeNode{conn, path, period.String()}
	_, err := target.change()
	if err != nil {
		return nil, err
	}
	p := &Pinger{
		conn:      conn,
		target:    target,
		period:    period,
		stayAlive: make(chan struct{}),
	}
	go p.loop()
	return p, nil
}

// loop calls change on p.target every p.period nanoseconds until p is stopped.
func (p *Pinger) loop() {
	defer p.finish()
	var wait time.Duration
	var then time.Time
	for {
		select {
		case <-p.tomb.Dying():
			log.Debugf("presence: %s pinger died after %s wait", p.target.path, time.Now().Sub(then))
			return
		case now := <-time.After(wait):
			mtime, err := p.target.change()
			if err != nil {
				p.tomb.Kill(err)
				return
			}
			next := then.Add(p.period)
			wait = next.Sub(now)
			log.Debugf("presence: pinged %s at %s; waiting %s", p.target.path, mtime, wait)
			then = now
		}
	}
}

// finish marks the Pinger as finally dead; it may also trigger a final write
// to the target node, if the Pinger has been cleanly shut down via the Stop
// method.
func (p *Pinger) finish() {
	if p.tomb.Err() == nil {
		select {
		case <-p.stayAlive:
			if _, err := p.target.change(); err != nil {
				p.tomb.Kill(err)
			}
		default:
		}
	}
	p.tomb.Done()
}

// Dying returns a channel that will be closed when the Pinger is no longer
// operating.
func (p *Pinger) Dying() <-chan struct{} {
	return p.tomb.Dying()
}

// Stop stops updating the node; AliveW watches will not notice any change
// until they time out. If the pinger is shut down without errors, a final
// write to the target node will occur, to ensure that remote watchers take
// as long as possible to detect the Pinger's absence.
func (p *Pinger) Stop() error {
	close(p.stayAlive)
	p.tomb.Kill(nil)
	return p.tomb.Wait()
}

// Kill stops updating and deletes the node, causing any AliveW watches
// to observe its departure (almost) immediately.
func (p *Pinger) Kill() (err error) {
	p.tomb.Kill(nil)
	if err = p.tomb.Wait(); err == nil {
		if err = p.conn.Delete(p.target.path, -1); zk.IsError(err, zk.ZNONODE) {
			return nil
		}
	}
	return
}

// node represents the state of a presence node from a watcher's perspective.
type node struct {
	conn    *zk.Conn
	path    string
	alive   bool
	timeout time.Duration
}

// setState sets the current values of n.alive and n.timeout, given the
// content and stat of the zookeeper node (which must exist). firstTime
// should be true iff n.setState has not already been called.
func (n *node) setState(content string, stat *zk.Stat, firstTime bool) error {
	// Always check and reset timeout; it could change.
	period, err := time.ParseDuration(content)
	if err != nil {
		return fmt.Errorf("%s presence node has bad data: %q", n.path, content)
	}
	n.timeout = period * 3
	if firstTime {
		clock := changeNode{n.conn, "/clock", ""}
		now, err := clock.change()
		if err != nil {
			return err
		}
		mtime := stat.MTime()
		delay := now.Sub(mtime)
		n.alive = delay < n.timeout
		log.Debugf(`
presence: initial diagnosis of %s
  now (zk)          %s
  last write (zk)   %s
  apparent delay    %s
  timeout           %s
  alive             %t
`[1:], n.path, now, mtime, delay, n.timeout, n.alive)
	} else {
		// If this method is not being run for the first time, we know that
		// the node has just changed, so we know that it's alive and there's
		// no need to check for staleness with the clock node.
		n.alive = true
	}
	return nil
}

// update reads from ZooKeeper the current values of n.alive and n.timeout.
func (n *node) update() error {
	content, stat, err := n.conn.Get(n.path)
	if zk.IsError(err, zk.ZNONODE) {
		n.alive = false
		n.timeout = 0
		return nil
	} else if err != nil {
		return err
	}
	return n.setState(content, stat, true)
}

// updateW reads from ZooKeeper the current values of n.alive and n.timeout,
// and returns a ZooKeeper watch that will be notified of data or existence
// changes. firstTime should be true iff n.updateW has not already been called.
func (n *node) updateW(firstTime bool) (<-chan zk.Event, error) {
	content, stat, watch, err := n.conn.GetW(n.path)
	if zk.IsError(err, zk.ZNONODE) {
		n.alive = false
		n.timeout = 0
		stat, watch, err = n.conn.ExistsW(n.path)
		if err != nil {
			return nil, err
		}
		if stat != nil {
			// Someone *just* created the node; try again.
			return n.updateW(firstTime)
		}
		return watch, nil
	} else if err != nil {
		return nil, err
	}
	return watch, n.setState(content, stat, firstTime)
}

// waitDead sends false to watch when the node is deleted, or when it has
// not been updated recently enough to still qualify as alive. zkWatch must
// be observing changes on the existent node at n.path.
func (n *node) waitDead(zkWatch <-chan zk.Event, watch chan bool) {
	for n.alive {
		select {
		case <-time.After(n.timeout):
			if err := n.update(); err != nil {
				close(watch)
				return
			}
		case event := <-zkWatch:
			if !event.Ok() {
				close(watch)
				return
			}
			switch event.Type {
			case zk.EVENT_DELETED:
				n.alive = false
			case zk.EVENT_CHANGED:
				var err error
				zkWatch, err = n.updateW(false)
				if err != nil {
					close(watch)
					return
				}
			default:
				panic(fmt.Errorf("Unexpected event: %v", event))
			}
		}
	}
	watch <- false
}

// waitAlive sends true to watch when the node is changed or created. zkWatch
// must be observing changes on the stale/nonexistent node at n.path.
func (n *node) waitAlive(zkWatch <-chan zk.Event, watch chan bool) {
	for !n.alive {
		event := <-zkWatch
		if !event.Ok() {
			close(watch)
			return
		}
		switch event.Type {
		case zk.EVENT_CREATED, zk.EVENT_CHANGED:
			n.alive = true
		case zk.EVENT_DELETED:
			// The pinger is still dead (just differently dead); start a new watch.
			var err error
			zkWatch, err = n.updateW(false)
			if err != nil {
				close(watch)
				return
			}
		default:
			panic(fmt.Errorf("Unexpected event: %v", event))
		}
	}
	watch <- true
}

// Alive returns whether the Pinger at the given path seems to be alive.
func Alive(conn *zk.Conn, path string) (bool, error) {
	n := &node{conn: conn, path: path}
	err := n.update()
	return n.alive, err
}

// AliveW returns whether the Pinger at the given path seems to be alive.
// It also returns a channel that will receive the new status when it changes.
// If an error is encountered after AliveW returns, the channel will be closed.
func AliveW(conn *zk.Conn, path string) (bool, <-chan bool, error) {
	n := &node{conn: conn, path: path}
	zkWatch, err := n.updateW(true)
	if err != nil {
		return false, nil, err
	}
	alive := n.alive
	watch := make(chan bool)
	if alive {
		go n.waitDead(zkWatch, watch)
	} else {
		go n.waitAlive(zkWatch, watch)
	}
	return alive, watch, nil
}

// WaitAlive blocks until the node at the given path
// has been recently pinged or a timeout occurs.
func WaitAlive(conn *zk.Conn, path string, timeout time.Duration) error {
	alive, watch, err := AliveW(conn, path)
	if err != nil {
		return err
	}
	if alive {
		return nil
	}
	select {
	case alive, ok := <-watch:
		if !ok {
			return fmt.Errorf("presence: channel closed while waiting")
		}
		if !alive {
			return fmt.Errorf("presence: alive watch misbehaved while waiting")
		}
	case <-time.After(timeout):
		return fmt.Errorf("presence: still not alive after timeout")
	}
	return nil
}
<commit_msg>more clarity, better docs, in presence.go<commit_after>// The presence package is intended as a replacement for zookeeper ephemeral
// nodes; the primary difference is that node timeout is unrelated to session
// timeout, and this allows us to restart a presence-enabled process "silently"
// (from the perspective of the rest of the system) without dealing with the
// complication of session re-establishment.

package presence

import (
	"fmt"
	zk "launchpad.net/gozk/zookeeper"
	"launchpad.net/juju-core/log"
	"launchpad.net/tomb"
	"time"
)

// changeNode wraps a zookeeper node and can induce watches on that node to fire.
type changeNode struct {
	conn    *zk.Conn
	path    string
	content string
}

// change sets the zookeeper node's content (creating it if it doesn't exist) and
// returns the node's new MTime.
func (n *changeNode) change() (mtime time.Time, err error) {
	stat, err := n.conn.Set(n.path, n.content, -1)
	if zk.IsError(err, zk.ZNONODE) {
		_, err = n.conn.Create(n.path, n.content, 0, zk.WorldACL(zk.PERM_ALL))
		if err == nil || zk.IsError(err, zk.ZNODEEXISTS) {
			// *Someone* created the node anyway; just try again.
			return n.change()
		}
	}
	if err != nil {
		return
	}
	return stat.MTime(), nil
}

// Pinger periodically updates a node in zookeeper while it is running.
type Pinger struct {
	conn      *zk.Conn
	tomb      tomb.Tomb
	target    changeNode
	period    time.Duration
	lastPing  time.Time
	stayAlive chan struct{}
}

// StartPinger returns a new Pinger that refreshes the content of path periodically.
func StartPinger(conn *zk.Conn, path string, period time.Duration) (*Pinger, error) {
	target := changeNode{conn, path, period.String()}
	_, err := target.change()
	if err != nil {
		return nil, err
	}
	p := &Pinger{
		conn:      conn,
		target:    target,
		period:    period,
		lastPing:  time.Now(),
		stayAlive: make(chan struct{}),
	}
	go p.loop()
	return p, nil
}

// loop calls change on p.target every p.period nanoseconds until p is stopped.
func (p *Pinger) loop() {
	defer p.finish()
	tick := p.getTick()
	for {
		select {
		case <-p.tomb.Dying():
			log.Debugf("presence: %s pinger died after %s wait", p.target.path, time.Now().Sub(p.lastPing))
			return
		case now := <-tick:
			log.Debugf("presence: %s pinger awakened after %s wait", p.target.path, time.Now().Sub(p.lastPing))
			p.lastPing = now
			mtime, err := p.target.change()
			if err != nil {
				p.tomb.Kill(err)
				return
			}
			log.Debugf("presence: wrote to %s at (zk) %s", p.target.path, mtime)
			tick = p.getTick()
		}
	}
}

// getTick returns a channel that will receive an event when the pinger is next
// due to fire.
func (p *Pinger) getTick() <-chan time.Time {
	next := p.lastPing.Add(p.period)
	wait := next.Sub(time.Now())
	if wait <= 0 {
		wait = 0
	}
	log.Debugf("presence: anticipating ping of %s in %s", p.target.path, wait)
	return time.After(wait)
}

// finish marks the Pinger as finally dead; it may also trigger a final write
// to the target node, if the Pinger has been cleanly shut down via the Stop
// method.
func (p *Pinger) finish() {
	if p.tomb.Err() == nil {
		select {
		case <-p.stayAlive:
			if _, err := p.target.change(); err != nil {
				p.tomb.Kill(err)
			}
		default:
		}
	}
	p.tomb.Done()
}

// Dying returns a channel that will be closed when the Pinger is no longer
// operating.
func (p *Pinger) Dying() <-chan struct{} {
	return p.tomb.Dying()
}

// Stop stops updating the node; AliveW watches will not notice any change
// until they time out. If the pinger is shut down without errors, a final
// write to the target node will occur, to ensure that remote watchers take
// as long as possible to detect the Pinger's absence.
func (p *Pinger) Stop() error {
	close(p.stayAlive)
	p.tomb.Kill(nil)
	return p.tomb.Wait()
}

// Kill stops updating and deletes the node, causing any AliveW watches
// to observe its departure (almost) immediately.
func (p *Pinger) Kill() (err error) {
	p.tomb.Kill(nil)
	if err = p.tomb.Wait(); err == nil {
		if err = p.conn.Delete(p.target.path, -1); zk.IsError(err, zk.ZNONODE) {
			return nil
		}
	}
	return
}

// node represents the state of a presence node from a watcher's perspective.
type node struct {
	conn    *zk.Conn
	path    string
	alive   bool
	timeout time.Duration
}

// setState sets the current values of n.alive and n.timeout, given the
// content and stat of the zookeeper node (which must exist). firstTime
// should be true iff n.setState has not already been called.
func (n *node) setState(content string, stat *zk.Stat, firstTime bool) error {
	// Always check and reset timeout; it could change.
	period, err := time.ParseDuration(content)
	if err != nil {
		return fmt.Errorf("%s presence node has bad data: %q", n.path, content)
	}
	// Worst observed combination of GC pauses and scheduler weirdness led to
	// a gap between "500ms" pings of >1000ms; timeout should comfortably beat
	// that. This is largely an artifact of the notably short periods used in
	// tests; actual distributed Pingers should use longer periods in order to
	// compensate for network flakiness... and to allow a comfortable buffer
	// for Pinger re-establishment by a replacement process when (eg when
	// upgrading agent binaries).
	n.timeout = period * 3
	if firstTime {
		clock := changeNode{n.conn, "/clock", ""}
		now, err := clock.change()
		if err != nil {
			return err
		}
		mtime := stat.MTime()
		delay := now.Sub(mtime)
		n.alive = delay < n.timeout
		log.Debugf(`
presence: initial diagnosis of %s
  now (zk)          %s
  last write (zk)   %s
  apparent delay    %s
  timeout           %s
  alive             %t
`[1:], n.path, now, mtime, delay, n.timeout, n.alive)
	} else {
		// If this method is not being run for the first time, we know that
		// the node has just changed, so we know that it's alive and there's
		// no need to check for staleness with the clock node.
		n.alive = true
	}
	return nil
}

// update reads from ZooKeeper the current values of n.alive and n.timeout.
func (n *node) update() error {
	content, stat, err := n.conn.Get(n.path)
	if zk.IsError(err, zk.ZNONODE) {
		n.alive = false
		n.timeout = 0
		return nil
	} else if err != nil {
		return err
	}
	return n.setState(content, stat, true)
}

// updateW reads from ZooKeeper the current values of n.alive and n.timeout,
// and returns a ZooKeeper watch that will be notified of data or existence
// changes. firstTime should be true iff n.updateW has not already been called.
func (n *node) updateW(firstTime bool) (<-chan zk.Event, error) {
	content, stat, watch, err := n.conn.GetW(n.path)
	if zk.IsError(err, zk.ZNONODE) {
		n.alive = false
		n.timeout = 0
		stat, watch, err = n.conn.ExistsW(n.path)
		if err != nil {
			return nil, err
		}
		if stat != nil {
			// Someone *just* created the node; try again.
			return n.updateW(firstTime)
		}
		return watch, nil
	} else if err != nil {
		return nil, err
	}
	return watch, n.setState(content, stat, firstTime)
}

// waitDead sends false to watch when the node is deleted, or when it has
// not been updated recently enough to still qualify as alive. zkWatch must
// be observing changes on the existent node at n.path.
func (n *node) waitDead(zkWatch <-chan zk.Event, watch chan bool) {
	for n.alive {
		select {
		case <-time.After(n.timeout):
			if err := n.update(); err != nil {
				close(watch)
				return
			}
		case event := <-zkWatch:
			if !event.Ok() {
				close(watch)
				return
			}
			switch event.Type {
			case zk.EVENT_DELETED:
				n.alive = false
			case zk.EVENT_CHANGED:
				var err error
				zkWatch, err = n.updateW(false)
				if err != nil {
					close(watch)
					return
				}
			default:
				panic(fmt.Errorf("Unexpected event: %v", event))
			}
		}
	}
	watch <- false
}

// waitAlive sends true to watch when the node is changed or created. zkWatch
// must be observing changes on the stale/nonexistent node at n.path.
func (n *node) waitAlive(zkWatch <-chan zk.Event, watch chan bool) {
	for !n.alive {
		event := <-zkWatch
		if !event.Ok() {
			close(watch)
			return
		}
		switch event.Type {
		case zk.EVENT_CREATED, zk.EVENT_CHANGED:
			n.alive = true
		case zk.EVENT_DELETED:
			// The pinger is still dead (just differently dead); start a new watch.
			var err error
			zkWatch, err = n.updateW(false)
			if err != nil {
				close(watch)
				return
			}
		default:
			panic(fmt.Errorf("Unexpected event: %v", event))
		}
	}
	watch <- true
}

// Alive returns whether the Pinger at the given path seems to be alive.
func Alive(conn *zk.Conn, path string) (bool, error) {
	n := &node{conn: conn, path: path}
	err := n.update()
	return n.alive, err
}

// AliveW returns whether the Pinger at the given path seems to be alive.
// It also returns a channel that will receive the new status when it changes.
// If an error is encountered after AliveW returns, the channel will be closed.
func AliveW(conn *zk.Conn, path string) (bool, <-chan bool, error) {
	n := &node{conn: conn, path: path}
	zkWatch, err := n.updateW(true)
	if err != nil {
		return false, nil, err
	}
	alive := n.alive
	watch := make(chan bool)
	if alive {
		go n.waitDead(zkWatch, watch)
	} else {
		go n.waitAlive(zkWatch, watch)
	}
	return alive, watch, nil
}

// WaitAlive blocks until the node at the given path
// has been recently pinged or a timeout occurs.
func WaitAlive(conn *zk.Conn, path string, timeout time.Duration) error {
	alive, watch, err := AliveW(conn, path)
	if err != nil {
		return err
	}
	if alive {
		return nil
	}
	select {
	case alive, ok := <-watch:
		if !ok {
			return fmt.Errorf("presence: channel closed while waiting")
		}
		if !alive {
			return fmt.Errorf("presence: alive watch misbehaved while waiting")
		}
	case <-time.After(timeout):
		return fmt.Errorf("presence: still not alive after timeout")
	}
	return nil
}
<|endoftext|>


<commit_before>package inigo_test

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"fmt"
	"github.com/cloudfoundry-incubator/inigo/loggredile"
	"github.com/fraenkel/candiedyaml"
	"io"
	"io/ioutil"

	"github.com/cloudfoundry-incubator/inigo/inigo_server"
	"github.com/cloudfoundry-incubator/inigo/stager_runner"
	"github.com/cloudfoundry-incubator/inigo/zipper"
	"github.com/cloudfoundry-incubator/runtime-schema/models/factories"
	"github.com/cloudfoundry/yagnats"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Stager", func() {
	var otherStagerRunner *stager_runner.StagerRunner

	BeforeEach(func() {
		fileServerRunner.Start()
		otherStagerRunner = stager_runner.New(
			stagerPath,
			etcdRunner.NodeURLS(),
			[]string{fmt.Sprintf("127.0.0.1:%d", natsPort)},
		)
	})

	Context("when unable to find an appropriate compiler", func() {
		BeforeEach(func() {
			executorRunner.Start()
			stagerRunner.Start()
		})

		It("returns an error", func() {
			receivedMessages := make(chan *yagnats.Message)
			natsRunner.MessageBus.Subscribe("compiler-stagers-test", func(message *yagnats.Message) {
				receivedMessages <- message
			})

			natsRunner.MessageBus.PublishWithReplyTo(
				"diego.staging.start",
				"compiler-stagers-test",
				[]byte(`{"app_id": "some-app-guid", "task_id": "some-task-id", "stack": "no-compiler"}`))

			var receivedMessage *yagnats.Message
			Eventually(receivedMessages, 2.0).Should(Receive(&receivedMessage))
			Ω(receivedMessage.Payload).Should(ContainSubstring("No compiler defined for requested stack"))
			Consistently(receivedMessages, 2.0).ShouldNot(Receive())
		})
	})

	Describe("Staging", func() {
		var outputGuid string
		var stagingMessage []byte
		var buildpackToUse string

		BeforeEach(func() {
			executorRunner.Start()

			buildpackToUse = "admin_buildpack.zip"
			outputGuid = factories.GenerateGuid()

			fileServerRunner.ServeFile("smelter.zip", smelterZipPath)

			//make and upload an app
			var appFiles = []zipper.ZipFile{
				{"my-app", "scooby-doo"},
			}
			zipper.CreateZipFile("/tmp/app.zip", appFiles)
			inigoserver.UploadFile("app.zip", "/tmp/app.zip")

			//make and upload a buildpack
			var adminBuildpackFiles = []zipper.ZipFile{
				{"bin/detect", `#!/bin/bash
				echo My Buildpack
				`},
				{"bin/compile", `#!/bin/bash
				echo COMPILING BUILDPACK
				echo $SOME_STAGING_ENV
				touch $1/compiled
				`},
				{"bin/release", `#!/bin/bash
cat <<EOF
---
default_process_types:
  web: start-command
EOF
				`},
			}
			zipper.CreateZipFile("/tmp/admin_buildpack.zip", adminBuildpackFiles)
			inigoserver.UploadFile("admin_buildpack.zip", "/tmp/admin_buildpack.zip")

			var bustedAdminBuildpackFiles = []zipper.ZipFile{
				{"bin/detect", `#!/bin/bash]
				exit 1
				`},
				{"bin/compile", `#!/bin/bash`},
				{"bin/release", `#!/bin/bash`},
			}
			zipper.CreateZipFile("/tmp/busted_admin_buildpack.zip", bustedAdminBuildpackFiles)
			inigoserver.UploadFile("busted_admin_buildpack.zip", "/tmp/busted_admin_buildpack.zip")
		})

		JustBeforeEach(func() {
			stagingMessage = []byte(
				fmt.Sprintf(
					`{
						"app_id": "some-app-guid",
						"task_id": "some-task-id",
						"memory_mb": 128,
						"disk_mb": 128,
						"file_descriptors": 1024,
						"stack": "default",
						"app_bits_download_uri": "%s",
						"buildpacks" : [{ "key": "test-buildpack", "url": "%s" }],
						"environment": [["SOME_STAGING_ENV", "%s"]]
					}`,
					inigoserver.DownloadUrl("app.zip"),
					inigoserver.DownloadUrl(buildpackToUse),
					outputGuid,
				),
			)
		})

		Context("with one stager running", func() {
			BeforeEach(func() {
				stagerRunner.Start("--compilers", `{"default":"smelter.zip"}`)
			})

			It("runs the compiler on the executor with the correct environment variables, bits and log tag, and responds with the detected buildpack", func() {
				//listen for NATS response
				payloads := make(chan []byte)

				natsRunner.MessageBus.Subscribe("stager-test", func(msg *yagnats.Message) {
					payloads <- msg.Payload
				})

				//stream logs
				messages, stop := loggredile.StreamMessages(
					loggregatorRunner.Config.OutgoingPort,
					"/tail/?app=some-app-guid",
				)
				defer close(stop)

				logOutput := ""
				go func() {
					for message := range messages {
						Ω(message.GetSourceName()).To(Equal("STG"))
						logOutput += string(message.GetMessage()) + "\n"
					}
				}()

				//publish the staging message
				err := natsRunner.MessageBus.PublishWithReplyTo(
					"diego.staging.start",
					"stager-test",
					stagingMessage,
				)
				Ω(err).ShouldNot(HaveOccurred())

				//wait for staging to complete
				var payload []byte
				Eventually(payloads, 10.0).Should(Receive(&payload))

				//Assert on the staging output (detected buildpack)
				Ω(string(payload)).Should(Equal(`{"detected_buildpack":"My Buildpack"}`))

				//Asser the user saw reasonable output
				Eventually(func() string {
					return logOutput
				}).Should(ContainSubstring("COMPILING BUILDPACK"))
				Ω(logOutput).Should(ContainSubstring(outputGuid))

				//Fetch the compiled droplet from the fakeCC
				dropletData, ok := fakeCC.UploadedDroplets["some-app-guid"]
				Ω(ok).Should(BeTrue())
				Ω(dropletData).ShouldNot(BeEmpty())

				//Unzip the droplet
				ungzippedDropletData, err := gzip.NewReader(bytes.NewReader(dropletData))
				Ω(err).ShouldNot(HaveOccurred())

				//Untar the droplet
				untarredDropletData := tar.NewReader(ungzippedDropletData)
				dropletContents := map[string][]byte{}
				for {
					hdr, err := untarredDropletData.Next()
					if err == io.EOF {
						break
					}
					Ω(err).ShouldNot(HaveOccurred())

					content, err := ioutil.ReadAll(untarredDropletData)
					Ω(err).ShouldNot(HaveOccurred())

					dropletContents[hdr.Name] = content
				}

				//Assert the droplet has the right files in it
				Ω(dropletContents).Should(HaveKey("./"))
				Ω(dropletContents).Should(HaveKey("./staging_info.yml"))
				Ω(dropletContents).Should(HaveKey("./logs/"))
				Ω(dropletContents).Should(HaveKey("./tmp/"))
				Ω(dropletContents).Should(HaveKey("./app/"))
				Ω(dropletContents).Should(HaveKey("./app/my-app"))
				Ω(dropletContents).Should(HaveKey("./app/compiled"))

				//Assert the files contain the right content
				Ω(string(dropletContents["./app/my-app"])).Should(Equal("scooby-doo"))

				//In particular, staging_info.yml should have the correct detected_buildpack and start_command
				yamlDecoder := candiedyaml.NewDecoder(bytes.NewReader(dropletContents["./staging_info.yml"]))
				stagingInfo := map[string]string{}
				err = yamlDecoder.Decode(&stagingInfo)
				Ω(err).ShouldNot(HaveOccurred())

				Ω(stagingInfo["detected_buildpack"]).Should(Equal("My Buildpack"))
				Ω(stagingInfo["start_command"]).Should(Equal("start-command"))

				//Assert nothing else crept into the droplet
				Ω(dropletContents).Should(HaveLen(7))
			})

			Context("when compilation fails", func() {
				BeforeEach(func() {
					buildpackToUse = "busted_admin_buildpack.zip"
				})

				It("responds with the error, and no detected buildpack present", func() {
					payloads := make(chan []byte)

					natsRunner.MessageBus.Subscribe("stager-test", func(msg *yagnats.Message) {
						payloads <- msg.Payload
					})

					messages, stop := loggredile.StreamMessages(
						loggregatorRunner.Config.OutgoingPort,
						"/tail/?app=some-app-guid",
					)
					defer close(stop)

					logOutput := ""
					go func() {
						for message := range messages {
							Ω(message.GetSourceName()).To(Equal("STG"))
							logOutput += string(message.GetMessage()) + "\n"
						}
					}()

					err := natsRunner.MessageBus.PublishWithReplyTo(
						"diego.staging.start",
						"stager-test",
						stagingMessage,
					)
					Ω(err).ShouldNot(HaveOccurred())

					var payload []byte
					Eventually(payloads, 10.0).Should(Receive(&payload))
					Ω(string(payload)).Should(Equal(`{"error":"Process returned with exit value: 1"}`))

					Eventually(func() string {
						return logOutput
					}, 5.0).Should(ContainSubstring("no buildpack detected"))
				})
			})
		})

		Context("with two stagers running", func() {
			BeforeEach(func() {
				stagerRunner.Start("--compilers", `{"default":"smelter.zip"}`)
				otherStagerRunner.Start("--compilers", `{"default":"smelter.zip"}`)
			})

			AfterEach(func() {
				otherStagerRunner.Stop()
			})

			It("only one returns a staging completed response", func() {
				received := make(chan bool)
				natsRunner.MessageBus.Subscribe("two-stagers-test", func(message *yagnats.Message) {
					received <- true
				})

				natsRunner.MessageBus.PublishWithReplyTo(
					"diego.staging.start",
					"two-stagers-test",
					stagingMessage,
				)

				Eventually(received, 10.0).Should(Receive())
				Consistently(received, 2.0).ShouldNot(Receive())
			})
		})
	})
})
<commit_msg>lowercase error messages<commit_after>package inigo_test

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"fmt"
	"github.com/cloudfoundry-incubator/inigo/loggredile"
	"github.com/fraenkel/candiedyaml"
	"io"
	"io/ioutil"

	"github.com/cloudfoundry-incubator/inigo/inigo_server"
	"github.com/cloudfoundry-incubator/inigo/stager_runner"
	"github.com/cloudfoundry-incubator/inigo/zipper"
	"github.com/cloudfoundry-incubator/runtime-schema/models/factories"
	"github.com/cloudfoundry/yagnats"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Stager", func() {
	var otherStagerRunner *stager_runner.StagerRunner

	BeforeEach(func() {
		fileServerRunner.Start()
		otherStagerRunner = stager_runner.New(
			stagerPath,
			etcdRunner.NodeURLS(),
			[]string{fmt.Sprintf("127.0.0.1:%d", natsPort)},
		)
	})

	Context("when unable to find an appropriate compiler", func() {
		BeforeEach(func() {
			executorRunner.Start()
			stagerRunner.Start()
		})

		It("returns an error", func() {
			receivedMessages := make(chan *yagnats.Message)
			natsRunner.MessageBus.Subscribe("compiler-stagers-test", func(message *yagnats.Message) {
				receivedMessages <- message
			})

			natsRunner.MessageBus.PublishWithReplyTo(
				"diego.staging.start",
				"compiler-stagers-test",
				[]byte(`{"app_id": "some-app-guid", "task_id": "some-task-id", "stack": "no-compiler"}`))

			var receivedMessage *yagnats.Message
			Eventually(receivedMessages, 2.0).Should(Receive(&receivedMessage))
			Ω(receivedMessage.Payload).Should(ContainSubstring("no compiler defined for requested stack"))
			Consistently(receivedMessages, 2.0).ShouldNot(Receive())
		})
	})

	Describe("Staging", func() {
		var outputGuid string
		var stagingMessage []byte
		var buildpackToUse string

		BeforeEach(func() {
			executorRunner.Start()

			buildpackToUse = "admin_buildpack.zip"
			outputGuid = factories.GenerateGuid()

			fileServerRunner.ServeFile("smelter.zip", smelterZipPath)

			//make and upload an app
			var appFiles = []zipper.ZipFile{
				{"my-app", "scooby-doo"},
			}
			zipper.CreateZipFile("/tmp/app.zip", appFiles)
			inigoserver.UploadFile("app.zip", "/tmp/app.zip")

			//make and upload a buildpack
			var adminBuildpackFiles = []zipper.ZipFile{
				{"bin/detect", `#!/bin/bash
				echo My Buildpack
				`},
				{"bin/compile", `#!/bin/bash
				echo COMPILING BUILDPACK
				echo $SOME_STAGING_ENV
				touch $1/compiled
				`},
				{"bin/release", `#!/bin/bash
cat <<EOF
---
default_process_types:
  web: start-command
EOF
				`},
			}
			zipper.CreateZipFile("/tmp/admin_buildpack.zip", adminBuildpackFiles)
			inigoserver.UploadFile("admin_buildpack.zip", "/tmp/admin_buildpack.zip")

			var bustedAdminBuildpackFiles = []zipper.ZipFile{
				{"bin/detect", `#!/bin/bash]
				exit 1
				`},
				{"bin/compile", `#!/bin/bash`},
				{"bin/release", `#!/bin/bash`},
			}
			zipper.CreateZipFile("/tmp/busted_admin_buildpack.zip", bustedAdminBuildpackFiles)
			inigoserver.UploadFile("busted_admin_buildpack.zip", "/tmp/busted_admin_buildpack.zip")
		})

		JustBeforeEach(func() {
			stagingMessage = []byte(
				fmt.Sprintf(
					`{
						"app_id": "some-app-guid",
						"task_id": "some-task-id",
						"memory_mb": 128,
						"disk_mb": 128,
						"file_descriptors": 1024,
						"stack": "default",
						"app_bits_download_uri": "%s",
						"buildpacks" : [{ "key": "test-buildpack", "url": "%s" }],
						"environment": [["SOME_STAGING_ENV", "%s"]]
					}`,
					inigoserver.DownloadUrl("app.zip"),
					inigoserver.DownloadUrl(buildpackToUse),
					outputGuid,
				),
			)
		})

		Context("with one stager running", func() {
			BeforeEach(func() {
				stagerRunner.Start("--compilers", `{"default":"smelter.zip"}`)
			})

			It("runs the compiler on the executor with the correct environment variables, bits and log tag, and responds with the detected buildpack", func() {
				//listen for NATS response
				payloads := make(chan []byte)

				natsRunner.MessageBus.Subscribe("stager-test", func(msg *yagnats.Message) {
					payloads <- msg.Payload
				})

				//stream logs
				messages, stop := loggredile.StreamMessages(
					loggregatorRunner.Config.OutgoingPort,
					"/tail/?app=some-app-guid",
				)
				defer close(stop)

				logOutput := ""
				go func() {
					for message := range messages {
						Ω(message.GetSourceName()).To(Equal("STG"))
						logOutput += string(message.GetMessage()) + "\n"
					}
				}()

				//publish the staging message
				err := natsRunner.MessageBus.PublishWithReplyTo(
					"diego.staging.start",
					"stager-test",
					stagingMessage,
				)
				Ω(err).ShouldNot(HaveOccurred())

				//wait for staging to complete
				var payload []byte
				Eventually(payloads, 10.0).Should(Receive(&payload))

				//Assert on the staging output (detected buildpack)
				Ω(string(payload)).Should(Equal(`{"detected_buildpack":"My Buildpack"}`))

				//Asser the user saw reasonable output
				Eventually(func() string {
					return logOutput
				}).Should(ContainSubstring("COMPILING BUILDPACK"))
				Ω(logOutput).Should(ContainSubstring(outputGuid))

				//Fetch the compiled droplet from the fakeCC
				dropletData, ok := fakeCC.UploadedDroplets["some-app-guid"]
				Ω(ok).Should(BeTrue())
				Ω(dropletData).ShouldNot(BeEmpty())

				//Unzip the droplet
				ungzippedDropletData, err := gzip.NewReader(bytes.NewReader(dropletData))
				Ω(err).ShouldNot(HaveOccurred())

				//Untar the droplet
				untarredDropletData := tar.NewReader(ungzippedDropletData)
				dropletContents := map[string][]byte{}
				for {
					hdr, err := untarredDropletData.Next()
					if err == io.EOF {
						break
					}
					Ω(err).ShouldNot(HaveOccurred())

					content, err := ioutil.ReadAll(untarredDropletData)
					Ω(err).ShouldNot(HaveOccurred())

					dropletContents[hdr.Name] = content
				}

				//Assert the droplet has the right files in it
				Ω(dropletContents).Should(HaveKey("./"))
				Ω(dropletContents).Should(HaveKey("./staging_info.yml"))
				Ω(dropletContents).Should(HaveKey("./logs/"))
				Ω(dropletContents).Should(HaveKey("./tmp/"))
				Ω(dropletContents).Should(HaveKey("./app/"))
				Ω(dropletContents).Should(HaveKey("./app/my-app"))
				Ω(dropletContents).Should(HaveKey("./app/compiled"))

				//Assert the files contain the right content
				Ω(string(dropletContents["./app/my-app"])).Should(Equal("scooby-doo"))

				//In particular, staging_info.yml should have the correct detected_buildpack and start_command
				yamlDecoder := candiedyaml.NewDecoder(bytes.NewReader(dropletContents["./staging_info.yml"]))
				stagingInfo := map[string]string{}
				err = yamlDecoder.Decode(&stagingInfo)
				Ω(err).ShouldNot(HaveOccurred())

				Ω(stagingInfo["detected_buildpack"]).Should(Equal("My Buildpack"))
				Ω(stagingInfo["start_command"]).Should(Equal("start-command"))

				//Assert nothing else crept into the droplet
				Ω(dropletContents).Should(HaveLen(7))
			})

			Context("when compilation fails", func() {
				BeforeEach(func() {
					buildpackToUse = "busted_admin_buildpack.zip"
				})

				It("responds with the error, and no detected buildpack present", func() {
					payloads := make(chan []byte)

					natsRunner.MessageBus.Subscribe("stager-test", func(msg *yagnats.Message) {
						payloads <- msg.Payload
					})

					messages, stop := loggredile.StreamMessages(
						loggregatorRunner.Config.OutgoingPort,
						"/tail/?app=some-app-guid",
					)
					defer close(stop)

					logOutput := ""
					go func() {
						for message := range messages {
							Ω(message.GetSourceName()).To(Equal("STG"))
							logOutput += string(message.GetMessage()) + "\n"
						}
					}()

					err := natsRunner.MessageBus.PublishWithReplyTo(
						"diego.staging.start",
						"stager-test",
						stagingMessage,
					)
					Ω(err).ShouldNot(HaveOccurred())

					var payload []byte
					Eventually(payloads, 10.0).Should(Receive(&payload))
					Ω(string(payload)).Should(Equal(`{"error":"process exited with status 1"}`))

					Eventually(func() string {
						return logOutput
					}, 5.0).Should(ContainSubstring("no buildpack detected"))
				})
			})
		})

		Context("with two stagers running", func() {
			BeforeEach(func() {
				stagerRunner.Start("--compilers", `{"default":"smelter.zip"}`)
				otherStagerRunner.Start("--compilers", `{"default":"smelter.zip"}`)
			})

			AfterEach(func() {
				otherStagerRunner.Stop()
			})

			It("only one returns a staging completed response", func() {
				received := make(chan bool)
				natsRunner.MessageBus.Subscribe("two-stagers-test", func(message *yagnats.Message) {
					received <- true
				})

				natsRunner.MessageBus.PublishWithReplyTo(
					"diego.staging.start",
					"two-stagers-test",
					stagingMessage,
				)

				Eventually(received, 10.0).Should(Receive())
				Consistently(received, 2.0).ShouldNot(Receive())
			})
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2016 Apcera Inc. All rights reserved.

package main

import (
	"flag"
	"os"
	"runtime"
	"strings"

	"fmt"
	natsd "github.com/nats-io/gnatsd/server"
	"github.com/nats-io/stan"
	stand "github.com/nats-io/stan-server/server"
	"github.com/nats-io/stan-server/stores"
)

func main() {

	// Parse flags
	sOpts, nOpts := parseFlags()
	stand.EnableDefaultLogger(&nOpts)
	stand.Noticef("Starting stan-server[%s] version %s", sOpts.ID, stan.Version)
	stand.RunServerWithOpts(&sOpts, &nOpts)

	runtime.Goexit()
}

func parseFlags() (stand.Options, natsd.Options) {
	var showVersion bool
	var debugAndTrace bool
	var showTLSHelp bool
	var configFile string

	stanOpts := stand.DefaultOptions
	flag.StringVar(&stanOpts.ID, "cluster_id", stand.DefaultClusterID, "Cluster ID.")
	flag.StringVar(&stanOpts.StoreType, "store", stores.TypeMemory, fmt.Sprintf("Store type: (%s|%s)", stores.TypeMemory, stores.TypeFile))
	flag.StringVar(&stanOpts.FilestoreDir, "dir", "", "Root directory")
	flag.IntVar(&stanOpts.MaxChannels, "max_channels", stand.DefaultChannelLimit, "Max number of channels")
	flag.IntVar(&stanOpts.MaxSubscriptions, "max_subs", stand.DefaultSubStoreLimit, "Max number of subscriptions per channel")
	flag.IntVar(&stanOpts.MaxMsgs, "max_msgs", stand.DefaultMsgStoreLimit, "Max number of messages per channel")
	flag.Uint64Var(&stanOpts.MaxBytes, "max_bytes", stand.DefaultMsgSizeStoreLimit, "Max messages total size per channel")

	natsOpts := natsd.Options{}

	// TODO: Expose gnatsd parsing into server options
	// (cls) This is a development placeholder until gnatsd
	// options parsing is exposed, if we go that way.
	flag.IntVar(&natsOpts.Port, "port", 0, "Port to listen on.")
	flag.IntVar(&natsOpts.Port, "p", 0, "Port to listen on.")
	flag.StringVar(&natsOpts.Host, "addr", "", "Network host to listen on.")
	flag.StringVar(&natsOpts.Host, "a", "", "Network host to listen on.")
	flag.StringVar(&natsOpts.Host, "net", "", "Network host to listen on.")
	flag.BoolVar(&natsOpts.Debug, "D", false, "Enable Debug logging.")
	flag.BoolVar(&natsOpts.Debug, "debug", false, "Enable Debug logging.")
	flag.BoolVar(&natsOpts.Trace, "V", false, "Enable Trace logging.")
	flag.BoolVar(&natsOpts.Trace, "trace", false, "Enable Trace logging.")
	flag.BoolVar(&debugAndTrace, "DV", false, "Enable Debug and Trace logging.")
	flag.BoolVar(&natsOpts.Logtime, "T", true, "Timestamp log entries.")
	flag.BoolVar(&natsOpts.Logtime, "logtime", true, "Timestamp log entries.")
	flag.StringVar(&natsOpts.Username, "user", "", "Username required for connection.")
	flag.StringVar(&natsOpts.Password, "pass", "", "Password required for connection.")
	flag.StringVar(&natsOpts.Authorization, "auth", "", "Authorization token required for connection.")
	flag.IntVar(&natsOpts.HTTPPort, "m", 0, "HTTP Port for /varz, /connz endpoints.")
	flag.IntVar(&natsOpts.HTTPPort, "http_port", 0, "HTTP Port for /varz, /connz endpoints.")
	flag.IntVar(&natsOpts.HTTPSPort, "ms", 0, "HTTPS Port for /varz, /connz endpoints.")
	flag.IntVar(&natsOpts.HTTPSPort, "https_port", 0, "HTTPS Port for /varz, /connz endpoints.")
	flag.StringVar(&configFile, "c", "", "Configuration file.")
	flag.StringVar(&configFile, "config", "", "Configuration file.")
	flag.StringVar(&natsOpts.PidFile, "P", "", "File to store process pid.")
	flag.StringVar(&natsOpts.PidFile, "pid", "", "File to store process pid.")
	flag.StringVar(&natsOpts.LogFile, "l", "", "File to store logging output.")
	flag.StringVar(&natsOpts.LogFile, "log", "", "File to store logging output.")
	flag.BoolVar(&natsOpts.Syslog, "s", false, "Enable syslog as log method.")
	flag.BoolVar(&natsOpts.Syslog, "syslog", false, "Enable syslog as log method..")
	flag.StringVar(&natsOpts.RemoteSyslog, "r", "", "Syslog server addr (udp://localhost:514).")
	flag.StringVar(&natsOpts.RemoteSyslog, "remote_syslog", "", "Syslog server addr (udp://localhost:514).")
	flag.BoolVar(&showVersion, "version", false, "Print version information.")
	flag.BoolVar(&showVersion, "v", false, "Print version information.")
	flag.IntVar(&natsOpts.ProfPort, "profile", 0, "Profiling HTTP port")
	flag.StringVar(&natsOpts.RoutesStr, "routes", "", "Routes to actively solicit a connection.")
	flag.StringVar(&natsOpts.ClusterListenStr, "cluster", "", "Cluster url from which members can solicit routes.")
	flag.StringVar(&natsOpts.ClusterListenStr, "cluster_listen", "", "Cluster url from which members can solicit routes.")
	flag.BoolVar(&showTLSHelp, "help_tls", false, "TLS help.")
	flag.BoolVar(&natsOpts.TLS, "tls", false, "Enable TLS.")
	flag.BoolVar(&natsOpts.TLSVerify, "tlsverify", false, "Enable TLS with client verification.")
	flag.StringVar(&natsOpts.TLSCert, "tlscert", "", "Server certificate file.")
	flag.StringVar(&natsOpts.TLSKey, "tlskey", "", "Private key for server certificate.")
	flag.StringVar(&natsOpts.TLSCaCert, "tlscacert", "", "Client certificate CA for verification.")

	// Not public per se, will be replaced with dynamic system, but can be used to lower memory footprint when
	// lots of connections present.
	flag.IntVar(&natsOpts.BufSize, "bs", 0, "Read/Write buffer size per client connection.")

	flag.Usage = stand.Usage
	flag.Parse()

	// Show version and exit
	if showVersion {
		fmt.Printf("stan-server version %s, ", stan.Version)
		natsd.PrintServerAndExit()
	}

	if showTLSHelp {
		natsd.PrintTLSHelpAndDie()
	}

	// Ensure some options are set based on selected store type
	checkStoreOpts(&stanOpts)

	// One flag can set multiple options.
	if debugAndTrace {
		natsOpts.Trace, natsOpts.Debug = true, true
	}

	return stanOpts, natsOpts
}

func checkStoreOpts(opts *stand.Options) {
	// Convert the user input to upper case
	storeType := strings.ToUpper(opts.StoreType)

	// If FILE, check some parameters
	if storeType == stores.TypeFile {
		if opts.FilestoreDir == "" {
			fmt.Printf("\nFor %v stores, option \"-dir\" must be specified\n", stores.TypeFile)
			flag.Usage()
			os.Exit(0)
		}
	}
}
<commit_msg>Remove buffer size parameter.<commit_after>// Copyright 2016 Apcera Inc. All rights reserved.

package main

import (
	"flag"
	"os"
	"runtime"
	"strings"

	"fmt"
	natsd "github.com/nats-io/gnatsd/server"
	"github.com/nats-io/stan"
	stand "github.com/nats-io/stan-server/server"
	"github.com/nats-io/stan-server/stores"
)

func main() {

	// Parse flags
	sOpts, nOpts := parseFlags()
	stand.EnableDefaultLogger(&nOpts)
	stand.Noticef("Starting stan-server[%s] version %s", sOpts.ID, stan.Version)
	stand.RunServerWithOpts(&sOpts, &nOpts)

	runtime.Goexit()
}

func parseFlags() (stand.Options, natsd.Options) {
	var showVersion bool
	var debugAndTrace bool
	var showTLSHelp bool
	var configFile string

	stanOpts := stand.DefaultOptions
	flag.StringVar(&stanOpts.ID, "cluster_id", stand.DefaultClusterID, "Cluster ID.")
	flag.StringVar(&stanOpts.StoreType, "store", stores.TypeMemory, fmt.Sprintf("Store type: (%s|%s)", stores.TypeMemory, stores.TypeFile))
	flag.StringVar(&stanOpts.FilestoreDir, "dir", "", "Root directory")
	flag.IntVar(&stanOpts.MaxChannels, "max_channels", stand.DefaultChannelLimit, "Max number of channels")
	flag.IntVar(&stanOpts.MaxSubscriptions, "max_subs", stand.DefaultSubStoreLimit, "Max number of subscriptions per channel")
	flag.IntVar(&stanOpts.MaxMsgs, "max_msgs", stand.DefaultMsgStoreLimit, "Max number of messages per channel")
	flag.Uint64Var(&stanOpts.MaxBytes, "max_bytes", stand.DefaultMsgSizeStoreLimit, "Max messages total size per channel")

	natsOpts := natsd.Options{}

	// TODO: Expose gnatsd parsing into server options
	// (cls) This is a development placeholder until gnatsd
	// options parsing is exposed, if we go that way.
	flag.IntVar(&natsOpts.Port, "port", 0, "Port to listen on.")
	flag.IntVar(&natsOpts.Port, "p", 0, "Port to listen on.")
	flag.StringVar(&natsOpts.Host, "addr", "", "Network host to listen on.")
	flag.StringVar(&natsOpts.Host, "a", "", "Network host to listen on.")
	flag.StringVar(&natsOpts.Host, "net", "", "Network host to listen on.")
	flag.BoolVar(&natsOpts.Debug, "D", false, "Enable Debug logging.")
	flag.BoolVar(&natsOpts.Debug, "debug", false, "Enable Debug logging.")
	flag.BoolVar(&natsOpts.Trace, "V", false, "Enable Trace logging.")
	flag.BoolVar(&natsOpts.Trace, "trace", false, "Enable Trace logging.")
	flag.BoolVar(&debugAndTrace, "DV", false, "Enable Debug and Trace logging.")
	flag.BoolVar(&natsOpts.Logtime, "T", true, "Timestamp log entries.")
	flag.BoolVar(&natsOpts.Logtime, "logtime", true, "Timestamp log entries.")
	flag.StringVar(&natsOpts.Username, "user", "", "Username required for connection.")
	flag.StringVar(&natsOpts.Password, "pass", "", "Password required for connection.")
	flag.StringVar(&natsOpts.Authorization, "auth", "", "Authorization token required for connection.")
	flag.IntVar(&natsOpts.HTTPPort, "m", 0, "HTTP Port for /varz, /connz endpoints.")
	flag.IntVar(&natsOpts.HTTPPort, "http_port", 0, "HTTP Port for /varz, /connz endpoints.")
	flag.IntVar(&natsOpts.HTTPSPort, "ms", 0, "HTTPS Port for /varz, /connz endpoints.")
	flag.IntVar(&natsOpts.HTTPSPort, "https_port", 0, "HTTPS Port for /varz, /connz endpoints.")
	flag.StringVar(&configFile, "c", "", "Configuration file.")
	flag.StringVar(&configFile, "config", "", "Configuration file.")
	flag.StringVar(&natsOpts.PidFile, "P", "", "File to store process pid.")
	flag.StringVar(&natsOpts.PidFile, "pid", "", "File to store process pid.")
	flag.StringVar(&natsOpts.LogFile, "l", "", "File to store logging output.")
	flag.StringVar(&natsOpts.LogFile, "log", "", "File to store logging output.")
	flag.BoolVar(&natsOpts.Syslog, "s", false, "Enable syslog as log method.")
	flag.BoolVar(&natsOpts.Syslog, "syslog", false, "Enable syslog as log method..")
	flag.StringVar(&natsOpts.RemoteSyslog, "r", "", "Syslog server addr (udp://localhost:514).")
	flag.StringVar(&natsOpts.RemoteSyslog, "remote_syslog", "", "Syslog server addr (udp://localhost:514).")
	flag.BoolVar(&showVersion, "version", false, "Print version information.")
	flag.BoolVar(&showVersion, "v", false, "Print version information.")
	flag.IntVar(&natsOpts.ProfPort, "profile", 0, "Profiling HTTP port")
	flag.StringVar(&natsOpts.RoutesStr, "routes", "", "Routes to actively solicit a connection.")
	flag.StringVar(&natsOpts.ClusterListenStr, "cluster", "", "Cluster url from which members can solicit routes.")
	flag.StringVar(&natsOpts.ClusterListenStr, "cluster_listen", "", "Cluster url from which members can solicit routes.")
	flag.BoolVar(&showTLSHelp, "help_tls", false, "TLS help.")
	flag.BoolVar(&natsOpts.TLS, "tls", false, "Enable TLS.")
	flag.BoolVar(&natsOpts.TLSVerify, "tlsverify", false, "Enable TLS with client verification.")
	flag.StringVar(&natsOpts.TLSCert, "tlscert", "", "Server certificate file.")
	flag.StringVar(&natsOpts.TLSKey, "tlskey", "", "Private key for server certificate.")
	flag.StringVar(&natsOpts.TLSCaCert, "tlscacert", "", "Client certificate CA for verification.")

	// Not public per se, will be replaced with dynamic system, but can be used to lower memory footprint when

	flag.Usage = stand.Usage
	flag.Parse()

	// Show version and exit
	if showVersion {
		fmt.Printf("stan-server version %s, ", stan.Version)
		natsd.PrintServerAndExit()
	}

	if showTLSHelp {
		natsd.PrintTLSHelpAndDie()
	}

	// Ensure some options are set based on selected store type
	checkStoreOpts(&stanOpts)

	// One flag can set multiple options.
	if debugAndTrace {
		natsOpts.Trace, natsOpts.Debug = true, true
	}

	return stanOpts, natsOpts
}

func checkStoreOpts(opts *stand.Options) {
	// Convert the user input to upper case
	storeType := strings.ToUpper(opts.StoreType)

	// If FILE, check some parameters
	if storeType == stores.TypeFile {
		if opts.FilestoreDir == "" {
			fmt.Printf("\nFor %v stores, option \"-dir\" must be specified\n", stores.TypeFile)
			flag.Usage()
			os.Exit(0)
		}
	}
}
<|endoftext|>


<commit_before>// Package store provides a distributed SQLite instance.
//
// Distributed consensus is provided via the Raft algorithm.
package store

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	"github.com/hashicorp/raft-boltdb"
	sql "github.com/otoolep/rqlite/db"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 10 * time.Second
)

type command struct {
	Tx      bool     `json:"tx,omitempty"`
	Queries []string `json:"queries,omitempty"`
}

// Store is a SQLite database, where all changes are made via Raft consensus.
type Store struct {
	raftDir  string
	raftBind string

	mu sync.Mutex

	raft   *raft.Raft // The consensus mechanism
	dbConf *sql.Config
	dbPath string
	db     *sql.DB // The underlying SQLite store

	logger *log.Logger

	SQLiteDB string
}

// New returns a new Store.
func New(dbConf *sql.Config, dir, bind string) *Store {
	return &Store{
		raftDir:  dir,
		raftBind: bind,
		dbConf:   dbConf,
		dbPath:   filepath.Join(dir, "db.sqlite3"),
		logger:   log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

// Open opens the store. If enableSingle is set, and there are no existing peers,
// then this node becomesthe first node, and therefore leader, of the cluster.
func (s *Store) Open(enableSingle bool) error {
	// Create the database. It must be deleted as it will be rebuilt from
	// (possibly) a snapshot and committed log entries.
	if err := os.Remove(s.dbPath); !os.IsNotExist(err) {
		return err
	}

	db, err := sql.OpenWithConfiguration(s.dbPath, s.dbConf)
	if err != nil {
		return err
	}
	s.db = db

	// Setup Raft configuration.
	config := raft.DefaultConfig()

	// Check for any existing peers.
	peers, err := readPeersJSON(filepath.Join(s.raftDir, "peers.json"))
	if err != nil {
		return err
	}

	// Allow the node to entry single-mode, potentially electing itself, if
	// explicitly enabled and there is only 1 node in the cluster already.
	if enableSingle && len(peers) <= 1 {
		s.logger.Println("enabling single-node mode")
		config.EnableSingleNode = true
		config.DisableBootstrapAfterElect = false
	}

	// Setup Raft communication.
	addr, err := net.ResolveTCPAddr("tcp", s.raftBind)
	if err != nil {
		return err
	}
	transport, err := raft.NewTCPTransport(s.raftBind, addr, 3, 10*time.Second, os.Stderr)
	if err != nil {
		return err
	}

	// Create peer storage.
	peerStore := raft.NewJSONPeers(s.raftDir, transport)

	// Create the snapshot store. This allows the Raft to truncate the log.
	snapshots, err := raft.NewFileSnapshotStore(s.raftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	// Create the log store and stable store.
	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.raftDir, "raft.db"))
	if err != nil {
		return fmt.Errorf("new bolt store: %s", err)
	}

	// Instantiate the Raft systems.
	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, peerStore, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra

	return nil
}

func (s *Store) Close() error {
	return s.db.Close()
}

// Execute executes queries that return no rows.
func (s *Store) Execute(queries []string, tx bool) ([]*sql.Result, error) {
	if s.raft.State() != raft.Leader {
		return nil, fmt.Errorf("not leader")
	}

	c := &command{
		Tx:      tx,
		Queries: queries,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return nil, err
	}

	f := s.raft.Apply(b, raftTimeout)
	if e := f.(raft.Future); e.Error() != nil {
		return nil, e.Error()
	}

	r := f.Response().(*fsmResponse)
	return r.results, r.error
}

// Query execute queries that return rows.
func (s *Store) Query(queries []string, tx bool) ([]*sql.Rows, error) {
	r, err := s.db.Query(queries, tx)
	return r, err
}

// Join joins a node, located at addr, to this store. The node must be ready to
// respond to Raft communications at that address.
func (s *Store) Join(addr string) error {
	s.logger.Printf("received join request for remote node as %s", addr)

	f := s.raft.AddPeer(addr)
	if f.Error() != nil {
		return f.Error()
	}
	s.logger.Printf("node at %s joined successfully", addr)
	return nil
}

type fsm Store

type fsmResponse struct {
	results []*sql.Result
	error   error
}

// Apply applies a Raft log entry to the database.
func (f *fsm) Apply(l *raft.Log) interface{} {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	r, err := f.db.Execute(c.Queries, c.Tx)

	return &fsmResponse{results: r, error: err}
}

// Snapshot returns a snapshot of the database. The caller must ensure that
// no transaction is taking place during this call.
//
// http://sqlite.org/howtocorrupt.html states it is safe to do this
// as long as no transaction is in progress.
func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	b, err := ioutil.ReadFile(f.dbPath)
	if err != nil {
		log.Printf("Failed to generate snapshot: %s", err.Error())
		return nil, err
	}
	return &fsmSnapshot{data: b}, nil
}

// Restore restores the database to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	// Need to write bytes to SQLite file.
	// Open it with desired DSN params.
	// Implies the Open call above, should always blow away existing database, since
	// it will be restored completely from log, or from log and remaining log entries.
	return nil
}

type fsmSnapshot struct {
	data []byte
}

func (f *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	err := func() error {
		// Write data to sink.
		if _, err := sink.Write(f.data); err != nil {
			return err
		}

		// Close the sink.
		if err := sink.Close(); err != nil {
			return err
		}

		return nil
	}()

	if err != nil {
		sink.Cancel()
		return err
	}

	return nil
}

func (f *fsmSnapshot) Release() {}

func readPeersJSON(path string) ([]string, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	if len(b) == 0 {
		return nil, nil
	}

	var peers []string
	dec := json.NewDecoder(bytes.NewReader(b))
	if err := dec.Decode(&peers); err != nil {
		return nil, err
	}

	return peers, nil
}
<commit_msg>Correctly handle database removal<commit_after>// Package store provides a distributed SQLite instance.
//
// Distributed consensus is provided via the Raft algorithm.
package store

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	"github.com/hashicorp/raft-boltdb"
	sql "github.com/otoolep/rqlite/db"
)

const (
	retainSnapshotCount = 2
	raftTimeout         = 10 * time.Second
)

type command struct {
	Tx      bool     `json:"tx,omitempty"`
	Queries []string `json:"queries,omitempty"`
}

// Store is a SQLite database, where all changes are made via Raft consensus.
type Store struct {
	raftDir  string
	raftBind string

	mu sync.Mutex

	raft   *raft.Raft // The consensus mechanism
	dbConf *sql.Config
	dbPath string
	db     *sql.DB // The underlying SQLite store

	logger *log.Logger

	SQLiteDB string
}

// New returns a new Store.
func New(dbConf *sql.Config, dir, bind string) *Store {
	return &Store{
		raftDir:  dir,
		raftBind: bind,
		dbConf:   dbConf,
		dbPath:   filepath.Join(dir, "db.sqlite3"),
		logger:   log.New(os.Stderr, "[store] ", log.LstdFlags),
	}
}

// Open opens the store. If enableSingle is set, and there are no existing peers,
// then this node becomesthe first node, and therefore leader, of the cluster.
func (s *Store) Open(enableSingle bool) error {
	// Create the database. It must be deleted as it will be rebuilt from
	// (possibly) a snapshot and committed log entries.
	if err := os.Remove(s.dbPath); err != nil && !os.IsNotExist(err) {
		return err
	}

	db, err := sql.OpenWithConfiguration(s.dbPath, s.dbConf)
	if err != nil {
		return err
	}
	s.db = db

	// Setup Raft configuration.
	config := raft.DefaultConfig()

	// Check for any existing peers.
	peers, err := readPeersJSON(filepath.Join(s.raftDir, "peers.json"))
	if err != nil {
		return err
	}

	// Allow the node to entry single-mode, potentially electing itself, if
	// explicitly enabled and there is only 1 node in the cluster already.
	if enableSingle && len(peers) <= 1 {
		s.logger.Println("enabling single-node mode")
		config.EnableSingleNode = true
		config.DisableBootstrapAfterElect = false
	}

	// Setup Raft communication.
	addr, err := net.ResolveTCPAddr("tcp", s.raftBind)
	if err != nil {
		return err
	}
	transport, err := raft.NewTCPTransport(s.raftBind, addr, 3, 10*time.Second, os.Stderr)
	if err != nil {
		return err
	}

	// Create peer storage.
	peerStore := raft.NewJSONPeers(s.raftDir, transport)

	// Create the snapshot store. This allows the Raft to truncate the log.
	snapshots, err := raft.NewFileSnapshotStore(s.raftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	// Create the log store and stable store.
	logStore, err := raftboltdb.NewBoltStore(filepath.Join(s.raftDir, "raft.db"))
	if err != nil {
		return fmt.Errorf("new bolt store: %s", err)
	}

	// Instantiate the Raft systems.
	ra, err := raft.NewRaft(config, (*fsm)(s), logStore, logStore, snapshots, peerStore, transport)
	if err != nil {
		return fmt.Errorf("new raft: %s", err)
	}
	s.raft = ra

	return nil
}

func (s *Store) Close() error {
	return s.db.Close()
}

// Execute executes queries that return no rows.
func (s *Store) Execute(queries []string, tx bool) ([]*sql.Result, error) {
	if s.raft.State() != raft.Leader {
		return nil, fmt.Errorf("not leader")
	}

	c := &command{
		Tx:      tx,
		Queries: queries,
	}
	b, err := json.Marshal(c)
	if err != nil {
		return nil, err
	}

	f := s.raft.Apply(b, raftTimeout)
	if e := f.(raft.Future); e.Error() != nil {
		return nil, e.Error()
	}

	r := f.Response().(*fsmResponse)
	return r.results, r.error
}

// Query execute queries that return rows.
func (s *Store) Query(queries []string, tx bool) ([]*sql.Rows, error) {
	r, err := s.db.Query(queries, tx)
	return r, err
}

// Join joins a node, located at addr, to this store. The node must be ready to
// respond to Raft communications at that address.
func (s *Store) Join(addr string) error {
	s.logger.Printf("received join request for remote node as %s", addr)

	f := s.raft.AddPeer(addr)
	if f.Error() != nil {
		return f.Error()
	}
	s.logger.Printf("node at %s joined successfully", addr)
	return nil
}

type fsm Store

type fsmResponse struct {
	results []*sql.Result
	error   error
}

// Apply applies a Raft log entry to the database.
func (f *fsm) Apply(l *raft.Log) interface{} {
	var c command
	if err := json.Unmarshal(l.Data, &c); err != nil {
		panic(fmt.Sprintf("failed to unmarshal command: %s", err.Error()))
	}

	r, err := f.db.Execute(c.Queries, c.Tx)

	return &fsmResponse{results: r, error: err}
}

// Snapshot returns a snapshot of the database. The caller must ensure that
// no transaction is taking place during this call.
//
// http://sqlite.org/howtocorrupt.html states it is safe to do this
// as long as no transaction is in progress.
func (f *fsm) Snapshot() (raft.FSMSnapshot, error) {
	b, err := ioutil.ReadFile(f.dbPath)
	if err != nil {
		log.Printf("Failed to generate snapshot: %s", err.Error())
		return nil, err
	}
	return &fsmSnapshot{data: b}, nil
}

// Restore restores the database to a previous state.
func (f *fsm) Restore(rc io.ReadCloser) error {
	// Need to write bytes to SQLite file.
	// Open it with desired DSN params.
	// Implies the Open call above, should always blow away existing database, since
	// it will be restored completely from log, or from log and remaining log entries.
	return nil
}

type fsmSnapshot struct {
	data []byte
}

func (f *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	err := func() error {
		// Write data to sink.
		if _, err := sink.Write(f.data); err != nil {
			return err
		}

		// Close the sink.
		if err := sink.Close(); err != nil {
			return err
		}

		return nil
	}()

	if err != nil {
		sink.Cancel()
		return err
	}

	return nil
}

func (f *fsmSnapshot) Release() {}

func readPeersJSON(path string) ([]string, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil && !os.IsNotExist(err) {
		return nil, err
	}

	if len(b) == 0 {
		return nil, nil
	}

	var peers []string
	dec := json.NewDecoder(bytes.NewReader(b))
	if err := dec.Decode(&peers); err != nil {
		return nil, err
	}

	return peers, nil
}
<|endoftext|>