<commit_before>// Copyright 2015 The rkt Authors
//
// 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 main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/coreos/rkt/tests/testutils"
)

var volTests = []struct {
	rktCmd string
	expect string
}{
	// Check that we can read files in the ACI
	{
		`/bin/sh -c "export FILE=/dir1/file ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true ^READ_FILE^"`,
		`<<<dir1>>>`,
	},
	// Check that we can read files from a volume (both ro and rw)
	{
		`/bin/sh -c "export FILE=/dir1/file ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RW_READ_FILE^"`,
		`<<<host>>>`,
	},
	{
		`/bin/sh -c "export FILE=/dir1/file ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RO_READ_FILE^"`,
		`<<<host>>>`,
	},
	// Check that we can write to files in the ACI
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=1 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true ^WRITE_FILE^"`,
		`<<<1>>>`,
	},
	// Check that we can write files to a volume (both ro and rw)
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=2 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RW_WRITE_FILE^"`,
		`<<<2>>>`,
	},
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=3 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RO_WRITE_FILE^"`,
		`Cannot write to file "/dir1/file": open /dir1/file: read-only file system`,
	},
	// Check that the volume still contains the file previously written
	{
		`/bin/sh -c "export FILE=/dir1/file ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RO_READ_FILE^"`,
		`<<<2>>>`,
	},
	// Check that injecting a rw mount/volume works without any mountpoint in the image manifest
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=1 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ ^VOL_ADD_MOUNT_RW^ --mount=volume=dir1,target=dir1"`,
		`<<<1>>>`,
	},
	// Check that injecting a ro mount/volume works without any mountpoint in the image manifest
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=1 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^,readOnly=true ^VOL_ADD_MOUNT_RO^ --mount=volume=dir1,target=dir1"`,
		`Cannot write to file "/dir1/file": open /dir1/file: read-only file system`,
	},
}

func TestVolumes(t *testing.T) {
	readFileImage := patchTestACI("rkt-inspect-read-file.aci", "--exec=/inspect --read-file")
	defer os.Remove(readFileImage)
	writeFileImage := patchTestACI("rkt-inspect-write-file.aci", "--exec=/inspect --write-file --read-file")
	defer os.Remove(writeFileImage)
	volRwReadFileImage := patchTestACI("rkt-inspect-vol-rw-read-file.aci", "--exec=/inspect --read-file", "--mounts=dir1,path=/dir1,readOnly=false")
	defer os.Remove(volRwReadFileImage)
	volRwWriteFileImage := patchTestACI("rkt-inspect-vol-rw-write-file.aci", "--exec=/inspect --write-file --read-file", "--mounts=dir1,path=/dir1,readOnly=false")
	defer os.Remove(volRwWriteFileImage)
	volRoReadFileImage := patchTestACI("rkt-inspect-vol-ro-read-file.aci", "--exec=/inspect --read-file", "--mounts=dir1,path=/dir1,readOnly=true")
	defer os.Remove(volRoReadFileImage)
	volRoWriteFileImage := patchTestACI("rkt-inspect-vol-ro-write-file.aci", "--exec=/inspect --write-file --read-file", "--mounts=dir1,path=/dir1,readOnly=true")
	defer os.Remove(volRoWriteFileImage)
	volAddMountRwImage := patchTestACI("rkt-inspect-vol-add-mount-rw.aci", "--exec=/inspect --write-file --read-file")
	defer os.Remove(volAddMountRwImage)
	volAddMountRoImage := patchTestACI("rkt-inspect-vol-add-mount-ro.aci", "--exec=/inspect --write-file --read-file")
	defer os.Remove(volAddMountRoImage)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	tmpdir := createTempDirOrPanic("rkt-tests.")
	defer os.RemoveAll(tmpdir)

	tmpfile := filepath.Join(tmpdir, "file")
	if err := ioutil.WriteFile(tmpfile, []byte("host"), 0600); err != nil {
		t.Fatalf("Cannot create temporary file: %v", err)
	}

	for i, tt := range volTests {
		cmd := strings.Replace(tt.rktCmd, "^TMPDIR^", tmpdir, -1)
		cmd = strings.Replace(cmd, "^RKT_BIN^", ctx.Cmd(), -1)
		cmd = strings.Replace(cmd, "^READ_FILE^", readFileImage, -1)
		cmd = strings.Replace(cmd, "^WRITE_FILE^", writeFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_RO_READ_FILE^", volRoReadFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_RO_WRITE_FILE^", volRoWriteFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_RW_READ_FILE^", volRwReadFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_RW_WRITE_FILE^", volRwWriteFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_ADD_MOUNT_RW^", volAddMountRwImage, -1)
		cmd = strings.Replace(cmd, "^VOL_ADD_MOUNT_RO^", volAddMountRoImage, -1)

		t.Logf("Running test #%v", i)
		child := spawnOrFail(t, cmd)
		defer waitOrFail(t, child, true)

		if err := expectTimeoutWithOutput(child, tt.expect, time.Minute); err != nil {
			fmt.Printf("Command: %s\n", cmd)
			t.Fatalf("Expected %q but not found #%v: %v", tt.expect, i, err)
		}
	}
}
<commit_msg>functional tests: add test for implicit empty volumes<commit_after>// Copyright 2015 The rkt Authors
//
// 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 main

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/coreos/rkt/tests/testutils"
)

var volTests = []struct {
	rktCmd string
	expect string
}{
	// Check that we can read files in the ACI
	{
		`/bin/sh -c "export FILE=/dir1/file ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true ^READ_FILE^"`,
		`<<<dir1>>>`,
	},
	// Check that we can read files from a volume (both ro and rw)
	{
		`/bin/sh -c "export FILE=/dir1/file ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RW_READ_FILE^"`,
		`<<<host>>>`,
	},
	{
		`/bin/sh -c "export FILE=/dir1/file ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RO_READ_FILE^"`,
		`<<<host>>>`,
	},
	// Check that we can write to files in the ACI
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=1 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true ^WRITE_FILE^"`,
		`<<<1>>>`,
	},
	// Check that we can write files to a volume (both ro and rw)
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=2 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RW_WRITE_FILE^"`,
		`<<<2>>>`,
	},
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=3 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RO_WRITE_FILE^"`,
		`Cannot write to file "/dir1/file": open /dir1/file: read-only file system`,
	},
	// Check that the volume still contains the file previously written
	{
		`/bin/sh -c "export FILE=/dir1/file ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ --mount=volume=dir1,target=dir1 ^VOL_RO_READ_FILE^"`,
		`<<<2>>>`,
	},
	// Check that injecting a rw mount/volume works without any mountpoint in the image manifest
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=1 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^ ^VOL_ADD_MOUNT_RW^ --mount=volume=dir1,target=dir1"`,
		`<<<1>>>`,
	},
	// Check that injecting a ro mount/volume works without any mountpoint in the image manifest
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=1 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true --volume=dir1,kind=host,source=^TMPDIR^,readOnly=true ^VOL_ADD_MOUNT_RO^ --mount=volume=dir1,target=dir1"`,
		`Cannot write to file "/dir1/file": open /dir1/file: read-only file system`,
	},
	// Check that an implicit empty volume is created if the user didn't provide it but there's a mount point in the app
	{
		`/bin/sh -c "export FILE=/dir1/file CONTENT=1 ; ^RKT_BIN^ --debug --insecure-skip-verify run --mds-register=false --inherit-env=true ^WRITE_FILE^"`,
		`<<<1>>>`,
	},
}

func TestVolumes(t *testing.T) {
	readFileImage := patchTestACI("rkt-inspect-read-file.aci", "--exec=/inspect --read-file")
	defer os.Remove(readFileImage)
	writeFileImage := patchTestACI("rkt-inspect-write-file.aci", "--exec=/inspect --write-file --read-file")
	defer os.Remove(writeFileImage)
	volRwReadFileImage := patchTestACI("rkt-inspect-vol-rw-read-file.aci", "--exec=/inspect --read-file", "--mounts=dir1,path=/dir1,readOnly=false")
	defer os.Remove(volRwReadFileImage)
	volRwWriteFileImage := patchTestACI("rkt-inspect-vol-rw-write-file.aci", "--exec=/inspect --write-file --read-file", "--mounts=dir1,path=/dir1,readOnly=false")
	defer os.Remove(volRwWriteFileImage)
	volRoReadFileImage := patchTestACI("rkt-inspect-vol-ro-read-file.aci", "--exec=/inspect --read-file", "--mounts=dir1,path=/dir1,readOnly=true")
	defer os.Remove(volRoReadFileImage)
	volRoWriteFileImage := patchTestACI("rkt-inspect-vol-ro-write-file.aci", "--exec=/inspect --write-file --read-file", "--mounts=dir1,path=/dir1,readOnly=true")
	defer os.Remove(volRoWriteFileImage)
	volAddMountRwImage := patchTestACI("rkt-inspect-vol-add-mount-rw.aci", "--exec=/inspect --write-file --read-file")
	defer os.Remove(volAddMountRwImage)
	volAddMountRoImage := patchTestACI("rkt-inspect-vol-add-mount-ro.aci", "--exec=/inspect --write-file --read-file")
	defer os.Remove(volAddMountRoImage)
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	tmpdir := createTempDirOrPanic("rkt-tests.")
	defer os.RemoveAll(tmpdir)

	tmpfile := filepath.Join(tmpdir, "file")
	if err := ioutil.WriteFile(tmpfile, []byte("host"), 0600); err != nil {
		t.Fatalf("Cannot create temporary file: %v", err)
	}

	for i, tt := range volTests {
		cmd := strings.Replace(tt.rktCmd, "^TMPDIR^", tmpdir, -1)
		cmd = strings.Replace(cmd, "^RKT_BIN^", ctx.Cmd(), -1)
		cmd = strings.Replace(cmd, "^READ_FILE^", readFileImage, -1)
		cmd = strings.Replace(cmd, "^WRITE_FILE^", writeFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_RO_READ_FILE^", volRoReadFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_RO_WRITE_FILE^", volRoWriteFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_RW_READ_FILE^", volRwReadFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_RW_WRITE_FILE^", volRwWriteFileImage, -1)
		cmd = strings.Replace(cmd, "^VOL_ADD_MOUNT_RW^", volAddMountRwImage, -1)
		cmd = strings.Replace(cmd, "^VOL_ADD_MOUNT_RO^", volAddMountRoImage, -1)

		t.Logf("Running test #%v", i)
		child := spawnOrFail(t, cmd)
		defer waitOrFail(t, child, true)

		if err := expectTimeoutWithOutput(child, tt.expect, time.Minute); err != nil {
			fmt.Printf("Command: %s\n", cmd)
			t.Fatalf("Expected %q but not found #%v: %v", tt.expect, i, err)
		}
	}
}
<|endoftext|>


<commit_before>package state

import (
	"bytes"
	"strconv"
	"strings"

	"github.com/mesos/mesos-go/upid"
)

// Resources holds resources as defined in the /state.json Mesos HTTP endpoint.
type Resources struct {
	PortRanges string `json:"ports"`
}

// Ports returns a slice of individual ports expanded from PortRanges.
func (r Resources) Ports() []string {
	if r.PortRanges == "" || r.PortRanges == "[]" {
		return []string{}
	}

	rhs := strings.Split(r.PortRanges, "[")[1]
	lhs := strings.Split(rhs, "]")[0]

	yports := []string{}

	mports := strings.Split(lhs, ",")
	for _, port := range mports {
		tmp := strings.TrimSpace(port)
		pz := strings.Split(tmp, "-")
		lo, _ := strconv.Atoi(pz[0])
		hi, _ := strconv.Atoi(pz[1])

		for t := lo; t <= hi; t++ {
			yports = append(yports, strconv.Itoa(t))
		}
	}
	return yports
}

// Label holds a label as defined in the /state.json Mesos HTTP endpoint.
type Label struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

// Status holds a task status as defined in the /state.json Mesos HTTP endpoint.
type Status struct {
	Timestamp       float64         `json:"timestamp"`
	State           string          `json:"state"`
	Labels          []Label         `json:"labels,omitempty"`
	ContainerStatus ContainerStatus `json:"container_status"`
}

// ContainerStatus holds container metadata as defined in the /state.json
// Mesos HTTP endpoint.
type ContainerStatus struct {
	NetworkInfos []NetworkInfo `json:"network_infos"`
}

// NetworkInfo holds a network address resolution as defined in the
// /state.json Mesos HTTP endpoint.
type NetworkInfo struct {
	IPAddress string `json:"ip_address"`
}

// Task holds a task as defined in the /state.json Mesos HTTP endpoint.
type Task struct {
	FrameworkID   string   `json:"framework_id"`
	ID            string   `json:"id"`
	Name          string   `json:"name"`
	SlaveID       string   `json:"slave_id"`
	State         string   `json:"state"`
	Statuses      []Status `json:"statuses"`
	Resources     `json:"resources"`
	DiscoveryInfo DiscoveryInfo `json:"discovery"`

	SlaveIP string `json:"-"`
}

var ipLabels = map[string]string{
	"docker": "Docker.NetworkSettings.IPAddress",
	"mesos":  "MesosContainerizer.NetworkSettings.IPAddress",
}

// IP extracts the IP from a task given the prioritized list of IP sources
func (t *Task) IP(srcs ...string) string {
	for _, src := range srcs {
		switch src {
		case "host":
			return t.SlaveIP
		case "docker", "mesos":
			if ip := t.containerIP(src); ip != "" {
				return ip
			}
		}
	}
	return ""
}

// HasDiscoveryInfo return whether the DiscoveryInfo was provided in the state.json
func (t *Task) HasDiscoveryInfo() bool {
	return t.DiscoveryInfo.Name != ""
}

// containerIP extracts a container ip from a Mesos state.json task. If no
// container ip is provided, an empty string is returned.
func (t *Task) containerIP(src string) string {
	ipLabel := ipLabels[src]

	// find TASK_RUNNING statuses
	var latestContainerIP string
	var latestTimestamp float64
	for _, status := range t.Statuses {
		if status.State != "TASK_RUNNING" || status.Timestamp <= latestTimestamp {
			continue
		}

		// first try to extract the address from the NetworkInfo structure
		if len(status.ContainerStatus.NetworkInfos) > 0 {
			// TODO(CD): handle multiple NetworkInfo objects
			// TODO(CD): only create A records if the address is IPv4
			// TODO(CD): create AAAA records if the address is IPv6
			latestContainerIP = status.ContainerStatus.NetworkInfos[0].IPAddress
			latestTimestamp = status.Timestamp
			break
		}

		// next, fall back to the docker-inspect label
		// TODO(CD): deprecate and then remove this address discovery method
		for _, label := range status.Labels {
			if label.Key == ipLabel {
				latestContainerIP = label.Value
				latestTimestamp = status.Timestamp
				break
			}
		}
	}

	return latestContainerIP
}

// Framework holds a framework as defined in the /state.json Mesos HTTP endpoint.
type Framework struct {
	Tasks    []Task `json:"tasks"`
	PID      PID    `json:"pid"`
	Name     string `json:"name"`
	Hostname string `json:"hostname"`
}

// HostPort returns the hostname and port where a framework's scheduler is
// listening on.
func (f Framework) HostPort() (string, string) {
	if f.PID.UPID != nil {
		return f.PID.Host, f.PID.Port
	}
	return f.Hostname, ""
}

// Slave holds a slave as defined in the /state.json Mesos HTTP endpoint.
type Slave struct {
	ID       string `json:"id"`
	Hostname string `json:"hostname"`
	PID      PID    `json:"pid"`
}

// PID holds a Mesos PID and implements the json.Unmarshaler interface.
type PID struct{ *upid.UPID }

// UnmarshalJSON implements the json.Unmarshaler interface for PIDs.
func (p *PID) UnmarshalJSON(data []byte) (err error) {
	p.UPID, err = upid.Parse(string(bytes.Trim(data, `" `)))
	return err
}

// State holds the state defined in the /state.json Mesos HTTP endpoint.
type State struct {
	Frameworks []Framework `json:"frameworks"`
	Slaves     []Slave     `json:"slaves"`
	Leader     string      `json:"leader"`
}

// DiscoveryInfo holds the discovery meta data for a task defined in the /state.json Mesos HTTP endpoint.
type DiscoveryInfo struct {
	Visibilty   string `json:"visibility"`
	Version     string `json:"version,omitempty"`
	Name        string `json:"name,omitempty"`
	Location    string `json:"location,omitempty"`
	Environment string `json:"environment,omitempty"`
	Labels      struct {
		Labels `json:"labels"`
	} `json:"labels"`
	Ports struct {
		DiscoveryPorts `json:"ports"`
	} `json:"ports"`
}

// Labels holds the key/value labels of a task defined in the /state.json Mesos HTTP endpoint.
type Labels []struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

// DiscoveryPorts holds the ports for a task defined in the /state.json Mesos HTTP endpoint.
type DiscoveryPorts []struct {
	Protocol string `json:"protocol"`
	Number   int    `json:"number"`
	Name     string `json:"name"`
}
<commit_msg>Added omitempty annotation as appropriate for state parsers.<commit_after>package state

import (
	"bytes"
	"strconv"
	"strings"

	"github.com/mesos/mesos-go/upid"
)

// Resources holds resources as defined in the /state.json Mesos HTTP endpoint.
type Resources struct {
	PortRanges string `json:"ports"`
}

// Ports returns a slice of individual ports expanded from PortRanges.
func (r Resources) Ports() []string {
	if r.PortRanges == "" || r.PortRanges == "[]" {
		return []string{}
	}

	rhs := strings.Split(r.PortRanges, "[")[1]
	lhs := strings.Split(rhs, "]")[0]

	yports := []string{}

	mports := strings.Split(lhs, ",")
	for _, port := range mports {
		tmp := strings.TrimSpace(port)
		pz := strings.Split(tmp, "-")
		lo, _ := strconv.Atoi(pz[0])
		hi, _ := strconv.Atoi(pz[1])

		for t := lo; t <= hi; t++ {
			yports = append(yports, strconv.Itoa(t))
		}
	}
	return yports
}

// Label holds a label as defined in the /state.json Mesos HTTP endpoint.
type Label struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

// Status holds a task status as defined in the /state.json Mesos HTTP endpoint.
type Status struct {
	Timestamp       float64         `json:"timestamp"`
	State           string          `json:"state"`
	Labels          []Label         `json:"labels,omitempty"`
	ContainerStatus ContainerStatus `json:"container_status,omitempty"`
}

// ContainerStatus holds container metadata as defined in the /state.json
// Mesos HTTP endpoint.
type ContainerStatus struct {
	NetworkInfos []NetworkInfo `json:"network_infos,omitempty"`
}

// NetworkInfo holds a network address resolution as defined in the
// /state.json Mesos HTTP endpoint.
type NetworkInfo struct {
	IPAddress string `json:"ip_address,omitempty"`
}

// Task holds a task as defined in the /state.json Mesos HTTP endpoint.
type Task struct {
	FrameworkID   string   `json:"framework_id"`
	ID            string   `json:"id"`
	Name          string   `json:"name"`
	SlaveID       string   `json:"slave_id"`
	State         string   `json:"state"`
	Statuses      []Status `json:"statuses"`
	Resources     `json:"resources"`
	DiscoveryInfo DiscoveryInfo `json:"discovery"`

	SlaveIP string `json:"-"`
}

var ipLabels = map[string]string{
	"docker": "Docker.NetworkSettings.IPAddress",
	"mesos":  "MesosContainerizer.NetworkSettings.IPAddress",
}

// IP extracts the IP from a task given the prioritized list of IP sources
func (t *Task) IP(srcs ...string) string {
	for _, src := range srcs {
		switch src {
		case "host":
			return t.SlaveIP
		case "docker", "mesos":
			if ip := t.containerIP(src); ip != "" {
				return ip
			}
		}
	}
	return ""
}

// HasDiscoveryInfo return whether the DiscoveryInfo was provided in the state.json
func (t *Task) HasDiscoveryInfo() bool {
	return t.DiscoveryInfo.Name != ""
}

// containerIP extracts a container ip from a Mesos state.json task. If no
// container ip is provided, an empty string is returned.
func (t *Task) containerIP(src string) string {
	ipLabel := ipLabels[src]

	// find TASK_RUNNING statuses
	var latestContainerIP string
	var latestTimestamp float64
	for _, status := range t.Statuses {
		if status.State != "TASK_RUNNING" || status.Timestamp <= latestTimestamp {
			continue
		}

		// first try to extract the address from the NetworkInfo structure
		if len(status.ContainerStatus.NetworkInfos) > 0 {
			// TODO(CD): handle multiple NetworkInfo objects
			// TODO(CD): only create A records if the address is IPv4
			// TODO(CD): create AAAA records if the address is IPv6
			latestContainerIP = status.ContainerStatus.NetworkInfos[0].IPAddress
			latestTimestamp = status.Timestamp
			break
		}

		// next, fall back to the docker-inspect label
		// TODO(CD): deprecate and then remove this address discovery method
		for _, label := range status.Labels {
			if label.Key == ipLabel {
				latestContainerIP = label.Value
				latestTimestamp = status.Timestamp
				break
			}
		}
	}

	return latestContainerIP
}

// Framework holds a framework as defined in the /state.json Mesos HTTP endpoint.
type Framework struct {
	Tasks    []Task `json:"tasks"`
	PID      PID    `json:"pid"`
	Name     string `json:"name"`
	Hostname string `json:"hostname"`
}

// HostPort returns the hostname and port where a framework's scheduler is
// listening on.
func (f Framework) HostPort() (string, string) {
	if f.PID.UPID != nil {
		return f.PID.Host, f.PID.Port
	}
	return f.Hostname, ""
}

// Slave holds a slave as defined in the /state.json Mesos HTTP endpoint.
type Slave struct {
	ID       string `json:"id"`
	Hostname string `json:"hostname"`
	PID      PID    `json:"pid"`
}

// PID holds a Mesos PID and implements the json.Unmarshaler interface.
type PID struct{ *upid.UPID }

// UnmarshalJSON implements the json.Unmarshaler interface for PIDs.
func (p *PID) UnmarshalJSON(data []byte) (err error) {
	p.UPID, err = upid.Parse(string(bytes.Trim(data, `" `)))
	return err
}

// State holds the state defined in the /state.json Mesos HTTP endpoint.
type State struct {
	Frameworks []Framework `json:"frameworks"`
	Slaves     []Slave     `json:"slaves"`
	Leader     string      `json:"leader"`
}

// DiscoveryInfo holds the discovery meta data for a task defined in the /state.json Mesos HTTP endpoint.
type DiscoveryInfo struct {
	Visibilty   string `json:"visibility"`
	Version     string `json:"version,omitempty"`
	Name        string `json:"name,omitempty"`
	Location    string `json:"location,omitempty"`
	Environment string `json:"environment,omitempty"`
	Labels      struct {
		Labels `json:"labels"`
	} `json:"labels"`
	Ports struct {
		DiscoveryPorts `json:"ports"`
	} `json:"ports"`
}

// Labels holds the key/value labels of a task defined in the /state.json Mesos HTTP endpoint.
type Labels []struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

// DiscoveryPorts holds the ports for a task defined in the /state.json Mesos HTTP endpoint.
type DiscoveryPorts []struct {
	Protocol string `json:"protocol"`
	Number   int    `json:"number"`
	Name     string `json:"name"`
}
<|endoftext|>


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

// The git-cleanup command deletes branches which have already
// been merged to the Gerrit server.
package main

import (
	"bufio"
	"bytes"
	"log"
	"os/exec"
	"regexp"
	"strings"
)

func main() {
	curBranch := ""
	out, err := exec.Command("git", "branch").Output()
	if err != nil {
		log.Fatalf("Error running git branch: %v", err)
	}
	var branches []string
	for _, line := range strings.Split(string(out), "\n") {
		line = strings.TrimSpace(line)
		cur := false
		if strings.HasPrefix(line, "* ") {
			line = line[2:]
			cur = true
		}
		if line == "" {
			continue
		}
		branches = append(branches, line)
		if cur {
			curBranch = line
		}
	}
	if curBranch != "master" {
		log.Fatalf("On branch %s; must be on master", curBranch)
	}

	for _, br := range branches {
		if br == "master" {
			continue
		}
		if isSubmitted(branchChangeID(br)) {
			// Display a ref for the branch we're about to delete,
			// so that if we screw up (never!), the user can get it back easily.
			short, err := exec.Command("git", "rev-parse", "--short", "refs/heads/"+br).Output()
			if err != nil {
				log.Fatalf("Error running git rev-parse: %v", err)
			}
			short = bytes.TrimSpace(short)
			log.Printf("Removing branch %s (%s) ...", br, short)
			if out, err := exec.Command("git", "branch", "-D", br).CombinedOutput(); err != nil {
				log.Printf("Error removing branch %s: %v, %s", br, err, out)
			}
		}
	}
}

var changeRx = regexp.MustCompile(`(?m)^\s*Change-Id: (I[0-9a-f]+)`)

// returns change-id or the empty string
func branchChangeID(br string) string {
	out, _ := exec.Command("git", "show", br).CombinedOutput()
	if m := changeRx.FindSubmatch(out); m != nil {
		return string(m[1])
	}
	return ""
}

func isSubmitted(changeID string) bool {
	if changeID == "" {
		return false
	}
	for _, v := range changeIDLog() {
		if v == changeID {
			return true
		}
	}
	return false
}

var changeIDLogCache []string

func changeIDLog() (ret []string) {
	if v := changeIDLogCache; v != nil {
		return v
	}
	defer func() { changeIDLogCache = ret }()
	cmd := exec.Command("git", "log")
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatalf("pipe error: %v", err)
	}
	if err := cmd.Start(); err != nil {
		log.Fatalf("error running git log: %v", err)
	}
	bs := bufio.NewScanner(stdout)
	wantLine := []byte("Change-Id: ")
	for bs.Scan() {
		if !bytes.Contains(bs.Bytes(), wantLine) {
			continue
		}
		changeID := strings.TrimPrefix(strings.TrimSpace(bs.Text()), "Change-Id: ")
		ret = append(ret, changeID)
	}
	if err := bs.Err(); err != nil {
		log.Fatalf("error running git log: %v", err)
	}
	if err := cmd.Wait(); err != nil {
		log.Fatalf("error running git log: %v", err)
	}
	return
}
<commit_msg>Delete .mailed tags<commit_after>// Copyright 2014 Google. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The git-cleanup command deletes branches which have already
// been merged to the Gerrit server.
package main

import (
	"bufio"
	"bytes"
	"log"
	"os/exec"
	"regexp"
	"strings"
)

func main() {
	curBranch := ""
	out, err := exec.Command("git", "branch").Output()
	if err != nil {
		log.Fatalf("Error running git branch: %v", err)
	}
	var branches []string
	for _, line := range strings.Split(string(out), "\n") {
		line = strings.TrimSpace(line)
		cur := false
		if strings.HasPrefix(line, "* ") {
			line = line[2:]
			cur = true
		}
		if line == "" {
			continue
		}
		branches = append(branches, line)
		if cur {
			curBranch = line
		}
	}
	if curBranch != "master" {
		log.Fatalf("On branch %s; must be on master", curBranch)
	}

	for _, br := range branches {
		if br == "master" {
			continue
		}
		if isSubmitted(branchChangeID(br)) {
			// Display a ref for the branch we're about to delete,
			// so that if we screw up (never!), the user can get it back easily.
			short, err := exec.Command("git", "rev-parse", "--short", "refs/heads/"+br).Output()
			if err != nil {
				log.Fatalf("Error running git rev-parse: %v", err)
			}
			short = bytes.TrimSpace(short)
			log.Printf("Removing branch %s (%s) ...", br, short)
			if out, err := exec.Command("git", "branch", "-D", br).CombinedOutput(); err != nil {
				log.Printf("Error removing branch %s: %v, %s", br, err, out)
			}

			// Remove corresponding tag, if there is one.
			tag := br + ".mailed"
			short, err = exec.Command("git", "rev-parse", "--short", "refs/tags/"+tag).Output()
			if err != nil {
				continue
			}
			short = bytes.TrimSpace(short)
			log.Printf("Removing tag %s (%s) ...", tag, short)
			if out, err := exec.Command("git", "tag", "-d", tag).CombinedOutput(); err != nil {
				log.Printf("Error removing tag %s: %v, %s", tag, err, out)
			}
		}
	}
}

var changeRx = regexp.MustCompile(`(?m)^\s*Change-Id: (I[0-9a-f]+)`)

// returns change-id or the empty string
func branchChangeID(br string) string {
	out, _ := exec.Command("git", "show", br).CombinedOutput()
	if m := changeRx.FindSubmatch(out); m != nil {
		return string(m[1])
	}
	return ""
}

func isSubmitted(changeID string) bool {
	if changeID == "" {
		return false
	}
	for _, v := range changeIDLog() {
		if v == changeID {
			return true
		}
	}
	return false
}

var changeIDLogCache []string

func changeIDLog() (ret []string) {
	if v := changeIDLogCache; v != nil {
		return v
	}
	defer func() { changeIDLogCache = ret }()
	cmd := exec.Command("git", "log")
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatalf("pipe error: %v", err)
	}
	if err := cmd.Start(); err != nil {
		log.Fatalf("error running git log: %v", err)
	}
	bs := bufio.NewScanner(stdout)
	wantLine := []byte("Change-Id: ")
	for bs.Scan() {
		if !bytes.Contains(bs.Bytes(), wantLine) {
			continue
		}
		changeID := strings.TrimPrefix(strings.TrimSpace(bs.Text()), "Change-Id: ")
		ret = append(ret, changeID)
	}
	if err := bs.Err(); err != nil {
		log.Fatalf("error running git log: %v", err)
	}
	if err := cmd.Wait(); err != nil {
		log.Fatalf("error running git log: %v", err)
	}
	return
}
<|endoftext|>


<commit_before>package git

import (
	"errors"
	"log"
	"net/http"
	"strings"

	api "github.com/google/go-github/github"
	"golang.org/x/oauth2"
)

var ErrorNotFound = errors.New("not found")

type GithubRepositories struct {
	client *api.Client
}

func NewGithubRepositories(token string) *GithubRepositories {
	tokenSource := oauth2.StaticTokenSource(&oauth2.Token{
		AccessToken: token,
	})
	oauthClient := oauth2.NewClient(oauth2.NoContext, tokenSource)

	return &GithubRepositories{
		client: api.NewClient(oauthClient),
	}
}

func (service *GithubRepositories) All() ([]*Repository, error) {
	opts := &api.RepositoryListOptions{
		ListOptions: api.ListOptions{
			PerPage: 50,
		},
	}

	var allRepos []*Repository

	paginatedRepos := make([]*Repository, opts.ListOptions.PerPage)
	for {
		githubRepos, response, err := service.client.Repositories.List("", opts)
		if err != nil {
			return nil, err
		}

		for i, githubRepo := range githubRepos {
			if githubRepo.FullName == nil {
				log.Printf("[WARN] excluding GitHub repository without full_name %v", githubRepo)
				continue
			}

			if githubRepo.GitURL == nil {
				log.Printf("[WARN] excluding GitHub repository without git_url %v", githubRepo)
				continue
			}

			paginatedRepos[i] = &Repository{
				FullName: *githubRepo.FullName,
				Master:   "master",
			}

			if githubRepo.Description != nil {
				paginatedRepos[i].Description = *githubRepo.Description
			}

			if githubRepo.DefaultBranch != nil {
				paginatedRepos[i].Master = *githubRepo.DefaultBranch
			}

			if githubRepo.HTMLURL != nil {
				paginatedRepos[i].HTMLURL = *githubRepo.HTMLURL
			}
		}

		allRepos = append(allRepos, paginatedRepos[0:len(githubRepos)]...)
		if response.NextPage == 0 {
			break
		}

		opts.ListOptions.Page = response.NextPage
	}

	return allRepos, nil
}

func (service *GithubRepositories) Get(fullName string) (*Repository, error) {
	repoOwner, repoName := ParseRepositoryName(fullName)

	githubRepo, response, err := service.client.Repositories.Get(repoOwner, repoName)
	if err == nil {
		err = api.CheckResponse(response.Response)
	}

	if err != nil {
		if response.StatusCode == http.StatusNotFound {
			return nil, ErrorNotFound
		}

		return nil, err
	}

	masterBranch, _, err := service.client.Repositories.GetBranch(repoOwner, repoName, *githubRepo.DefaultBranch)
	if err != nil {
		return nil, err
	}

	lastCommit, _, err := service.client.Git.GetCommit(repoOwner, repoName, *masterBranch.Commit.SHA)
	if err != nil {
		return nil, err
	}

	repo := repositoryFromGithub(githubRepo)
	repo.LatestMasterCommit = commitFromGithub(lastCommit)

	return repo, nil
}

func (service *GithubRepositories) Track(fullName, callbackURL string) error {
	owner, name := ParseRepositoryName(fullName)
	return service.registerPushWebhook(owner, name, callbackURL)
}

func (service *GithubRepositories) registerPushWebhook(owner, repo, cbURL string) error {
	hook := &api.Hook{
		Name:   new(string),
		Active: new(bool),
		Events: []string{"push"},
		Config: map[string]interface{}{
			"url":          cbURL,
			"content_type": "json",
		},
	}
	*hook.Name = "web"
	*hook.Active = true

	_, response, err := service.client.Repositories.CreateHook(owner, repo, hook)
	if err == nil {
		err = api.CheckResponse(response.Response)
	}

	return err
}

// ParseRepositoryName returns owner and project name for given GitHub repository.
func ParseRepositoryName(fullName string) (string, string) {
	fields := strings.SplitN(fullName, "/", 2)
	return fields[0], fields[1]
}

func repositoryFromGithub(githubRepo *api.Repository) *Repository {
	return &Repository{
		FullName:    *githubRepo.FullName,
		Description: *githubRepo.Description,
		Master:      *githubRepo.DefaultBranch,
		HTMLURL:     *githubRepo.HTMLURL,
		GitURL:      *githubRepo.GitURL,
	}
}

func commitFromGithub(githubCommit *api.Commit) *Commit {
	return &Commit{
		SHA:     *githubCommit.SHA,
		Message: *githubCommit.Message,
		Author:  *githubCommit.Committer.Name,
		Date:    *githubCommit.Committer.Date,
	}
}
<commit_msg>Remove redundant calls to github.CheckResponse<commit_after>package git

import (
	"errors"
	"log"
	"net/http"
	"strings"

	api "github.com/google/go-github/github"
	"golang.org/x/oauth2"
)

var ErrorNotFound = errors.New("not found")

type GithubRepositories struct {
	client *api.Client
}

func NewGithubRepositories(token string) *GithubRepositories {
	tokenSource := oauth2.StaticTokenSource(&oauth2.Token{
		AccessToken: token,
	})
	oauthClient := oauth2.NewClient(oauth2.NoContext, tokenSource)

	return &GithubRepositories{
		client: api.NewClient(oauthClient),
	}
}

func (service *GithubRepositories) All() ([]*Repository, error) {
	opts := &api.RepositoryListOptions{
		ListOptions: api.ListOptions{
			PerPage: 50,
		},
	}

	var allRepos []*Repository

	paginatedRepos := make([]*Repository, opts.ListOptions.PerPage)
	for {
		githubRepos, response, err := service.client.Repositories.List("", opts)
		if err != nil {
			return nil, err
		}

		for i, githubRepo := range githubRepos {
			if githubRepo.FullName == nil {
				log.Printf("[WARN] excluding GitHub repository without full_name %v", githubRepo)
				continue
			}

			if githubRepo.GitURL == nil {
				log.Printf("[WARN] excluding GitHub repository without git_url %v", githubRepo)
				continue
			}

			paginatedRepos[i] = &Repository{
				FullName: *githubRepo.FullName,
				Master:   "master",
			}

			if githubRepo.Description != nil {
				paginatedRepos[i].Description = *githubRepo.Description
			}

			if githubRepo.DefaultBranch != nil {
				paginatedRepos[i].Master = *githubRepo.DefaultBranch
			}

			if githubRepo.HTMLURL != nil {
				paginatedRepos[i].HTMLURL = *githubRepo.HTMLURL
			}
		}

		allRepos = append(allRepos, paginatedRepos[0:len(githubRepos)]...)
		if response.NextPage == 0 {
			break
		}

		opts.ListOptions.Page = response.NextPage
	}

	return allRepos, nil
}

func (service *GithubRepositories) Get(fullName string) (*Repository, error) {
	repoOwner, repoName := ParseRepositoryName(fullName)

	githubRepo, response, err := service.client.Repositories.Get(repoOwner, repoName)
	if err != nil {
		if response.StatusCode == http.StatusNotFound {
			return nil, ErrorNotFound
		}

		return nil, err
	}

	masterBranch, _, err := service.client.Repositories.GetBranch(repoOwner, repoName, *githubRepo.DefaultBranch)
	if err != nil {
		return nil, err
	}

	lastCommit, _, err := service.client.Git.GetCommit(repoOwner, repoName, *masterBranch.Commit.SHA)
	if err != nil {
		return nil, err
	}

	repo := repositoryFromGithub(githubRepo)
	repo.LatestMasterCommit = commitFromGithub(lastCommit)

	return repo, nil
}

func (service *GithubRepositories) Track(fullName, callbackURL string) error {
	owner, name := ParseRepositoryName(fullName)
	return service.registerPushWebhook(owner, name, callbackURL)
}

func (service *GithubRepositories) registerPushWebhook(owner, repo, cbURL string) error {
	hook := &api.Hook{
		Name:   new(string),
		Active: new(bool),
		Events: []string{"push"},
		Config: map[string]interface{}{
			"url":          cbURL,
			"content_type": "json",
		},
	}
	*hook.Name = "web"
	*hook.Active = true

	_, _, err := service.client.Repositories.CreateHook(owner, repo, hook)

	return err
}

// ParseRepositoryName returns owner and project name for given GitHub repository.
func ParseRepositoryName(fullName string) (string, string) {
	fields := strings.SplitN(fullName, "/", 2)
	return fields[0], fields[1]
}

func repositoryFromGithub(githubRepo *api.Repository) *Repository {
	return &Repository{
		FullName:    *githubRepo.FullName,
		Description: *githubRepo.Description,
		Master:      *githubRepo.DefaultBranch,
		HTMLURL:     *githubRepo.HTMLURL,
		GitURL:      *githubRepo.GitURL,
	}
}

func commitFromGithub(githubCommit *api.Commit) *Commit {
	return &Commit{
		SHA:     *githubCommit.SHA,
		Message: *githubCommit.Message,
		Author:  *githubCommit.Committer.Name,
		Date:    *githubCommit.Committer.Date,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/flynn/go-crypto-ssh"
	"github.com/flynn/go-shlex"
)

const PrereceiveHook = `#!/bin/bash
set -eo pipefail; while read oldrev newrev refname; do
[[ $refname = "refs/heads/master" ]] && git archive $newrev | {{RECEIVER}} "$RECEIVE_USER" "$RECEIVE_REPO" "$RECEIVE_KEYNAME" "$RECEIVE_FINGERPRINT" | sed -$([[ $(uname) == "Darwin" ]] && echo l || echo u) "s/^/"$'\e[1G'"/"
done
`

var port *string = flag.String("p", "22", "port to listen on")
var repoPath *string = flag.String("r", "/tmp/repos", "path to repo cache")
var keyPath *string = flag.String("k", "/tmp/keys", "path to named keys")
var noAuth *bool = flag.Bool("n", false, "no client authentication")

var receiver string
var privateKey string
var keyNames = make(map[string]string)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage:  %v [options] <privatekey> <receiver>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if len(os.Args) < 2 {
		flag.Usage()
		return
	}
	privateKey = flag.Arg(0)
	receiver = flag.Arg(1)

	var config *ssh.ServerConfig
	if *noAuth {
		config = &ssh.ServerConfig{NoClientAuth: true}
	} else {
		config = &ssh.ServerConfig{PublicKeyCallback: keyCallback}
	}

	pemBytes, err := ioutil.ReadFile(privateKey)
	if err != nil {
		log.Fatal("Failed to load private key:", err)
	}
	if err = config.SetRSAPrivateKey(pemBytes); err != nil {
		log.Fatal("Failed to parse private key:", err)
	}

	listener, err := ssh.Listen("tcp", "0.0.0.0:"+*port, config)
	if err != nil {
		log.Fatal("failed to listen for connection")
	}
	for {
		// SSH connections just house multiplexed connections
		conn, err := listener.Accept()
		if err != nil {
			log.Println("failed to accept incoming connection")
			continue
		}
		if err := conn.Handshake(); err != nil {
			log.Println("failed to handshake")
			continue
		}
		go handleConnection(conn)
	}
}

func keyCallback(conn *ssh.ServerConn, user, algo string, pubkey []byte) bool {
	clientkey, _, ok := ssh.ParsePublicKey(pubkey)
	if !ok {
		return false
	}
	os.MkdirAll(*keyPath, 0755)
	files, err := ioutil.ReadDir(*keyPath)
	if err != nil {
		panic(err)
	}
	for _, file := range files {
		if !file.IsDir() {
			data, err := ioutil.ReadFile(*keyPath + "/" + file.Name())
			if err != nil {
				panic(err)
			}
			filekey, _, _, _, ok := ssh.ParseAuthorizedKey(data)
			if !ok {
				continue
			}
			if bytes.Equal(clientkey.Marshal(), filekey.Marshal()) {
				keyNames[publicKeyFingerprint(clientkey)] = file.Name()
				return true
			}
		}
	}
	return false
}

func handleConnection(conn *ssh.ServerConn) {
	defer conn.Close()
	for {
		// Accept reads from the connection, demultiplexes packets
		// to their corresponding channels and returns when a new
		// channel request is seen. Some goroutine must always be
		// calling Accept; otherwise no messages will be forwarded
		// to the channels.
		ch, err := conn.Accept()
		if err == io.EOF {
			return
		}
		if err != nil {
			log.Println("hanndleConnection Accept:", err)
			break
		}
		if ch.ChannelType() != "session" {
			ch.Reject(ssh.UnknownChannelType, "unknown channel type")
			break
		}
		go handleChannel(conn, ch)
	}
}

func handleChannel(conn *ssh.ServerConn, ch ssh.Channel) {
	err := ch.Accept()
	if err != nil {
		panic(err)
	}
	defer ch.Close()
	for {
		req, err := ch.ReadRequest()
		if err == io.EOF {
			return
		}
		if err != nil {
			log.Println("handleChannel read request err:", err)
			continue
		}
		switch req.Request {
		case "exec":
			if req.WantReply {
				ch.AckRequest(true)
			}
			cmdline := string(req.Payload[4:])
			cmdargs, err := shlex.Split(cmdline)
			if err != nil {
				panic(err)
			}
			if strings.HasPrefix(cmdargs[1], "/") {
				cmdargs[1] = cmdargs[1][1:]
			}
			ensureCacheRepo(cmdargs[1])
			var keyname, fingerprint string
			if *noAuth {
				fingerprint = ""
				keyname = ""
			} else {
				fingerprint = publicKeyFingerprint(conn.PublicKey)
				keyname = keyNames[fingerprint]
			}
			cmd := exec.Command(cmdargs[0], cmdargs[1:]...)
			cmd.Dir = *repoPath
			cmd.Env = []string{
				"RECEIVE_USER=" + conn.User,
				"RECEIVE_REPO=" + cmdargs[1],
				"RECEIVE_KEYNAME=" + keyname,
				"RECEIVE_FINGERPRINT=" + fingerprint,
			}
			errCh := attachCmd(cmd, ch, ch.Stderr(), ch)
			err = cmd.Start()
			if err != nil {
				panic(err)
			}
			exitCh := exitStatusCh(cmd)
			if err = <-errCh; err != nil {
				panic(err)
			}
			ch.Exit(<-exitCh)
		case "env":
			if req.WantReply {
				ch.AckRequest(true)
			}
		}
	}
}

func attachCmd(cmd *exec.Cmd, stdout, stderr io.Writer, stdin io.Reader) chan error {
	errCh := make(chan error)

	stdinIn, err := cmd.StdinPipe()
	if err != nil {
		panic(err)
	}
	stdoutOut, err := cmd.StdoutPipe()
	if err != nil {
		panic(err)
	}
	stderrOut, err := cmd.StderrPipe()
	if err != nil {
		panic(err)
	}

	go func() {
		_, e := io.Copy(stdinIn, stdin)
		errCh <- e
	}()
	go func() {
		_, e := io.Copy(stdout, stdoutOut)
		errCh <- e
	}()
	go func() {
		_, e := io.Copy(stderr, stderrOut)
		errCh <- e
	}()

	return errCh
}

func exitStatusCh(cmd *exec.Cmd) chan uint {
	exitCh := make(chan uint)
	go func() {
		err := cmd.Wait()
		if err != nil {
			if exiterr, ok := err.(*exec.ExitError); ok {
				// There is no plattform independent way to retrieve
				// the exit code, but the following will work on Unix
				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
					exitCh <- uint(status.ExitStatus())
				}
			} else {
				panic(err)
			}
			return
		}
		exitCh <- uint(0)
	}()
	return exitCh
}

func ensureCacheRepo(path string) {
	cachePath := *repoPath + "/" + path
	if _, err := os.Stat(cachePath); os.IsNotExist(err) {
		os.MkdirAll(cachePath, 0755)
		cmd := exec.Command("git", "init", "--bare")
		cmd.Dir = cachePath
		err = cmd.Run()
		if err != nil {
			panic(err)
		}
	}
	receiver, err := filepath.Abs(receiver)
	if err != nil {
		panic(err)
	}
	ioutil.WriteFile(
		cachePath+"/hooks/pre-receive",
		[]byte(strings.Replace(PrereceiveHook, "{{RECEIVER}}", receiver, 1)),
		0755)
}

func publicKeyFingerprint(key ssh.PublicKey) string {
	var values []string
	h := md5.New()
	h.Write(ssh.MarshalPublicKey(key))
	for _, value := range h.Sum(nil) {
		values = append(values, fmt.Sprintf("%x", value))
	}
	return strings.Join(values, ":")
}
<commit_msg>gitreceived: Fix logging typos<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/flynn/go-crypto-ssh"
	"github.com/flynn/go-shlex"
)

const PrereceiveHook = `#!/bin/bash
set -eo pipefail; while read oldrev newrev refname; do
[[ $refname = "refs/heads/master" ]] && git archive $newrev | {{RECEIVER}} "$RECEIVE_USER" "$RECEIVE_REPO" "$RECEIVE_KEYNAME" "$RECEIVE_FINGERPRINT" | sed -$([[ $(uname) == "Darwin" ]] && echo l || echo u) "s/^/"$'\e[1G'"/"
done
`

var port *string = flag.String("p", "22", "port to listen on")
var repoPath *string = flag.String("r", "/tmp/repos", "path to repo cache")
var keyPath *string = flag.String("k", "/tmp/keys", "path to named keys")
var noAuth *bool = flag.Bool("n", false, "no client authentication")

var receiver string
var privateKey string
var keyNames = make(map[string]string)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage:  %v [options] <privatekey> <receiver>\n\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if len(os.Args) < 2 {
		flag.Usage()
		return
	}
	privateKey = flag.Arg(0)
	receiver = flag.Arg(1)

	var config *ssh.ServerConfig
	if *noAuth {
		config = &ssh.ServerConfig{NoClientAuth: true}
	} else {
		config = &ssh.ServerConfig{PublicKeyCallback: keyCallback}
	}

	pemBytes, err := ioutil.ReadFile(privateKey)
	if err != nil {
		log.Fatal("Failed to load private key:", err)
	}
	if err = config.SetRSAPrivateKey(pemBytes); err != nil {
		log.Fatal("Failed to parse private key:", err)
	}

	listener, err := ssh.Listen("tcp", "0.0.0.0:"+*port, config)
	if err != nil {
		log.Fatal("failed to listen for connection")
	}
	for {
		// SSH connections just house multiplexed connections
		conn, err := listener.Accept()
		if err != nil {
			log.Println("failed to accept incoming connection:", err)
			continue
		}
		if err := conn.Handshake(); err != nil {
			log.Println("failed to handshake:", err)
			continue
		}
		go handleConnection(conn)
	}
}

func keyCallback(conn *ssh.ServerConn, user, algo string, pubkey []byte) bool {
	clientkey, _, ok := ssh.ParsePublicKey(pubkey)
	if !ok {
		return false
	}
	os.MkdirAll(*keyPath, 0755)
	files, err := ioutil.ReadDir(*keyPath)
	if err != nil {
		panic(err)
	}
	for _, file := range files {
		if !file.IsDir() {
			data, err := ioutil.ReadFile(*keyPath + "/" + file.Name())
			if err != nil {
				panic(err)
			}
			filekey, _, _, _, ok := ssh.ParseAuthorizedKey(data)
			if !ok {
				continue
			}
			if bytes.Equal(clientkey.Marshal(), filekey.Marshal()) {
				keyNames[publicKeyFingerprint(clientkey)] = file.Name()
				return true
			}
		}
	}
	return false
}

func handleConnection(conn *ssh.ServerConn) {
	defer conn.Close()
	for {
		// Accept reads from the connection, demultiplexes packets
		// to their corresponding channels and returns when a new
		// channel request is seen. Some goroutine must always be
		// calling Accept; otherwise no messages will be forwarded
		// to the channels.
		ch, err := conn.Accept()
		if err == io.EOF {
			return
		}
		if err != nil {
			log.Println("handleConnection Accept:", err)
			break
		}
		if ch.ChannelType() != "session" {
			ch.Reject(ssh.UnknownChannelType, "unknown channel type")
			break
		}
		go handleChannel(conn, ch)
	}
}

func handleChannel(conn *ssh.ServerConn, ch ssh.Channel) {
	err := ch.Accept()
	if err != nil {
		panic(err)
	}
	defer ch.Close()
	for {
		req, err := ch.ReadRequest()
		if err == io.EOF {
			return
		}
		if err != nil {
			log.Println("handleChannel read request err:", err)
			continue
		}
		switch req.Request {
		case "exec":
			if req.WantReply {
				ch.AckRequest(true)
			}
			cmdline := string(req.Payload[4:])
			cmdargs, err := shlex.Split(cmdline)
			if err != nil {
				panic(err)
			}
			if strings.HasPrefix(cmdargs[1], "/") {
				cmdargs[1] = cmdargs[1][1:]
			}
			ensureCacheRepo(cmdargs[1])
			var keyname, fingerprint string
			if *noAuth {
				fingerprint = ""
				keyname = ""
			} else {
				fingerprint = publicKeyFingerprint(conn.PublicKey)
				keyname = keyNames[fingerprint]
			}
			cmd := exec.Command(cmdargs[0], cmdargs[1:]...)
			cmd.Dir = *repoPath
			cmd.Env = []string{
				"RECEIVE_USER=" + conn.User,
				"RECEIVE_REPO=" + cmdargs[1],
				"RECEIVE_KEYNAME=" + keyname,
				"RECEIVE_FINGERPRINT=" + fingerprint,
			}
			errCh := attachCmd(cmd, ch, ch.Stderr(), ch)
			err = cmd.Start()
			if err != nil {
				panic(err)
			}
			exitCh := exitStatusCh(cmd)
			if err = <-errCh; err != nil {
				panic(err)
			}
			ch.Exit(<-exitCh)
		case "env":
			if req.WantReply {
				ch.AckRequest(true)
			}
		}
	}
}

func attachCmd(cmd *exec.Cmd, stdout, stderr io.Writer, stdin io.Reader) chan error {
	errCh := make(chan error)

	stdinIn, err := cmd.StdinPipe()
	if err != nil {
		panic(err)
	}
	stdoutOut, err := cmd.StdoutPipe()
	if err != nil {
		panic(err)
	}
	stderrOut, err := cmd.StderrPipe()
	if err != nil {
		panic(err)
	}

	go func() {
		_, e := io.Copy(stdinIn, stdin)
		errCh <- e
	}()
	go func() {
		_, e := io.Copy(stdout, stdoutOut)
		errCh <- e
	}()
	go func() {
		_, e := io.Copy(stderr, stderrOut)
		errCh <- e
	}()

	return errCh
}

func exitStatusCh(cmd *exec.Cmd) chan uint {
	exitCh := make(chan uint)
	go func() {
		err := cmd.Wait()
		if err != nil {
			if exiterr, ok := err.(*exec.ExitError); ok {
				// There is no plattform independent way to retrieve
				// the exit code, but the following will work on Unix
				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
					exitCh <- uint(status.ExitStatus())
				}
			} else {
				panic(err)
			}
			return
		}
		exitCh <- uint(0)
	}()
	return exitCh
}

func ensureCacheRepo(path string) {
	cachePath := *repoPath + "/" + path
	if _, err := os.Stat(cachePath); os.IsNotExist(err) {
		os.MkdirAll(cachePath, 0755)
		cmd := exec.Command("git", "init", "--bare")
		cmd.Dir = cachePath
		err = cmd.Run()
		if err != nil {
			panic(err)
		}
	}
	receiver, err := filepath.Abs(receiver)
	if err != nil {
		panic(err)
	}
	ioutil.WriteFile(
		cachePath+"/hooks/pre-receive",
		[]byte(strings.Replace(PrereceiveHook, "{{RECEIVER}}", receiver, 1)),
		0755)
}

func publicKeyFingerprint(key ssh.PublicKey) string {
	var values []string
	h := md5.New()
	h.Write(ssh.MarshalPublicKey(key))
	for _, value := range h.Sum(nil) {
		values = append(values, fmt.Sprintf("%x", value))
	}
	return strings.Join(values, ":")
}
<|endoftext|>


<commit_before><commit_msg>Initial template for TV API, seems a lot simpler than the normal QD API.<commit_after><|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strings"
)

const bindatafile = "bindata.go"

func isDebug(args []string) bool {
	flagset := flag.NewFlagSet("", flag.ContinueOnError)
	debug := flagset.Bool("debug", false, "")
	debugArgs := make([]string, 0)
	for _, arg := range args {
		if strings.HasPrefix(arg, "-debug") {
			debugArgs = append(debugArgs, arg)
		}
	}
	flagset.Parse(debugArgs)
	if debug == nil {
		return false
	}
	return *debug
}

func main() {
	if _, err := exec.LookPath("go-bindata"); err != nil {
		fmt.Println("Cannot find go-bindata executable in path")
		fmt.Println("Maybe you need: go get github.com/elazarl/go-bindata-assetfs/...")
		os.Exit(1)
	}
	cmd := exec.Command("go-bindata", os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		os.Exit(1)
	}
	in, err := os.Open(bindatafile)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Cannot read", bindatafile, err)
		return
	}
	out, err := os.Create("bindata_assetfs.go")
	if err != nil {
		fmt.Fprintln(os.Stderr, "Cannot write 'bindata_assetfs.go'", err)
		return
	}
	debug := isDebug(os.Args[1:])
	r := bufio.NewReader(in)
	done := false
	for line, isPrefix, err := r.ReadLine(); err == nil; line, isPrefix, err = r.ReadLine() {
		if !isPrefix {
			line = append(line, '\n')
		}
		if _, err := out.Write(line); err != nil {
			fmt.Fprintln(os.Stderr, "Cannot write to 'bindata_assetfs.go'", err)
			return
		}
		if !done && !isPrefix && bytes.HasPrefix(line, []byte("import (")) {
			if debug {
				fmt.Fprintln(out, "\t\"net/http\"")
			} else {
				fmt.Fprintln(out, "\t\"github.com/elazarl/go-bindata-assetfs\"")
			}
			done = true
		}
	}
	if debug {
		fmt.Fprintln(out, `
func assetFS() http.FileSystem {
	for k := range _bintree.Children {
		return http.Dir(k)
	}
	panic("unreachable")
}`)
	} else {
		fmt.Fprintln(out, `
func assetFS() *assetfs.AssetFS {
	for k := range _bintree.Children {
		return &assetfs.AssetFS{Asset: Asset, AssetDir: AssetDir, AssetInfo: AssetInfo, Prefix: k}
	}
	panic("unreachable")
}`)
	}
	// Close files BEFORE remove calls (don't use defer).
	in.Close()
	out.Close()
	if err := os.Remove(bindatafile); err != nil {
		fmt.Fprintln(os.Stderr, "Cannot remove", bindatafile, err)
	}
}
<commit_msg>This fixes #39 for me: AssetInfo undefined<commit_after>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strings"
)

const bindatafile = "bindata.go"

func isDebug(args []string) bool {
	flagset := flag.NewFlagSet("", flag.ContinueOnError)
	debug := flagset.Bool("debug", false, "")
	debugArgs := make([]string, 0)
	for _, arg := range args {
		if strings.HasPrefix(arg, "-debug") {
			debugArgs = append(debugArgs, arg)
		}
	}
	flagset.Parse(debugArgs)
	if debug == nil {
		return false
	}
	return *debug
}

func main() {
	if _, err := exec.LookPath("go-bindata"); err != nil {
		fmt.Println("Cannot find go-bindata executable in path")
		fmt.Println("Maybe you need: go get github.com/elazarl/go-bindata-assetfs/...")
		os.Exit(1)
	}
	cmd := exec.Command("go-bindata", os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		os.Exit(1)
	}
	in, err := os.Open(bindatafile)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Cannot read", bindatafile, err)
		return
	}
	out, err := os.Create("bindata_assetfs.go")
	if err != nil {
		fmt.Fprintln(os.Stderr, "Cannot write 'bindata_assetfs.go'", err)
		return
	}
	debug := isDebug(os.Args[1:])
	r := bufio.NewReader(in)
	done := false
	for line, isPrefix, err := r.ReadLine(); err == nil; line, isPrefix, err = r.ReadLine() {
		if !isPrefix {
			line = append(line, '\n')
		}
		if _, err := out.Write(line); err != nil {
			fmt.Fprintln(os.Stderr, "Cannot write to 'bindata_assetfs.go'", err)
			return
		}
		if !done && !isPrefix && bytes.HasPrefix(line, []byte("import (")) {
			if debug {
				fmt.Fprintln(out, "\t\"net/http\"")
			} else {
				fmt.Fprintln(out, "\t\"github.com/elazarl/go-bindata-assetfs\"")
				fmt.Fprintln(out, "\t\"os\"")
			}
			done = true
		}
	}
	if debug {
		fmt.Fprintln(out, `
func assetFS() http.FileSystem {
	for k := range _bintree.Children {
		return http.Dir(k)
	}
	panic("unreachable")
}`)
	} else {
		fmt.Fprintln(out, `
func assetFS() *assetfs.AssetFS {
	assetInfo := func(path string) (os.FileInfo, error) {
		return os.Stat(path)
	}
	for k := range _bintree.Children {
		return &assetfs.AssetFS{Asset: Asset, AssetDir: AssetDir, AssetInfo: assetInfo, Prefix: k}
	}
	panic("unreachable")
}`)
	}
	// Close files BEFORE remove calls (don't use defer).
	in.Close()
	out.Close()
	if err := os.Remove(bindatafile); err != nil {
		fmt.Fprintln(os.Stderr, "Cannot remove", bindatafile, err)
	}
}
<|endoftext|>


<commit_before>package confdis

import (
	"testing"
	"time"
)

type SampleConfig struct {
	Name  string   `json:"name"`
	Users []string `json:"users"`
	Meta  struct {
		Researcher string `json:"researcher"`
		Grant      int    `json:"grant"`
	} `json:"meta"`
}

func NewConfDis(t *testing.T, rootKey string) *ConfDis {
	c, err := New("localhost:6379", "test:confdis:simple", SampleConfig{})
	if err != nil {
		t.Fatalf("Failed to connect to redis: %v", err)
	}
	return c
}

func redisDelay() {
	// Allow reasonable delay for network/redis latency in the above
	// save operation.
	time.Sleep(time.Duration(100 * time.Millisecond))
}

func TestSimple(t *testing.T) {
	c := NewConfDis(t, "test:confdis:simple")
	if err := c.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		config.Name = "primates"
		config.Users = []string{"chimp", "bonobo", "lemur"}
		config.Meta.Researcher = "Jane Goodall"
		config.Meta.Grant = 1200
		return nil
	}); err != nil {
		t.Fatal(err)
	}
}

func TestChangeNotification(t *testing.T) {
	// First client, with initial data.
	c := NewConfDis(t, "test:confdis:notify")
	go c.MustReceiveChanges()
	if err := c.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		config.Name = "primates-changes"
		config.Users = []string{"chimp", "bonobo", "lemur"}
		config.Meta.Researcher = "Jane Goodall"
		config.Meta.Grant = 1200
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	redisDelay()

	// Second client
	c2 := NewConfDis(t, "test:confdis:notify")
	go c2.MustReceiveChanges()

	if c2.Config.(*SampleConfig).Meta.Researcher != "Jane Goodall" {
		t.Fatal("different value")
	}

	// Trigger a change via the first client
	if err := c.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		config.Meta.Researcher = "Francine Patterson"
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	redisDelay()

	// Second client must get notified of that change
	if c2.Config.(*SampleConfig).Meta.Researcher != "Francine Patterson" {
		t.Fatal("did not receive change")
	}
}

func TestAtomicSave(t *testing.T) {
	// First client, with initial data.
	c := NewConfDis(t, "test:confdis:atomicsave")
	go c.MustReceiveChanges()
	if err := c.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		config.Name = "primates-changes"
		config.Users = []string{"chimp", "bonobo", "lemur"}
		config.Meta.Researcher = "Jane Goodall"
		config.Meta.Grant = 1200
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	redisDelay()

	// Trigger a change every 20 milliseconds
	go func() {
		for _ = range time.Tick(20 * time.Millisecond) {
			if err := c.AtomicSave(func(i interface{}) error {
				config := i.(*SampleConfig)
				config.Meta.Grant += 15
				return nil
			}); err != nil {
				t.Fatalf("Error in periodic-saving: %v", err)
			}
		}
	}()

	// Second client
	c2 := NewConfDis(t, "test:confdis:atomicsave")
	go c2.MustReceiveChanges()

	// Trigger a *slow* change, expecting write conflict.
	if err := c2.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		// Choose a delay value (50ms) greater than the frequency of
		// change (20ms) from the other client above.
		time.Sleep(50 * time.Millisecond)
		config.Meta.Researcher = "Francine Patterson"
		return nil
	}); err == nil {
		t.Fatal("Expecting this save to fail.")
	} else {
		// t.Logf("Failed as expected with: %v", err)
	}
}
<commit_msg>fix failing tests due to API change<commit_after>package confdis

import (
	"github.com/vmihailenco/redis"
	"net"
	"testing"
	"time"
)

type SampleConfig struct {
	Name  string   `json:"name"`
	Users []string `json:"users"`
	Meta  struct {
		Researcher string `json:"researcher"`
		Grant      int    `json:"grant"`
	} `json:"meta"`
}

func NewConfDis(t *testing.T, rootKey string) *ConfDis {
	redis, err := NewRedisClient("localhost:6379", "", 0)
	if err != nil {
		t.Fatalf("Unable to connect to redis: %v", err)
	}

	c, err := New(redis, "test:confdis:simple", SampleConfig{})
	if err != nil {
		t.Fatalf("Failed to connect to redis: %v", err)
	}
	return c
}

func redisDelay() {
	// Allow reasonable delay for network/redis latency in the above
	// save operation.
	time.Sleep(time.Duration(100 * time.Millisecond))
}

func TestSimple(t *testing.T) {
	c := NewConfDis(t, "test:confdis:simple")
	if err := c.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		config.Name = "primates"
		config.Users = []string{"chimp", "bonobo", "lemur"}
		config.Meta.Researcher = "Jane Goodall"
		config.Meta.Grant = 1200
		return nil
	}); err != nil {
		t.Fatal(err)
	}
}

func TestChangeNotification(t *testing.T) {
	// First client, with initial data.
	c := NewConfDis(t, "test:confdis:notify")
	go c.MustReceiveChanges()
	if err := c.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		config.Name = "primates-changes"
		config.Users = []string{"chimp", "bonobo", "lemur"}
		config.Meta.Researcher = "Jane Goodall"
		config.Meta.Grant = 1200
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	redisDelay()

	// Second client
	c2 := NewConfDis(t, "test:confdis:notify")
	go c2.MustReceiveChanges()

	if c2.Config.(*SampleConfig).Meta.Researcher != "Jane Goodall" {
		t.Fatal("different value")
	}

	// Trigger a change via the first client
	if err := c.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		config.Meta.Researcher = "Francine Patterson"
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	redisDelay()

	// Second client must get notified of that change
	if c2.Config.(*SampleConfig).Meta.Researcher != "Francine Patterson" {
		t.Fatal("did not receive change")
	}
}

func TestAtomicSave(t *testing.T) {
	// First client, with initial data.
	c := NewConfDis(t, "test:confdis:atomicsave")
	go c.MustReceiveChanges()
	if err := c.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		config.Name = "primates-changes"
		config.Users = []string{"chimp", "bonobo", "lemur"}
		config.Meta.Researcher = "Jane Goodall"
		config.Meta.Grant = 1200
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	redisDelay()

	// Trigger a change every 20 milliseconds
	go func() {
		for _ = range time.Tick(20 * time.Millisecond) {
			if err := c.AtomicSave(func(i interface{}) error {
				config := i.(*SampleConfig)
				config.Meta.Grant += 15
				return nil
			}); err != nil {
				t.Fatalf("Error in periodic-saving: %v", err)
			}
		}
	}()

	// Second client
	c2 := NewConfDis(t, "test:confdis:atomicsave")
	go c2.MustReceiveChanges()

	// Trigger a *slow* change, expecting write conflict.
	if err := c2.AtomicSave(func(i interface{}) error {
		config := i.(*SampleConfig)
		// Choose a delay value (50ms) greater than the frequency of
		// change (20ms) from the other client above.
		time.Sleep(50 * time.Millisecond)
		config.Meta.Researcher = "Francine Patterson"
		return nil
	}); err == nil {
		t.Fatal("Expecting this save to fail.")
	} else {
		// t.Logf("Failed as expected with: %v", err)
	}
}

// NewRedisClient connects to redis after ensuring that the server is
// indeed running.
func NewRedisClient(addr, password string, database int64) (*redis.Client, error) {
	// Bug #97459 -- is the redis client library faking connection for
	// the down server?
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}
	conn.Close()

	return redis.NewTCPClient(addr, password, database), nil
}
<|endoftext|>


<commit_before>package eg

// This file defines the AST rewriting pass.
// Most of it was plundered directly from
// $GOROOT/src/cmd/gofmt/rewrite.go (after convergent evolution).

import (
	"fmt"
	"go/ast"
	"go/token"
	"os"
	"reflect"
	"sort"
	"strconv"
	"strings"

	"code.google.com/p/go.tools/astutil"
	"code.google.com/p/go.tools/go/types"
)

// Transform applies the transformation to the specified parsed file,
// whose type information is supplied in info, and returns the number
// of replacements that were made.
//
// It mutates the AST in place (the identity of the root node is
// unchanged), and may add nodes for which no type information is
// available in info.
//
// Derived from rewriteFile in $GOROOT/src/cmd/gofmt/rewrite.go.
//
func (tr *Transformer) Transform(info *types.Info, pkg *types.Package, file *ast.File) int {
	if !tr.seenInfos[info] {
		tr.seenInfos[info] = true
		mergeTypeInfo(&tr.info.Info, info)
	}
	tr.currentPkg = pkg
	tr.nsubsts = 0

	if tr.verbose {
		fmt.Fprintf(os.Stderr, "before: %s\n", astString(tr.fset, tr.before))
		fmt.Fprintf(os.Stderr, "after: %s\n", astString(tr.fset, tr.after))
	}

	var f func(rv reflect.Value) reflect.Value
	f = func(rv reflect.Value) reflect.Value {
		// don't bother if val is invalid to start with
		if !rv.IsValid() {
			return reflect.Value{}
		}

		rv = apply(f, rv)

		e := rvToExpr(rv)
		if e != nil {
			savedEnv := tr.env
			tr.env = make(map[string]ast.Expr) // inefficient!  Use a slice of k/v pairs

			if tr.matchExpr(tr.before, e) {
				if tr.verbose {
					fmt.Fprintf(os.Stderr, "%s matches %s",
						astString(tr.fset, tr.before), astString(tr.fset, e))
					if len(tr.env) > 0 {
						fmt.Fprintf(os.Stderr, " with:")
						for name, ast := range tr.env {
							fmt.Fprintf(os.Stderr, " %s->%s",
								name, astString(tr.fset, ast))
						}
					}
					fmt.Fprintf(os.Stderr, "\n")
				}
				tr.nsubsts++

				// Clone the replacement tree, performing parameter substitution.
				// We update all positions to n.Pos() to aid comment placement.
				rv = tr.subst(tr.env, reflect.ValueOf(tr.after),
					reflect.ValueOf(e.Pos()))
			}
			tr.env = savedEnv
		}

		return rv
	}
	file2 := apply(f, reflect.ValueOf(file)).Interface().(*ast.File)

	// By construction, the root node is unchanged.
	if file != file2 {
		panic("BUG")
	}

	// Add any necessary imports.
	// TODO(adonovan): remove no-longer needed imports too.
	if tr.nsubsts > 0 {
		pkgs := make(map[string]*types.Package)
		for obj := range tr.importedObjs {
			pkgs[obj.Pkg().Path()] = obj.Pkg()
		}

		for _, imp := range file.Imports {
			path, _ := strconv.Unquote(imp.Path.Value)
			delete(pkgs, path)
		}
		delete(pkgs, pkg.Path()) // don't import self

		// NB: AddImport may completely replace the AST!
		// It thus renders info and tr.info no longer relevant to file.
		var paths []string
		for path := range pkgs {
			paths = append(paths, path)
		}
		sort.Strings(paths)
		for _, path := range paths {
			astutil.AddImport(tr.fset, file, path)
		}
	}

	tr.currentPkg = nil

	return tr.nsubsts
}

// setValue is a wrapper for x.SetValue(y); it protects
// the caller from panics if x cannot be changed to y.
func setValue(x, y reflect.Value) {
	// don't bother if y is invalid to start with
	if !y.IsValid() {
		return
	}
	defer func() {
		if x := recover(); x != nil {
			if s, ok := x.(string); ok &&
				(strings.Contains(s, "type mismatch") || strings.Contains(s, "not assignable")) {
				// x cannot be set to y - ignore this rewrite
				return
			}
			panic(x)
		}
	}()
	x.Set(y)
}

// Values/types for special cases.
var (
	objectPtrNil = reflect.ValueOf((*ast.Object)(nil))
	scopePtrNil  = reflect.ValueOf((*ast.Scope)(nil))

	identType        = reflect.TypeOf((*ast.Ident)(nil))
	selectorExprType = reflect.TypeOf((*ast.SelectorExpr)(nil))
	objectPtrType    = reflect.TypeOf((*ast.Object)(nil))
	positionType     = reflect.TypeOf(token.NoPos)
	callExprType     = reflect.TypeOf((*ast.CallExpr)(nil))
	scopePtrType     = reflect.TypeOf((*ast.Scope)(nil))
)

// apply replaces each AST field x in val with f(x), returning val.
// To avoid extra conversions, f operates on the reflect.Value form.
func apply(f func(reflect.Value) reflect.Value, val reflect.Value) reflect.Value {
	if !val.IsValid() {
		return reflect.Value{}
	}

	// *ast.Objects introduce cycles and are likely incorrect after
	// rewrite; don't follow them but replace with nil instead
	if val.Type() == objectPtrType {
		return objectPtrNil
	}

	// similarly for scopes: they are likely incorrect after a rewrite;
	// replace them with nil
	if val.Type() == scopePtrType {
		return scopePtrNil
	}

	switch v := reflect.Indirect(val); v.Kind() {
	case reflect.Slice:
		for i := 0; i < v.Len(); i++ {
			e := v.Index(i)
			setValue(e, f(e))
		}
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			e := v.Field(i)
			setValue(e, f(e))
		}
	case reflect.Interface:
		e := v.Elem()
		setValue(v, f(e))
	}
	return val
}

// subst returns a copy of (replacement) pattern with values from env
// substituted in place of wildcards and pos used as the position of
// tokens from the pattern.  if env == nil, subst returns a copy of
// pattern and doesn't change the line number information.
func (tr *Transformer) subst(env map[string]ast.Expr, pattern, pos reflect.Value) reflect.Value {
	if !pattern.IsValid() {
		return reflect.Value{}
	}

	// *ast.Objects introduce cycles and are likely incorrect after
	// rewrite; don't follow them but replace with nil instead
	if pattern.Type() == objectPtrType {
		return objectPtrNil
	}

	// similarly for scopes: they are likely incorrect after a rewrite;
	// replace them with nil
	if pattern.Type() == scopePtrType {
		return scopePtrNil
	}

	// Wildcard gets replaced with map value.
	if env != nil && pattern.Type() == identType {
		id := pattern.Interface().(*ast.Ident)
		if old, ok := env[id.Name]; ok {
			return tr.subst(nil, reflect.ValueOf(old), reflect.Value{})
		}
	}

	// Emit qualified identifiers in the pattern by appropriate
	// (possibly qualified) identifier in the input.
	//
	// The template cannot contain dot imports, so all identifiers
	// for imported objects are explicitly qualified.
	//
	// We assume (unsoundly) that there are no dot or named
	// imports in the input code, nor are any imported package
	// names shadowed, so the usual normal qualified identifier
	// syntax may be used.
	// TODO(adonovan): fix: avoid this assumption.
	//
	// A refactoring may be applied to a package referenced by the
	// template.  Objects belonging to the current package are
	// denoted by unqualified identifiers.
	//
	if tr.importedObjs != nil && pattern.Type() == selectorExprType {
		obj := isRef(pattern.Interface().(*ast.SelectorExpr), &tr.info)
		if obj != nil {
			if sel, ok := tr.importedObjs[obj]; ok {
				var id ast.Expr
				if obj.Pkg() == tr.currentPkg {
					id = sel.Sel // unqualified
				} else {
					id = sel // pkg-qualified
				}

				// Return a clone of id.
				saved := tr.importedObjs
				tr.importedObjs = nil // break cycle
				r := tr.subst(nil, reflect.ValueOf(id), pos)
				tr.importedObjs = saved
				return r
			}
		}
	}

	if pos.IsValid() && pattern.Type() == positionType {
		// use new position only if old position was valid in the first place
		if old := pattern.Interface().(token.Pos); !old.IsValid() {
			return pattern
		}
		return pos
	}

	// Otherwise copy.
	switch p := pattern; p.Kind() {
	case reflect.Slice:
		v := reflect.MakeSlice(p.Type(), p.Len(), p.Len())
		for i := 0; i < p.Len(); i++ {
			v.Index(i).Set(tr.subst(env, p.Index(i), pos))
		}
		return v

	case reflect.Struct:
		v := reflect.New(p.Type()).Elem()
		for i := 0; i < p.NumField(); i++ {
			v.Field(i).Set(tr.subst(env, p.Field(i), pos))
		}
		return v

	case reflect.Ptr:
		v := reflect.New(p.Type()).Elem()
		if elem := p.Elem(); elem.IsValid() {
			v.Set(tr.subst(env, elem, pos).Addr())
		}

		// Duplicate type information for duplicated ast.Expr.
		// All ast.Node implementations are *structs,
		// so this case catches them all.
		if e := rvToExpr(v); e != nil {
			updateTypeInfo(&tr.info.Info, e, p.Interface().(ast.Expr))
		}
		return v

	case reflect.Interface:
		v := reflect.New(p.Type()).Elem()
		if elem := p.Elem(); elem.IsValid() {
			v.Set(tr.subst(env, elem, pos))
		}
		return v
	}

	return pattern
}

// -- utilitiies -------------------------------------------------------

func rvToExpr(rv reflect.Value) ast.Expr {
	if rv.CanInterface() {
		if e, ok := rv.Interface().(ast.Expr); ok {
			return e
		}
	}
	return nil
}

// updateTypeInfo duplicates type information for the existing AST old
// so that it also applies to duplicated AST new.
func updateTypeInfo(info *types.Info, new, old ast.Expr) {
	switch new := new.(type) {
	case *ast.Ident:
		orig := old.(*ast.Ident)
		if obj, ok := info.Defs[orig]; ok {
			info.Defs[new] = obj
		}
		if obj, ok := info.Uses[orig]; ok {
			info.Uses[new] = obj
		}

	case *ast.SelectorExpr:
		orig := old.(*ast.SelectorExpr)
		if sel, ok := info.Selections[orig]; ok {
			info.Selections[new] = sel
		}
	}

	if tv, ok := info.Types[old]; ok {
		info.Types[new] = tv
	}
}

func F() {}
func G() {}

func init() {
	F()
}
<commit_msg>go.tools/refactor/eg: remove debugging code accidentally committed.<commit_after>package eg

// This file defines the AST rewriting pass.
// Most of it was plundered directly from
// $GOROOT/src/cmd/gofmt/rewrite.go (after convergent evolution).

import (
	"fmt"
	"go/ast"
	"go/token"
	"os"
	"reflect"
	"sort"
	"strconv"
	"strings"

	"code.google.com/p/go.tools/astutil"
	"code.google.com/p/go.tools/go/types"
)

// Transform applies the transformation to the specified parsed file,
// whose type information is supplied in info, and returns the number
// of replacements that were made.
//
// It mutates the AST in place (the identity of the root node is
// unchanged), and may add nodes for which no type information is
// available in info.
//
// Derived from rewriteFile in $GOROOT/src/cmd/gofmt/rewrite.go.
//
func (tr *Transformer) Transform(info *types.Info, pkg *types.Package, file *ast.File) int {
	if !tr.seenInfos[info] {
		tr.seenInfos[info] = true
		mergeTypeInfo(&tr.info.Info, info)
	}
	tr.currentPkg = pkg
	tr.nsubsts = 0

	if tr.verbose {
		fmt.Fprintf(os.Stderr, "before: %s\n", astString(tr.fset, tr.before))
		fmt.Fprintf(os.Stderr, "after: %s\n", astString(tr.fset, tr.after))
	}

	var f func(rv reflect.Value) reflect.Value
	f = func(rv reflect.Value) reflect.Value {
		// don't bother if val is invalid to start with
		if !rv.IsValid() {
			return reflect.Value{}
		}

		rv = apply(f, rv)

		e := rvToExpr(rv)
		if e != nil {
			savedEnv := tr.env
			tr.env = make(map[string]ast.Expr) // inefficient!  Use a slice of k/v pairs

			if tr.matchExpr(tr.before, e) {
				if tr.verbose {
					fmt.Fprintf(os.Stderr, "%s matches %s",
						astString(tr.fset, tr.before), astString(tr.fset, e))
					if len(tr.env) > 0 {
						fmt.Fprintf(os.Stderr, " with:")
						for name, ast := range tr.env {
							fmt.Fprintf(os.Stderr, " %s->%s",
								name, astString(tr.fset, ast))
						}
					}
					fmt.Fprintf(os.Stderr, "\n")
				}
				tr.nsubsts++

				// Clone the replacement tree, performing parameter substitution.
				// We update all positions to n.Pos() to aid comment placement.
				rv = tr.subst(tr.env, reflect.ValueOf(tr.after),
					reflect.ValueOf(e.Pos()))
			}
			tr.env = savedEnv
		}

		return rv
	}
	file2 := apply(f, reflect.ValueOf(file)).Interface().(*ast.File)

	// By construction, the root node is unchanged.
	if file != file2 {
		panic("BUG")
	}

	// Add any necessary imports.
	// TODO(adonovan): remove no-longer needed imports too.
	if tr.nsubsts > 0 {
		pkgs := make(map[string]*types.Package)
		for obj := range tr.importedObjs {
			pkgs[obj.Pkg().Path()] = obj.Pkg()
		}

		for _, imp := range file.Imports {
			path, _ := strconv.Unquote(imp.Path.Value)
			delete(pkgs, path)
		}
		delete(pkgs, pkg.Path()) // don't import self

		// NB: AddImport may completely replace the AST!
		// It thus renders info and tr.info no longer relevant to file.
		var paths []string
		for path := range pkgs {
			paths = append(paths, path)
		}
		sort.Strings(paths)
		for _, path := range paths {
			astutil.AddImport(tr.fset, file, path)
		}
	}

	tr.currentPkg = nil

	return tr.nsubsts
}

// setValue is a wrapper for x.SetValue(y); it protects
// the caller from panics if x cannot be changed to y.
func setValue(x, y reflect.Value) {
	// don't bother if y is invalid to start with
	if !y.IsValid() {
		return
	}
	defer func() {
		if x := recover(); x != nil {
			if s, ok := x.(string); ok &&
				(strings.Contains(s, "type mismatch") || strings.Contains(s, "not assignable")) {
				// x cannot be set to y - ignore this rewrite
				return
			}
			panic(x)
		}
	}()
	x.Set(y)
}

// Values/types for special cases.
var (
	objectPtrNil = reflect.ValueOf((*ast.Object)(nil))
	scopePtrNil  = reflect.ValueOf((*ast.Scope)(nil))

	identType        = reflect.TypeOf((*ast.Ident)(nil))
	selectorExprType = reflect.TypeOf((*ast.SelectorExpr)(nil))
	objectPtrType    = reflect.TypeOf((*ast.Object)(nil))
	positionType     = reflect.TypeOf(token.NoPos)
	callExprType     = reflect.TypeOf((*ast.CallExpr)(nil))
	scopePtrType     = reflect.TypeOf((*ast.Scope)(nil))
)

// apply replaces each AST field x in val with f(x), returning val.
// To avoid extra conversions, f operates on the reflect.Value form.
func apply(f func(reflect.Value) reflect.Value, val reflect.Value) reflect.Value {
	if !val.IsValid() {
		return reflect.Value{}
	}

	// *ast.Objects introduce cycles and are likely incorrect after
	// rewrite; don't follow them but replace with nil instead
	if val.Type() == objectPtrType {
		return objectPtrNil
	}

	// similarly for scopes: they are likely incorrect after a rewrite;
	// replace them with nil
	if val.Type() == scopePtrType {
		return scopePtrNil
	}

	switch v := reflect.Indirect(val); v.Kind() {
	case reflect.Slice:
		for i := 0; i < v.Len(); i++ {
			e := v.Index(i)
			setValue(e, f(e))
		}
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			e := v.Field(i)
			setValue(e, f(e))
		}
	case reflect.Interface:
		e := v.Elem()
		setValue(v, f(e))
	}
	return val
}

// subst returns a copy of (replacement) pattern with values from env
// substituted in place of wildcards and pos used as the position of
// tokens from the pattern.  if env == nil, subst returns a copy of
// pattern and doesn't change the line number information.
func (tr *Transformer) subst(env map[string]ast.Expr, pattern, pos reflect.Value) reflect.Value {
	if !pattern.IsValid() {
		return reflect.Value{}
	}

	// *ast.Objects introduce cycles and are likely incorrect after
	// rewrite; don't follow them but replace with nil instead
	if pattern.Type() == objectPtrType {
		return objectPtrNil
	}

	// similarly for scopes: they are likely incorrect after a rewrite;
	// replace them with nil
	if pattern.Type() == scopePtrType {
		return scopePtrNil
	}

	// Wildcard gets replaced with map value.
	if env != nil && pattern.Type() == identType {
		id := pattern.Interface().(*ast.Ident)
		if old, ok := env[id.Name]; ok {
			return tr.subst(nil, reflect.ValueOf(old), reflect.Value{})
		}
	}

	// Emit qualified identifiers in the pattern by appropriate
	// (possibly qualified) identifier in the input.
	//
	// The template cannot contain dot imports, so all identifiers
	// for imported objects are explicitly qualified.
	//
	// We assume (unsoundly) that there are no dot or named
	// imports in the input code, nor are any imported package
	// names shadowed, so the usual normal qualified identifier
	// syntax may be used.
	// TODO(adonovan): fix: avoid this assumption.
	//
	// A refactoring may be applied to a package referenced by the
	// template.  Objects belonging to the current package are
	// denoted by unqualified identifiers.
	//
	if tr.importedObjs != nil && pattern.Type() == selectorExprType {
		obj := isRef(pattern.Interface().(*ast.SelectorExpr), &tr.info)
		if obj != nil {
			if sel, ok := tr.importedObjs[obj]; ok {
				var id ast.Expr
				if obj.Pkg() == tr.currentPkg {
					id = sel.Sel // unqualified
				} else {
					id = sel // pkg-qualified
				}

				// Return a clone of id.
				saved := tr.importedObjs
				tr.importedObjs = nil // break cycle
				r := tr.subst(nil, reflect.ValueOf(id), pos)
				tr.importedObjs = saved
				return r
			}
		}
	}

	if pos.IsValid() && pattern.Type() == positionType {
		// use new position only if old position was valid in the first place
		if old := pattern.Interface().(token.Pos); !old.IsValid() {
			return pattern
		}
		return pos
	}

	// Otherwise copy.
	switch p := pattern; p.Kind() {
	case reflect.Slice:
		v := reflect.MakeSlice(p.Type(), p.Len(), p.Len())
		for i := 0; i < p.Len(); i++ {
			v.Index(i).Set(tr.subst(env, p.Index(i), pos))
		}
		return v

	case reflect.Struct:
		v := reflect.New(p.Type()).Elem()
		for i := 0; i < p.NumField(); i++ {
			v.Field(i).Set(tr.subst(env, p.Field(i), pos))
		}
		return v

	case reflect.Ptr:
		v := reflect.New(p.Type()).Elem()
		if elem := p.Elem(); elem.IsValid() {
			v.Set(tr.subst(env, elem, pos).Addr())
		}

		// Duplicate type information for duplicated ast.Expr.
		// All ast.Node implementations are *structs,
		// so this case catches them all.
		if e := rvToExpr(v); e != nil {
			updateTypeInfo(&tr.info.Info, e, p.Interface().(ast.Expr))
		}
		return v

	case reflect.Interface:
		v := reflect.New(p.Type()).Elem()
		if elem := p.Elem(); elem.IsValid() {
			v.Set(tr.subst(env, elem, pos))
		}
		return v
	}

	return pattern
}

// -- utilitiies -------------------------------------------------------

func rvToExpr(rv reflect.Value) ast.Expr {
	if rv.CanInterface() {
		if e, ok := rv.Interface().(ast.Expr); ok {
			return e
		}
	}
	return nil
}

// updateTypeInfo duplicates type information for the existing AST old
// so that it also applies to duplicated AST new.
func updateTypeInfo(info *types.Info, new, old ast.Expr) {
	switch new := new.(type) {
	case *ast.Ident:
		orig := old.(*ast.Ident)
		if obj, ok := info.Defs[orig]; ok {
			info.Defs[new] = obj
		}
		if obj, ok := info.Uses[orig]; ok {
			info.Uses[new] = obj
		}

	case *ast.SelectorExpr:
		orig := old.(*ast.SelectorExpr)
		if sel, ok := info.Selections[orig]; ok {
			info.Selections[new] = sel
		}
	}

	if tv, ok := info.Types[old]; ok {
		info.Types[new] = tv
	}
}
<|endoftext|>


<commit_before>// Copyright 2013, 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 worker

import (
	"flag"
	"fmt"
	"math/rand"
	"time"

	"github.com/youtube/vitess/go/vt/discovery"
	"github.com/youtube/vitess/go/vt/servenv"
	"github.com/youtube/vitess/go/vt/topo/topoproto"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

var (
	minHealthyEndPoints = flag.Int("min_healthy_rdonly_endpoints", 2, "minimum number of healthy rdonly endpoints before taking out one")

	// WaitForHealthyEndPointsTimeout intent is to wait for the
	// healthcheck to automatically return rdonly instances which
	// have been taken out by previous *Clone or *Diff runs.
	// Therefore, the default for this variable must be higher
	// than -health_check_interval.
	// (it is public for tests to override it)
	WaitForHealthyEndPointsTimeout = flag.Duration("wait_for_healthy_rdonly_endpoints_timeout", 60*time.Second, "maximum time to wait if less than --min_healthy_rdonly_endpoints are available")

	healthCheckTopologyRefresh = flag.Duration("worker_healthcheck_topology_refresh", 30*time.Second, "refresh interval for re-reading the topology")
	healthcheckRetryDelay      = flag.Duration("worker_healthcheck_retry_delay", 5*time.Second, "delay before retrying a failed healthcheck")
	healthCheckTimeout         = flag.Duration("worker_healthcheck_timeout", time.Minute, "the health check timeout period")
)

// FindHealthyRdonlyEndPoint returns a random healthy endpoint.
// Since we don't want to use them all, we require at least
// minHealthyEndPoints servers to be healthy.
// May block up to -wait_for_healthy_rdonly_endpoints_timeout.
func FindHealthyRdonlyEndPoint(ctx context.Context, wr *wrangler.Wrangler, cell, keyspace, shard string) (*topodatapb.TabletAlias, error) {
	busywaitCtx, busywaitCancel := context.WithTimeout(ctx, *WaitForHealthyEndPointsTimeout)
	defer busywaitCancel()

	// create a discovery healthcheck, wait for it to have one rdonly
	// endpoints at this point
	healthCheck := discovery.NewHealthCheck(*remoteActionsTimeout, *healthcheckRetryDelay, *healthCheckTimeout, "" /* statsSuffix */)
	watcher := discovery.NewShardReplicationWatcher(wr.TopoServer(), healthCheck, cell, keyspace, shard, *healthCheckTopologyRefresh, 5 /*topoReadConcurrency*/)
	defer watcher.Stop()
	defer healthCheck.Close()
	if err := discovery.WaitForEndPoints(ctx, healthCheck, cell, keyspace, shard, []topodatapb.TabletType{topodatapb.TabletType_RDONLY}); err != nil {
		return nil, fmt.Errorf("error waiting for rdonly endpoints for (%v,%v/%v): %v", cell, keyspace, shard, err)
	}

	var healthyEndpoints []*topodatapb.EndPoint
	for {
		select {
		case <-busywaitCtx.Done():
			return nil, fmt.Errorf("Not enough endpoints to choose from in (%v,%v/%v), have %v healthy ones, need at least %v Context Error: %v", cell, keyspace, shard, len(healthyEndpoints), *minHealthyEndPoints, busywaitCtx.Err())
		default:
		}

		addrs := healthCheck.GetEndPointStatsFromTarget(keyspace, shard, topodatapb.TabletType_RDONLY)
		healthyEndpoints = make([]*topodatapb.EndPoint, 0, len(addrs))
		for _, addr := range addrs {
			// Note we do not check the 'Serving' flag here.
			// This is mainly to avoid the case where we run a
			// Diff between a source and destination, and the source
			// is not serving (disabled by TabletControl).
			// When we switch the tablet to 'worker', it will
			// go back to serving state.
			if addr.Stats == nil || addr.Stats.HealthError != "" || addr.Stats.SecondsBehindMaster > 30 {
				continue
			}
			healthyEndpoints = append(healthyEndpoints, addr.EndPoint)
		}
		if len(healthyEndpoints) < *minHealthyEndPoints {
			deadlineForLog, _ := busywaitCtx.Deadline()
			wr.Logger().Infof("Waiting for enough endpoints to become available. available: %v required: %v Waiting up to %.1f more seconds.", len(healthyEndpoints), *minHealthyEndPoints, deadlineForLog.Sub(time.Now()).Seconds())
			// Block for 1 second because 2 seconds is the -health_check_interval flag value in integration tests.
			timer := time.NewTimer(1 * time.Second)
			select {
			case <-busywaitCtx.Done():
				timer.Stop()
			case <-timer.C:
			}
		} else {
			break
		}
	}

	// random server in the list is what we want
	index := rand.Intn(len(healthyEndpoints))
	return &topodatapb.TabletAlias{
		Cell: cell,
		Uid:  healthyEndpoints[index].Uid,
	}, nil
}

// FindWorkerTablet will:
// - find a rdonly instance in the keyspace / shard
// - mark it as worker
// - tag it with our worker process
func FindWorkerTablet(ctx context.Context, wr *wrangler.Wrangler, cleaner *wrangler.Cleaner, cell, keyspace, shard string) (*topodatapb.TabletAlias, error) {
	tabletAlias, err := FindHealthyRdonlyEndPoint(ctx, wr, cell, keyspace, shard)
	if err != nil {
		return nil, err
	}

	// We add the tag before calling ChangeSlaveType, so the destination
	// vttablet reloads the worker URL when it reloads the tablet.
	ourURL := servenv.ListeningURL.String()
	wr.Logger().Infof("Adding tag[worker]=%v to tablet %v", ourURL, topoproto.TabletAliasString(tabletAlias))
	shortCtx, cancel := context.WithTimeout(ctx, *remoteActionsTimeout)
	_, err = wr.TopoServer().UpdateTabletFields(shortCtx, tabletAlias, func(tablet *topodatapb.Tablet) error {
		if tablet.Tags == nil {
			tablet.Tags = make(map[string]string)
		}
		tablet.Tags["worker"] = ourURL
		return nil
	})
	cancel()
	if err != nil {
		return nil, err
	}
	// Using "defer" here because we remove the tag *before* calling
	// ChangeSlaveType back, so we need to record this tag change after the change
	// slave type change in the cleaner.
	defer wrangler.RecordTabletTagAction(cleaner, tabletAlias, "worker", "")

	wr.Logger().Infof("Changing tablet %v to '%v'", topoproto.TabletAliasString(tabletAlias), topodatapb.TabletType_WORKER)
	shortCtx, cancel = context.WithTimeout(ctx, *remoteActionsTimeout)
	err = wr.ChangeSlaveType(shortCtx, tabletAlias, topodatapb.TabletType_WORKER)
	cancel()
	if err != nil {
		return nil, err
	}

	// Record a clean-up action to take the tablet back to rdonly.
	// We will alter this one later on and let the tablet go back to
	// 'spare' if we have stopped replication for too long on it.
	wrangler.RecordChangeSlaveTypeAction(cleaner, tabletAlias, topodatapb.TabletType_RDONLY)
	return tabletAlias, nil
}

func init() {
	rand.Seed(time.Now().UnixNano())
}
<commit_msg>worker: Simplify code which waits for enough health rdonly tablets to become available.<commit_after>// Copyright 2013, 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 worker

import (
	"flag"
	"fmt"
	"math/rand"
	"time"

	"github.com/youtube/vitess/go/vt/discovery"
	"github.com/youtube/vitess/go/vt/servenv"
	"github.com/youtube/vitess/go/vt/topo/topoproto"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"

	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

var (
	minHealthyEndPoints = flag.Int("min_healthy_rdonly_endpoints", 2, "minimum number of healthy rdonly endpoints before taking out one")

	// WaitForHealthyEndPointsTimeout intent is to wait for the
	// healthcheck to automatically return rdonly instances which
	// have been taken out by previous *Clone or *Diff runs.
	// Therefore, the default for this variable must be higher
	// than -health_check_interval.
	// (it is public for tests to override it)
	WaitForHealthyEndPointsTimeout = flag.Duration("wait_for_healthy_rdonly_endpoints_timeout", 60*time.Second, "maximum time to wait if less than --min_healthy_rdonly_endpoints are available")

	healthCheckTopologyRefresh = flag.Duration("worker_healthcheck_topology_refresh", 30*time.Second, "refresh interval for re-reading the topology")
	healthcheckRetryDelay      = flag.Duration("worker_healthcheck_retry_delay", 5*time.Second, "delay before retrying a failed healthcheck")
	healthCheckTimeout         = flag.Duration("worker_healthcheck_timeout", time.Minute, "the health check timeout period")
)

// FindHealthyRdonlyEndPoint returns a random healthy endpoint.
// Since we don't want to use them all, we require at least
// minHealthyEndPoints servers to be healthy.
// May block up to -wait_for_healthy_rdonly_endpoints_timeout.
func FindHealthyRdonlyEndPoint(ctx context.Context, wr *wrangler.Wrangler, cell, keyspace, shard string) (*topodatapb.TabletAlias, error) {
	busywaitCtx, busywaitCancel := context.WithTimeout(ctx, *WaitForHealthyEndPointsTimeout)
	defer busywaitCancel()

	// create a discovery healthcheck, wait for it to have one rdonly
	// endpoints at this point
	healthCheck := discovery.NewHealthCheck(*remoteActionsTimeout, *healthcheckRetryDelay, *healthCheckTimeout, "" /* statsSuffix */)
	watcher := discovery.NewShardReplicationWatcher(wr.TopoServer(), healthCheck, cell, keyspace, shard, *healthCheckTopologyRefresh, 5 /*topoReadConcurrency*/)
	defer watcher.Stop()
	defer healthCheck.Close()
	if err := discovery.WaitForEndPoints(ctx, healthCheck, cell, keyspace, shard, []topodatapb.TabletType{topodatapb.TabletType_RDONLY}); err != nil {
		return nil, fmt.Errorf("error waiting for rdonly endpoints for (%v,%v/%v): %v", cell, keyspace, shard, err)
	}

	var healthyEndpoints []*topodatapb.EndPoint
	for {
		select {
		case <-busywaitCtx.Done():
			return nil, fmt.Errorf("Not enough endpoints to choose from in (%v,%v/%v), have %v healthy ones, need at least %v Context Error: %v", cell, keyspace, shard, len(healthyEndpoints), *minHealthyEndPoints, busywaitCtx.Err())
		default:
		}

		addrs := healthCheck.GetEndPointStatsFromTarget(keyspace, shard, topodatapb.TabletType_RDONLY)
		healthyEndpoints = make([]*topodatapb.EndPoint, 0, len(addrs))
		for _, addr := range addrs {
			// Note we do not check the 'Serving' flag here.
			// This is mainly to avoid the case where we run a
			// Diff between a source and destination, and the source
			// is not serving (disabled by TabletControl).
			// When we switch the tablet to 'worker', it will
			// go back to serving state.
			if addr.Stats == nil || addr.Stats.HealthError != "" || addr.Stats.SecondsBehindMaster > 30 {
				continue
			}
			healthyEndpoints = append(healthyEndpoints, addr.EndPoint)
		}

		if len(healthyEndpoints) >= *minHealthyEndPoints {
			break
		}

		deadlineForLog, _ := busywaitCtx.Deadline()
		wr.Logger().Infof("Waiting for enough endpoints to become available. available: %v required: %v Waiting up to %.1f more seconds.", len(healthyEndpoints), *minHealthyEndPoints, deadlineForLog.Sub(time.Now()).Seconds())
		// Block for 1 second because 2 seconds is the -health_check_interval flag value in integration tests.
		timer := time.NewTimer(1 * time.Second)
		select {
		case <-busywaitCtx.Done():
			timer.Stop()
		case <-timer.C:
		}
	}

	// random server in the list is what we want
	index := rand.Intn(len(healthyEndpoints))
	return &topodatapb.TabletAlias{
		Cell: cell,
		Uid:  healthyEndpoints[index].Uid,
	}, nil
}

// FindWorkerTablet will:
// - find a rdonly instance in the keyspace / shard
// - mark it as worker
// - tag it with our worker process
func FindWorkerTablet(ctx context.Context, wr *wrangler.Wrangler, cleaner *wrangler.Cleaner, cell, keyspace, shard string) (*topodatapb.TabletAlias, error) {
	tabletAlias, err := FindHealthyRdonlyEndPoint(ctx, wr, cell, keyspace, shard)
	if err != nil {
		return nil, err
	}

	// We add the tag before calling ChangeSlaveType, so the destination
	// vttablet reloads the worker URL when it reloads the tablet.
	ourURL := servenv.ListeningURL.String()
	wr.Logger().Infof("Adding tag[worker]=%v to tablet %v", ourURL, topoproto.TabletAliasString(tabletAlias))
	shortCtx, cancel := context.WithTimeout(ctx, *remoteActionsTimeout)
	_, err = wr.TopoServer().UpdateTabletFields(shortCtx, tabletAlias, func(tablet *topodatapb.Tablet) error {
		if tablet.Tags == nil {
			tablet.Tags = make(map[string]string)
		}
		tablet.Tags["worker"] = ourURL
		return nil
	})
	cancel()
	if err != nil {
		return nil, err
	}
	// Using "defer" here because we remove the tag *before* calling
	// ChangeSlaveType back, so we need to record this tag change after the change
	// slave type change in the cleaner.
	defer wrangler.RecordTabletTagAction(cleaner, tabletAlias, "worker", "")

	wr.Logger().Infof("Changing tablet %v to '%v'", topoproto.TabletAliasString(tabletAlias), topodatapb.TabletType_WORKER)
	shortCtx, cancel = context.WithTimeout(ctx, *remoteActionsTimeout)
	err = wr.ChangeSlaveType(shortCtx, tabletAlias, topodatapb.TabletType_WORKER)
	cancel()
	if err != nil {
		return nil, err
	}

	// Record a clean-up action to take the tablet back to rdonly.
	// We will alter this one later on and let the tablet go back to
	// 'spare' if we have stopped replication for too long on it.
	wrangler.RecordChangeSlaveTypeAction(cleaner, tabletAlias, topodatapb.TabletType_RDONLY)
	return tabletAlias, nil
}

func init() {
	rand.Seed(time.Now().UnixNano())
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Jennal(jennalcn@gmail.com). All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// 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 event

import "sync"

type IEvent interface {
	On(string, interface{}, EventFunc)
	Off(string, interface{})
	Once(string, interface{}, EventFunc)
	Emit(string, ...interface{})
}

type EventFunc interface{}

type Event struct {
	sync.Mutex
	cbs map[string][]*Method
}

func NewEvent() *Event {
	return &Event{
		cbs: make(map[string][]*Method),
	}
}

func (self Event) On(name string, ins interface{}, cb EventFunc) {
	self.Lock()
	defer self.Unlock()

	list, ok := self.cbs[name]
	if !ok {
		list = []*Method{
			NewMethod(ins, cb),
		}
	} else {
		list = append(list, NewMethod(ins, cb))
	}
	self.cbs[name] = list
}

func (self Event) Off(name string, ins interface{}) {
	self.Lock()
	defer self.Unlock()

	list, ok := self.cbs[name]
	if !ok {
		return
	}

	for i := len(list) - 1; i >= 0; i-- {
		item := list[i]
		if item.caller != ins {
			continue
		}

		if i+1 > len(list) {
			list = list[:i]
		} else {
			list = append(list[:i], list[i+1:]...)
		}
	}
	self.cbs[name] = list
}

func (self Event) Once(name string, ins interface{}, cb EventFunc) {
	self.On(name, ins, func(args ...interface{}) {
		NewMethod(ins, cb).Call(args...)
		self.Off(name, ins)
	})
}

func (self Event) Emit(name string, args ...interface{}) {
	list, ok := self.cbs[name]
	if !ok {
		return
	}

	for _, item := range list {
		item.Call(args...)
	}
}
<commit_msg>event comments<commit_after>// Copyright (C) 2017 Jennal(jennalcn@gmail.com). All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// 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 event provide a nodejs like event emmiter
package event

import "sync"

type IEvent interface {
	On(string, interface{}, EventFunc)
	Off(string, interface{})
	Once(string, interface{}, EventFunc)
	Emit(string, ...interface{})
}

type EventFunc interface{}

type Event struct {
	sync.Mutex
	cbs map[string][]*Method
}

func NewEvent() *Event {
	return &Event{
		cbs: make(map[string][]*Method),
	}
}

func (self Event) On(name string, ins interface{}, cb EventFunc) {
	self.Lock()
	defer self.Unlock()

	list, ok := self.cbs[name]
	if !ok {
		list = []*Method{
			NewMethod(ins, cb),
		}
	} else {
		list = append(list, NewMethod(ins, cb))
	}
	self.cbs[name] = list
}

func (self Event) Off(name string, ins interface{}) {
	self.Lock()
	defer self.Unlock()

	list, ok := self.cbs[name]
	if !ok {
		return
	}

	for i := len(list) - 1; i >= 0; i-- {
		item := list[i]
		if item.caller != ins {
			continue
		}

		if i+1 > len(list) {
			list = list[:i]
		} else {
			list = append(list[:i], list[i+1:]...)
		}
	}
	self.cbs[name] = list
}

func (self Event) Once(name string, ins interface{}, cb EventFunc) {
	self.On(name, ins, func(args ...interface{}) {
		NewMethod(ins, cb).Call(args...)
		self.Off(name, ins)
	})
}

func (self Event) Emit(name string, args ...interface{}) {
	list, ok := self.cbs[name]
	if !ok {
		return
	}

	for _, item := range list {
		item.Call(args...)
	}
}
<|endoftext|>