<commit_before>package main

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"strconv"
	"strings"

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

// goproxy main purpose is to be a daemon connecting the docker daemon
// (remote API) and the custom Minecraft server (cubrite using lua scripts).
// But the cuberite lua scripts also execute goproxy as a short lived process
// to send requests to the goproxy daemon. If the program is executed without
// additional arguments it is the daemon "mode".
//
// As a daemon, goproxy listens for events from the docker daemon and send
// them to the cuberite server. It also listen for requests from the
// cuberite server, convert them into docker daemon remote API calls and send
// them to the docker daemon.

// Daemon maintains state when the dockercraft daemon is running
type Daemon struct {
	// Client is an instance of the DockerClient
	Client *dockerclient.DockerClient
	// Version is the version of the Docker Daemon
	Version string
	// BinaryName is the name of the Docker Binary
	BinaryName string
	// previouscpustats is a map containing the previous cpu stats we got from the
	// docker daemon through the docker remote api
	previousCPUStats map[string]*CPUStats
}

// NewDaemon returns a new instance of Daemon
func NewDaemon() *Daemon {
	return &Daemon{
		previousCPUStats: make(map[string]*CPUStats),
	}
}

// CPUStats contains the Total and System CPU stats
type CPUStats struct {
	TotalUsage  uint64
	SystemUsage uint64
}

// ProxyCmd submits a command to a running Daemon
func ProxyCmd(args []string) error {
	if len(args) != 1 {
		return fmt.Errorf("Only 1 argument expected. Received %d", len(args))
	}
	reqPath := "http://127.0.0.1:8000/" + os.Args[1]
	resp, err := http.Get(reqPath)
	if err != nil {
		return fmt.Errorf("Error on request: %s, ERROR: %s", reqPath, err.Error())
	}
	log.Infof("Request sent: %s, StatusCode: %d", reqPath, resp.StatusCode)
	return nil
}

// Init initializes a Daemon
func (d *Daemon) Init() error {
	var err error
	d.Client, err = dockerclient.NewDockerClient("unix:///var/run/docker.sock", nil)
	if err != nil {
		return err
	}

	// get the version of the docker remote API
	version, err := d.Client.Version()
	if err != nil {
		return err
	}
	d.Version = version.Version
	return nil
}

// eventCallback receives and handles the docker events
func (d *Daemon) eventCallback(event *dockerclient.Event, ec chan error, args ...interface{}) {
	log.Debugf("--\n%+v", *event)

	id := event.ID

	switch event.Status {
	case "create":
		log.Debug("create event")

		repo, tag := splitRepoAndTag(event.From)
		containerName := "<name>"
		containerInfo, err := d.Client.InspectContainer(id)
		if err != nil {
			log.Errorf("InspectContainer error: %s", err.Error())
		} else {
			containerName = containerInfo.Name
		}

		data := url.Values{
			"action":    {"createContainer"},
			"id":        {id},
			"name":      {containerName},
			"imageRepo": {repo},
			"imageTag":  {tag}}

		CuberiteServerRequest(data)

	case "start":
		log.Debugln("start event")

		repo, tag := splitRepoAndTag(event.From)
		containerName := "<name>"
		containerInfo, err := d.Client.InspectContainer(id)
		if err != nil {
			log.Errorf("InspectContainer error: %s", err.Error())
		} else {
			containerName = containerInfo.Name
		}

		data := url.Values{
			"action":    {"startContainer"},
			"id":        {id},
			"name":      {containerName},
			"imageRepo": {repo},
			"imageTag":  {tag}}

		// Monitor stats
		d.Client.StartMonitorStats(id, d.statCallback, nil)
		CuberiteServerRequest(data)

	case "stop":
		// die event is enough
		// http://docs.docker.com/reference/api/docker_remote_api/#docker-events

	case "restart":
		// start event is enough
		// http://docs.docker.com/reference/api/docker_remote_api/#docker-events

	case "kill":
		// die event is enough
		// http://docs.docker.com/reference/api/docker_remote_api/#docker-events

	case "die":
		log.Debugln("die event")

		// same as stop event
		repo, tag := splitRepoAndTag(event.From)
		containerName := "<name>"
		containerInfo, err := d.Client.InspectContainer(id)
		if err != nil {
			log.Errorf("InspectContainer error: %s", err.Error())
		} else {
			containerName = containerInfo.Name
		}

		data := url.Values{
			"action":    {"stopContainer"},
			"id":        {id},
			"name":      {containerName},
			"imageRepo": {repo},
			"imageTag":  {tag}}

		CuberiteServerRequest(data)

	case "destroy":
		log.Debug("destroy event")

		data := url.Values{
			"action": {"destroyContainer"},
			"id":     {id},
		}

		CuberiteServerRequest(data)
	}
}

// statCallback receives the stats (cpu & ram) from containers and send them to
// the cuberite server
func (d *Daemon) statCallback(id string, stat *dockerclient.Stats, ec chan error, args ...interface{}) {

	// log.Debugln("STATS", id, stat)
	// log.Debugln("---")
	// log.Debugln("cpu :", float64(stat.CpuStats.CpuUsage.TotalUsage)/float64(stat.CpuStats.SystemUsage))
	// log.Debugln("ram :", stat.MemoryStats.Usage)

	memPercent := float64(stat.MemoryStats.Usage) / float64(stat.MemoryStats.Limit) * 100.0
	var cpuPercent float64
	if preCPUStats, exists := d.previousCPUStats[id]; exists {
		cpuPercent = calculateCPUPercent(preCPUStats, &stat.CpuStats)
	}

	d.previousCPUStats[id] = &CPUStats{TotalUsage: stat.CpuStats.CpuUsage.TotalUsage, SystemUsage: stat.CpuStats.SystemUsage}

	data := url.Values{
		"action": {"stats"},
		"id":     {id},
		"cpu":    {strconv.FormatFloat(cpuPercent, 'f', 2, 64) + "%"},
		"ram":    {strconv.FormatFloat(memPercent, 'f', 2, 64) + "%"}}

	CuberiteServerRequest(data)
}

// execCmd handles http requests received for the path "/exec"
func (d *Daemon) execCmd(w http.ResponseWriter, r *http.Request) {

	io.WriteString(w, "OK")

	go func() {
		cmd := r.URL.Query().Get("cmd")
		cmd, _ = url.QueryUnescape(cmd)
		arr := strings.Split(cmd, " ")
		if len(arr) > 0 {

			if arr[0] == "docker" {
				arr[0] = d.BinaryName
			}

			cmd := exec.Command(arr[0], arr[1:]...)
			// Stdout buffer
			// cmdOutput := &bytes.Buffer{}
			// Attach buffer to command
			// cmd.Stdout = cmdOutput
			// Execute command
			// printCommand(cmd)
			err := cmd.Run() // will wait for command to return
			if err != nil {
				log.Error(err.Error())
			}
		}
	}()
}

// listContainers handles and reply to http requests having the path "/containers"
func (d *Daemon) listContainers(w http.ResponseWriter, r *http.Request) {

	// answer right away to avoid dead locks in LUA
	io.WriteString(w, "OK")

	go func() {
		containers, err := d.Client.ListContainers(true, false, "")

		if err != nil {
			log.Error(err.Error())
			return
		}

		images, err := d.Client.ListImages(true)

		if err != nil {
			log.Error(err.Error())
			return
		}

		for i := 0; i < len(containers); i++ {

			id := containers[i].Id
			info, _ := d.Client.InspectContainer(id)
			name := info.Name[1:]
			imageRepo := ""
			imageTag := ""

			for _, image := range images {
				if image.Id == info.Image {
					if len(image.RepoTags) > 0 {
						imageRepo, imageTag = splitRepoAndTag(image.RepoTags[0])
					}
					break
				}
			}

			data := url.Values{
				"action":    {"containerInfos"},
				"id":        {id},
				"name":      {name},
				"imageRepo": {imageRepo},
				"imageTag":  {imageTag},
				"running":   {strconv.FormatBool(info.State.Running)},
			}

			CuberiteServerRequest(data)

			if info.State.Running {
				// Monitor stats
				d.Client.StartMonitorStats(id, d.statCallback, nil)
			}
		}
	}()
}

// Utility functions

func calculateCPUPercent(previousCPUStats *CPUStats, newCPUStats *dockerclient.CpuStats) float64 {
	var (
		cpuPercent = 0.0
		// calculate the change for the cpu usage of the container in between readings
		cpuDelta = float64(newCPUStats.CpuUsage.TotalUsage - previousCPUStats.TotalUsage)
		// calculate the change for the entire system between readings
		systemDelta = float64(newCPUStats.SystemUsage - previousCPUStats.SystemUsage)
	)

	if systemDelta > 0.0 && cpuDelta > 0.0 {
		cpuPercent = (cpuDelta / systemDelta) * float64(len(newCPUStats.CpuUsage.PercpuUsage)) * 100.0
	}
	return cpuPercent
}

func splitRepoAndTag(repoTag string) (string, string) {

	repo := ""
	tag := ""

	repoAndTag := strings.Split(repoTag, ":")

	if len(repoAndTag) > 0 {
		repo = repoAndTag[0]
	}

	if len(repoAndTag) > 1 {
		tag = repoAndTag[1]
	}

	return repo, tag
}

// CuberiteServerRequest send a POST request that will be handled
// by our Cuberite Docker plugin.
func CuberiteServerRequest(data url.Values) {
	client := &http.Client{}
	req, _ := http.NewRequest("POST", "http://127.0.0.1:8080/webadmin/Docker/Docker", strings.NewReader(data.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.SetBasicAuth("admin", "admin")
	client.Do(req)
}
<commit_msg>Fix concurrency issue with cpu stats<commit_after>package main

import (
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"

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

// goproxy main purpose is to be a daemon connecting the docker daemon
// (remote API) and the custom Minecraft server (cubrite using lua scripts).
// But the cuberite lua scripts also execute goproxy as a short lived process
// to send requests to the goproxy daemon. If the program is executed without
// additional arguments it is the daemon "mode".
//
// As a daemon, goproxy listens for events from the docker daemon and send
// them to the cuberite server. It also listen for requests from the
// cuberite server, convert them into docker daemon remote API calls and send
// them to the docker daemon.

// Daemon maintains state when the dockercraft daemon is running
type Daemon struct {
	// Client is an instance of the DockerClient
	Client *dockerclient.DockerClient
	// Version is the version of the Docker Daemon
	Version string
	// BinaryName is the name of the Docker Binary
	BinaryName string
	// previouscpustats is a map containing the previous cpu stats we got from the
	// docker daemon through the docker remote api
	previousCPUStats map[string]*CPUStats
	sync.Mutex
}

// NewDaemon returns a new instance of Daemon
func NewDaemon() *Daemon {
	return &Daemon{
		previousCPUStats: make(map[string]*CPUStats),
	}
}

// CPUStats contains the Total and System CPU stats
type CPUStats struct {
	TotalUsage  uint64
	SystemUsage uint64
}

// ProxyCmd submits a command to a running Daemon
func ProxyCmd(args []string) error {
	if len(args) != 1 {
		return fmt.Errorf("Only 1 argument expected. Received %d", len(args))
	}
	reqPath := "http://127.0.0.1:8000/" + os.Args[1]
	resp, err := http.Get(reqPath)
	if err != nil {
		return fmt.Errorf("Error on request: %s, ERROR: %s", reqPath, err.Error())
	}
	log.Infof("Request sent: %s, StatusCode: %d", reqPath, resp.StatusCode)
	return nil
}

// Init initializes a Daemon
func (d *Daemon) Init() error {
	var err error
	d.Client, err = dockerclient.NewDockerClient("unix:///var/run/docker.sock", nil)
	if err != nil {
		return err
	}

	// get the version of the docker remote API
	version, err := d.Client.Version()
	if err != nil {
		return err
	}
	d.Version = version.Version
	return nil
}

// eventCallback receives and handles the docker events
func (d *Daemon) eventCallback(event *dockerclient.Event, ec chan error, args ...interface{}) {
	log.Debugf("--\n%+v", *event)

	id := event.ID

	switch event.Status {
	case "create":
		log.Debug("create event")

		repo, tag := splitRepoAndTag(event.From)
		containerName := "<name>"
		containerInfo, err := d.Client.InspectContainer(id)
		if err != nil {
			log.Errorf("InspectContainer error: %s", err.Error())
		} else {
			containerName = containerInfo.Name
		}

		data := url.Values{
			"action":    {"createContainer"},
			"id":        {id},
			"name":      {containerName},
			"imageRepo": {repo},
			"imageTag":  {tag}}

		CuberiteServerRequest(data)

	case "start":
		log.Debugln("start event")

		repo, tag := splitRepoAndTag(event.From)
		containerName := "<name>"
		containerInfo, err := d.Client.InspectContainer(id)
		if err != nil {
			log.Errorf("InspectContainer error: %s", err.Error())
		} else {
			containerName = containerInfo.Name
		}

		data := url.Values{
			"action":    {"startContainer"},
			"id":        {id},
			"name":      {containerName},
			"imageRepo": {repo},
			"imageTag":  {tag}}

		// Monitor stats
		d.Client.StartMonitorStats(id, d.statCallback, nil)
		CuberiteServerRequest(data)

	case "stop":
		// die event is enough
		// http://docs.docker.com/reference/api/docker_remote_api/#docker-events

	case "restart":
		// start event is enough
		// http://docs.docker.com/reference/api/docker_remote_api/#docker-events

	case "kill":
		// die event is enough
		// http://docs.docker.com/reference/api/docker_remote_api/#docker-events

	case "die":
		log.Debugln("die event")

		// same as stop event
		repo, tag := splitRepoAndTag(event.From)
		containerName := "<name>"
		containerInfo, err := d.Client.InspectContainer(id)
		if err != nil {
			log.Errorf("InspectContainer error: %s", err.Error())
		} else {
			containerName = containerInfo.Name
		}

		data := url.Values{
			"action":    {"stopContainer"},
			"id":        {id},
			"name":      {containerName},
			"imageRepo": {repo},
			"imageTag":  {tag}}

		CuberiteServerRequest(data)

	case "destroy":
		log.Debug("destroy event")

		data := url.Values{
			"action": {"destroyContainer"},
			"id":     {id},
		}

		CuberiteServerRequest(data)
	}
}

// statCallback receives the stats (cpu & ram) from containers and send them to
// the cuberite server
func (d *Daemon) statCallback(id string, stat *dockerclient.Stats, ec chan error, args ...interface{}) {

	// log.Debugln("STATS", id, stat)
	// log.Debugln("---")
	// log.Debugln("cpu :", float64(stat.CpuStats.CpuUsage.TotalUsage)/float64(stat.CpuStats.SystemUsage))
	// log.Debugln("ram :", stat.MemoryStats.Usage)
	d.Lock()
	defer d.Unlock()
	memPercent := float64(stat.MemoryStats.Usage) / float64(stat.MemoryStats.Limit) * 100.0
	var cpuPercent float64
	if preCPUStats, exists := d.previousCPUStats[id]; exists {
		cpuPercent = calculateCPUPercent(preCPUStats, &stat.CpuStats)
	}

	d.previousCPUStats[id] = &CPUStats{TotalUsage: stat.CpuStats.CpuUsage.TotalUsage, SystemUsage: stat.CpuStats.SystemUsage}

	data := url.Values{
		"action": {"stats"},
		"id":     {id},
		"cpu":    {strconv.FormatFloat(cpuPercent, 'f', 2, 64) + "%"},
		"ram":    {strconv.FormatFloat(memPercent, 'f', 2, 64) + "%"}}

	CuberiteServerRequest(data)
}

// execCmd handles http requests received for the path "/exec"
func (d *Daemon) execCmd(w http.ResponseWriter, r *http.Request) {

	io.WriteString(w, "OK")

	go func() {
		cmd := r.URL.Query().Get("cmd")
		cmd, _ = url.QueryUnescape(cmd)
		arr := strings.Split(cmd, " ")
		if len(arr) > 0 {

			if arr[0] == "docker" {
				arr[0] = d.BinaryName
			}

			cmd := exec.Command(arr[0], arr[1:]...)
			// Stdout buffer
			// cmdOutput := &bytes.Buffer{}
			// Attach buffer to command
			// cmd.Stdout = cmdOutput
			// Execute command
			// printCommand(cmd)
			err := cmd.Run() // will wait for command to return
			if err != nil {
				log.Error(err.Error())
			}
		}
	}()
}

// listContainers handles and reply to http requests having the path "/containers"
func (d *Daemon) listContainers(w http.ResponseWriter, r *http.Request) {

	// answer right away to avoid dead locks in LUA
	io.WriteString(w, "OK")

	go func() {
		containers, err := d.Client.ListContainers(true, false, "")

		if err != nil {
			log.Error(err.Error())
			return
		}

		images, err := d.Client.ListImages(true)

		if err != nil {
			log.Error(err.Error())
			return
		}

		for i := 0; i < len(containers); i++ {

			id := containers[i].Id
			info, _ := d.Client.InspectContainer(id)
			name := info.Name[1:]
			imageRepo := ""
			imageTag := ""

			for _, image := range images {
				if image.Id == info.Image {
					if len(image.RepoTags) > 0 {
						imageRepo, imageTag = splitRepoAndTag(image.RepoTags[0])
					}
					break
				}
			}

			data := url.Values{
				"action":    {"containerInfos"},
				"id":        {id},
				"name":      {name},
				"imageRepo": {imageRepo},
				"imageTag":  {imageTag},
				"running":   {strconv.FormatBool(info.State.Running)},
			}

			CuberiteServerRequest(data)

			if info.State.Running {
				// Monitor stats
				d.Client.StartMonitorStats(id, d.statCallback, nil)
			}
		}
	}()
}

// Utility functions

func calculateCPUPercent(previousCPUStats *CPUStats, newCPUStats *dockerclient.CpuStats) float64 {
	var (
		cpuPercent = 0.0
		// calculate the change for the cpu usage of the container in between readings
		cpuDelta = float64(newCPUStats.CpuUsage.TotalUsage - previousCPUStats.TotalUsage)
		// calculate the change for the entire system between readings
		systemDelta = float64(newCPUStats.SystemUsage - previousCPUStats.SystemUsage)
	)

	if systemDelta > 0.0 && cpuDelta > 0.0 {
		cpuPercent = (cpuDelta / systemDelta) * float64(len(newCPUStats.CpuUsage.PercpuUsage)) * 100.0
	}
	return cpuPercent
}

func splitRepoAndTag(repoTag string) (string, string) {

	repo := ""
	tag := ""

	repoAndTag := strings.Split(repoTag, ":")

	if len(repoAndTag) > 0 {
		repo = repoAndTag[0]
	}

	if len(repoAndTag) > 1 {
		tag = repoAndTag[1]
	}

	return repo, tag
}

// CuberiteServerRequest send a POST request that will be handled
// by our Cuberite Docker plugin.
func CuberiteServerRequest(data url.Values) {
	client := &http.Client{}
	req, _ := http.NewRequest("POST", "http://127.0.0.1:8080/webadmin/Docker/Docker", strings.NewReader(data.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.SetBasicAuth("admin", "admin")
	client.Do(req)
}
<|endoftext|>


<commit_before>package corepxe

import (
	"log"
	"net/http"

	"github.com/elazarl/goproxy"
)

func proxySetup() {
	proxy := goproxy.NewProxyHttpServer()
	proxy.OnRequest(UrlIs("public.update.core-os.net/v1/update/")).HandleConnect(goproxy.AlwaysMitm)
	proxy.OnResponse().DoFunc(
		func(r *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
			r.Header.Set("X-COREPXE", "corepxe")
			return r
		})
	proxy.Verbose = true
	log.Fatal(http.ListenAndServe(":8080", proxy))
}
<commit_msg>init responseHandler<commit_after>package corepxe

import (
	"log"
	"net/http"

	"github.com/elazarl/goproxy"
)

func reponseHandler(r *http.Response, ctx *goproxy.ProxyCtx, corpxeChan chan []int) *http.Response {
	r.Header.Set("X-COREPXE", "corepxe")
	return r
}

func proxySetup(corpxeChan chan []int) {
	proxy := goproxy.NewProxyHttpServer()
	proxy.OnRequest(UrlIs("public.update.core-os.net/v1/update/")).HandleConnect(goproxy.AlwaysMitm)
	proxy.OnResponse().DoFunc(reponseHandler(corpxeChan))
	proxy.Verbose = true
	log.Fatal(http.ListenAndServe(":8080", proxy))
}
<|endoftext|>


<commit_before>package inputs

import (
	"html"
	"log"
	"strings"

	"github.com/hexbotio/hex/models"
	"github.com/nlopes/slack"
)

// Slack struct
type Slack struct {
}

// Read function
func (x Slack) Read(inputMsgs chan<- models.Message, service models.Service) {
	defer Recovery(service)
	api := slack.New(service.Config["Key"])
	var slackDebug = false
	if strings.ToLower(service.Config["SlackDebug"]) == "true" {
		slackDebug = true
	}
	api.SetDebug(slackDebug)

	// get channel array
	channels := make(map[string]string)
	channelList, err := api.GetChannels(true)
	if err != nil {
		log.Printf("%s\n", err)
	}
	for _, channel := range channelList {
		channels[channel.ID] = "#" + channel.Name
	}

	// get user array
	bot := ""
	users := make(map[string]string)
	userList, err := api.GetUsers()
	if err != nil {
		log.Printf("%s\n", err)
	}
	for _, user := range userList {
		users[user.ID] = user.Name
		if service.BotName == user.Name {
			bot = user.ID
		}
	}

	// validate bot
	if bot == "" {
		log.Print("WARNING - Bot does not seem to be configured in Slack as '" + service.BotName + "'")
	} else {
		bot = "<@" + bot + ">"
	}

	// listen to messages
	rtm := api.NewRTM()
	go rtm.ManageConnection()
	for {
		select {
		case msg := <-rtm.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.MessageEvent:
				if ev.User != "" {

					channel, found := channels[ev.Channel]
					if !found {
						channel = ev.Channel
					}

					if strings.HasPrefix(ev.Text, bot) {
						input := strings.Replace(html.UnescapeString(ev.Text), bot+" ", "", 1)
						message := models.MakeMessage(service.Type, service.Name, channel, users[ev.User], input)
						inputMsgs <- message
					}

				}
			}
		}
	}
}
<commit_msg>Adding a switch to talk to hex via @ or not with 'UseAt'<commit_after>package inputs

import (
	"html"
	"log"
	"strings"

	"github.com/hexbotio/hex/models"
	"github.com/nlopes/slack"
)

// Slack struct
type Slack struct {
}

// Read function
func (x Slack) Read(inputMsgs chan<- models.Message, service models.Service) {
	defer Recovery(service)
	var useAt = true
	if strings.ToLower(service.Config["UseAt"]) == "false" {
		useAt = false
	}
	api := slack.New(service.Config["Key"])
	var slackDebug = false
	if strings.ToLower(service.Config["SlackDebug"]) == "true" {
		slackDebug = true
	}
	api.SetDebug(slackDebug)

	// get channel array
	channels := make(map[string]string)
	channelList, err := api.GetChannels(true)
	if err != nil {
		log.Printf("%s\n", err)
	}
	for _, channel := range channelList {
		channels[channel.ID] = "#" + channel.Name
	}

	// get user array
	bot := ""
	users := make(map[string]string)
	userList, err := api.GetUsers()
	if err != nil {
		log.Printf("%s\n", err)
	}
	for _, user := range userList {
		users[user.ID] = user.Name
		if service.BotName == user.Name {
			bot = user.ID
		}
	}

	// validate bot
	if bot == "" {
		log.Print("WARNING - Bot does not seem to be configured in Slack as '" + service.BotName + "'")
	} else {
		bot = "<@" + bot + ">"
	}

	// listen to messages
	rtm := api.NewRTM()
	go rtm.ManageConnection()
	for {
		select {
		case msg := <-rtm.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.MessageEvent:
				if ev.User != "" {

					channel, found := channels[ev.Channel]
					if !found {
						channel = ev.Channel
					}

					if useAt && strings.HasPrefix(ev.Text, bot) {
						input := strings.Replace(html.UnescapeString(ev.Text), bot+" ", "", 1)
						message := models.MakeMessage(service.Type, service.Name, channel, users[ev.User], input)
						inputMsgs <- message
					}
					if !useAt && strings.HasPrefix(ev.Text, service.BotName) {
						input := strings.Replace(html.UnescapeString(ev.Text), service.BotName+" ", "", 1)
						message := models.MakeMessage(service.Type, service.Name, channel, users[ev.User], input)
						inputMsgs <- message
					}
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gogap/spirit"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strings"
	"sync"
	"syscall"
	"text/template"
	"time"
)

var (
	ErrNoURNPackageSourceFound = errors.New("no urn packages source found")
	ErrConfigFileNameIsEmpty   = errors.New("config file name is empty")
)

type SpiritHelper struct {
	conf           spirit.SpiritConfig
	configFile     string
	configFileName string
	originalConfig []byte

	RefURNs     []string
	RefPackages []Package
}

func (p *SpiritHelper) LoadSpiritConfig(filename string) (err error) {

	if filename == "" {
		err = ErrConfigFileNameIsEmpty
		return
	}

	if fi, e := os.Stat(filename); e != nil {
		err = e
		return
	} else {
		p.configFile = filename
		p.configFileName = fi.Name()
	}

	if p.originalConfig, err = ioutil.ReadFile(filename); err != nil {
		return
	}

	if err = json.Unmarshal(p.originalConfig, &p.conf); err != nil {
		return
	}

	return
}

func (p *SpiritHelper) CreateProject(createOpts CreateOptions, tmplArgs map[string]interface{}) (err error) {
	if err = createOpts.Validate(); err != nil {
		return
	}

	goSrc := path.Join(createOpts.GoPath, "src")

	if err = p.parse(goSrc, createOpts.Sources); err != nil {
		return
	}

	// download packages
	if createOpts.GetPackages {
		if err = p.GetPackages(createOpts.PackagesRevision, createOpts.UpdatePackages); err != nil {
			return
		}
	}

	// make project dir
	projectPath := path.Join(goSrc, createOpts.ProjectPath)
	if path.IsAbs(createOpts.ProjectPath) {
		projectPath = createOpts.ProjectPath
	}

	if fi, e := os.Stat(projectPath); e != nil {
		if !strings.Contains(e.Error(), "no such file or directory") &&
			!os.IsNotExist(e) {
			err = e
			return
		}
	} else if !fi.IsDir() {
		err = fmt.Errorf("your project path %s already exist, but it is not a directory", projectPath)
		return
	} else if createOpts.ForceWrite {
		spirit.Logger().Warnf("project path %s already exist, it will be overwrite", projectPath)
	} else {
		err = fmt.Errorf("your project path %s already exist", projectPath)
		return
	}

	if err = os.MkdirAll(projectPath, os.FileMode(0755)); err != nil {
		if !os.IsNotExist(err) {
			return
		} else if !createOpts.ForceWrite {
			return
		}
		err = nil
	}

	// render code template
	tmplPathFmt := "github.com/gogap/spirit-tool/template/%s/main.go"
	tmplArgsPathFmt := "github.com/gogap/spirit-tool/template/%s/args.json"

	tmplPath := path.Join(goSrc, fmt.Sprintf(tmplPathFmt, createOpts.TemplateName))
	spirit.Logger().Infof("using template of %s: %s", createOpts.TemplateName, tmplPath)

	tmplArgsPath := path.Join(goSrc, fmt.Sprintf(tmplArgsPathFmt, createOpts.TemplateName))
	spirit.Logger().Infof("using template args of %s: %s", createOpts.TemplateName, tmplArgsPath)

	var tmpl *template.Template
	if tmpl, err = template.New("main.go").Option("missingkey=error").Delims("//<-", "->//").ParseFiles(tmplPath); err != nil {
		return
	}

	internalArgs := map[string]interface{}{}
	if argData, e := ioutil.ReadFile(tmplArgsPath); e == nil {
		if err = json.Unmarshal(argData, &internalArgs); err != nil {
			return
		}
	}

	if tmplArgs != nil {
		for k, v := range tmplArgs {
			internalArgs[k] = v
		}
	}

	buffer := &bytes.Buffer{}
	if err = tmpl.Execute(buffer, map[string]interface{}{
		"create_options":  createOpts,
		"packages":        p.RefPackages,
		"config":          p.configFile,
		"config_filename": p.configFileName,
		"create_time":     time.Now(),
		"args":            internalArgs}); err != nil {
		return
	}

	srcPath := path.Join(projectPath, "main.go")
	if err = ioutil.WriteFile(srcPath, buffer.Bytes(), os.FileMode(0644)); err != nil {
		return
	}

	confPath := path.Join(projectPath, p.configFileName)
	if err = ioutil.WriteFile(confPath, p.originalConfig, os.FileMode(0644)); err != nil {
		return
	}

	// format code for sort import packages order
	if _, err = execCommand("go fmt " + srcPath); err != nil {
		return
	}

	spirit.Logger().Infof("project created at %s\n", projectPath)

	return
}

func (p *SpiritHelper) GetPackages(pkgRevision map[string]string, update bool) (err error) {
	for _, pkg := range p.RefPackages {
		if pkgRevision != nil {
			if revision, exist := pkgRevision[pkg.URI]; exist {
				pkg.Revision = revision
			}
		}
		if err = pkg.Get(update); err != nil {
			return
		}
	}
	return
}

func (p *SpiritHelper) RunProject(createOpts CreateOptions, tmplArgs map[string]interface{}) (err error) {

	if err = p.CreateProject(createOpts, tmplArgs); err != nil {
		return
	}

	if _, err = execCommandWithDir("go build -o main "+path.Join(createOpts.ProjectPath, "main.go"), createOpts.ProjectPath); err != nil {
		return
	}

	if cmder, e := execute(path.Join(createOpts.ProjectPath, "main"), createOpts.ProjectPath); e != nil {
		err = e
		return
	} else {

		wg := sync.WaitGroup{}

		wg.Add(1)
		waitSignal(cmder, &wg)

		wg.Wait()
	}

	return
}

func (p *SpiritHelper) parse(gosrc string, sources []string) (err error) {
	if sources == nil || len(sources) == 0 {
		err = ErrNoURNPackageSourceFound
		return
	}

	var urns []string

	if urns = parseActorsUsingURN(
		p.conf.InputTranslators,
		p.conf.OutputTranslators,
		p.conf.Inboxes,
		p.conf.Outboxes,
		p.conf.Receivers,
		p.conf.Senders,
		p.conf.Routers,
		p.conf.Components,
		p.conf.LabelMatchers,
		p.conf.URNRewriters,
	); err != nil {
		return
	}

	for _, readerPool := range p.conf.ReaderPools {
		urns = append(urns, parseActorUsingURN(readerPool.ActorConfig)...)
		if readerPool.Reader != nil {
			urns = append(urns, parseActorUsingURN(*readerPool.Reader)...)
		}
	}

	for _, writerPool := range p.conf.WriterPools {
		urns = append(urns, parseActorUsingURN(writerPool.ActorConfig)...)
		if writerPool.Writer != nil {
			urns = append(urns, parseActorUsingURN(*writerPool.Writer)...)
		}
	}

	p.RefURNs = urns

	if p.RefPackages, err = urnsToPackages(gosrc, urns, sources...); err != nil {
		return
	}

	return
}

func parseActorsUsingURN(confs ...[]spirit.ActorConfig) (urns []string) {
	for _, conf := range confs {
		for _, c := range conf {
			urns = append(urns, c.URN)
		}
	}
	return
}

func parseActorUsingURN(actorConfs ...spirit.ActorConfig) (urns []string) {
	for _, conf := range actorConfs {
		urns = append(urns, conf.URN)
	}
	return
}

func urnsToPackages(gosrc string, urns []string, sourceFiles ...string) (packages []Package, err error) {
	urnPkgMap := map[string]string{}

	for _, sourceFile := range sourceFiles {
		var data []byte

		if data, err = ioutil.ReadFile(sourceFile); err != nil {
			return
		}

		sourceConf := SourceConfig{}
		if err = json.Unmarshal(data, &sourceConf); err != nil {
			return
		}

		for _, urnPkg := range sourceConf.Packages {
			if oldVal, exist := urnPkgMap[urnPkg.URN]; exist {
				if oldVal != urnPkg.Pkg {
					err = fmt.Errorf("source have duplicate urn pkg, urn:%s, pkg1:%s, pkg2: %s, file: %s", urnPkg.URN, oldVal, urnPkg.Pkg, sourceFile)
					return
				}
			}
			urnPkgMap[urnPkg.URN] = urnPkg.Pkg
		}
	}

	pkgs := map[string]bool{}

	for _, urn := range urns {
		if pkg, exist := urnPkgMap[urn]; !exist {
			err = fmt.Errorf("urn of %s not exist", urn)
		} else {
			pkgs[pkg] = true
		}
	}

	for pkg, _ := range pkgs {
		packages = append(packages, Package{gosrc: gosrc, URI: pkg, Revision: ""})
	}

	return
}

func waitSignal(cmd *exec.Cmd, wg *sync.WaitGroup) {
	defer wg.Done()

	isStopping := false
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, syscall.SIGTERM)

	pid := cmd.Process.Pid

	for {
		select {
		case signal := <-interrupt:
			{
				switch signal {
				case os.Interrupt, syscall.SIGTERM:
					{
						if isStopping {
							killProcess(pid)
							spirit.Logger().Infof("kill process, pid: %d\n", pid)
							return
						}

						isStopping = true
						spirit.Logger().Infof("stop process, pid: %d\n", pid)

						cmd.Wait()

						return
					}
				}
			}
		case <-time.After(time.Second):
			{
				continue
			}
		}
	}
}
<commit_msg>use cmd.Wait() to wait child process shutdown<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gogap/spirit"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"text/template"
	"time"
)

var (
	ErrNoURNPackageSourceFound = errors.New("no urn packages source found")
	ErrConfigFileNameIsEmpty   = errors.New("config file name is empty")
)

type SpiritHelper struct {
	conf           spirit.SpiritConfig
	configFile     string
	configFileName string
	originalConfig []byte

	RefURNs     []string
	RefPackages []Package
}

func (p *SpiritHelper) LoadSpiritConfig(filename string) (err error) {

	if filename == "" {
		err = ErrConfigFileNameIsEmpty
		return
	}

	if fi, e := os.Stat(filename); e != nil {
		err = e
		return
	} else {
		p.configFile = filename
		p.configFileName = fi.Name()
	}

	if p.originalConfig, err = ioutil.ReadFile(filename); err != nil {
		return
	}

	if err = json.Unmarshal(p.originalConfig, &p.conf); err != nil {
		return
	}

	return
}

func (p *SpiritHelper) CreateProject(createOpts CreateOptions, tmplArgs map[string]interface{}) (err error) {
	if err = createOpts.Validate(); err != nil {
		return
	}

	goSrc := path.Join(createOpts.GoPath, "src")

	if err = p.parse(goSrc, createOpts.Sources); err != nil {
		return
	}

	// download packages
	if createOpts.GetPackages {
		if err = p.GetPackages(createOpts.PackagesRevision, createOpts.UpdatePackages); err != nil {
			return
		}
	}

	// make project dir
	projectPath := path.Join(goSrc, createOpts.ProjectPath)
	if path.IsAbs(createOpts.ProjectPath) {
		projectPath = createOpts.ProjectPath
	}

	if fi, e := os.Stat(projectPath); e != nil {
		if !strings.Contains(e.Error(), "no such file or directory") &&
			!os.IsNotExist(e) {
			err = e
			return
		}
	} else if !fi.IsDir() {
		err = fmt.Errorf("your project path %s already exist, but it is not a directory", projectPath)
		return
	} else if createOpts.ForceWrite {
		spirit.Logger().Warnf("project path %s already exist, it will be overwrite", projectPath)
	} else {
		err = fmt.Errorf("your project path %s already exist", projectPath)
		return
	}

	if err = os.MkdirAll(projectPath, os.FileMode(0755)); err != nil {
		if !os.IsNotExist(err) {
			return
		} else if !createOpts.ForceWrite {
			return
		}
		err = nil
	}

	// render code template
	tmplPathFmt := "github.com/gogap/spirit-tool/template/%s/main.go"
	tmplArgsPathFmt := "github.com/gogap/spirit-tool/template/%s/args.json"

	tmplPath := path.Join(goSrc, fmt.Sprintf(tmplPathFmt, createOpts.TemplateName))
	spirit.Logger().Infof("using template of %s: %s", createOpts.TemplateName, tmplPath)

	tmplArgsPath := path.Join(goSrc, fmt.Sprintf(tmplArgsPathFmt, createOpts.TemplateName))
	spirit.Logger().Infof("using template args of %s: %s", createOpts.TemplateName, tmplArgsPath)

	var tmpl *template.Template
	if tmpl, err = template.New("main.go").Option("missingkey=error").Delims("//<-", "->//").ParseFiles(tmplPath); err != nil {
		return
	}

	internalArgs := map[string]interface{}{}
	if argData, e := ioutil.ReadFile(tmplArgsPath); e == nil {
		if err = json.Unmarshal(argData, &internalArgs); err != nil {
			return
		}
	}

	if tmplArgs != nil {
		for k, v := range tmplArgs {
			internalArgs[k] = v
		}
	}

	buffer := &bytes.Buffer{}
	if err = tmpl.Execute(buffer, map[string]interface{}{
		"create_options":  createOpts,
		"packages":        p.RefPackages,
		"config":          p.configFile,
		"config_filename": p.configFileName,
		"create_time":     time.Now(),
		"args":            internalArgs}); err != nil {
		return
	}

	srcPath := path.Join(projectPath, "main.go")
	if err = ioutil.WriteFile(srcPath, buffer.Bytes(), os.FileMode(0644)); err != nil {
		return
	}

	confPath := path.Join(projectPath, p.configFileName)
	if err = ioutil.WriteFile(confPath, p.originalConfig, os.FileMode(0644)); err != nil {
		return
	}

	// format code for sort import packages order
	if _, err = execCommand("go fmt " + srcPath); err != nil {
		return
	}

	spirit.Logger().Infof("project created at %s\n", projectPath)

	return
}

func (p *SpiritHelper) GetPackages(pkgRevision map[string]string, update bool) (err error) {
	for _, pkg := range p.RefPackages {
		if pkgRevision != nil {
			if revision, exist := pkgRevision[pkg.URI]; exist {
				pkg.Revision = revision
			}
		}
		if err = pkg.Get(update); err != nil {
			return
		}
	}
	return
}

func (p *SpiritHelper) RunProject(createOpts CreateOptions, tmplArgs map[string]interface{}) (err error) {

	if err = p.CreateProject(createOpts, tmplArgs); err != nil {
		return
	}

	if _, err = execCommandWithDir("go build -o main "+path.Join(createOpts.ProjectPath, "main.go"), createOpts.ProjectPath); err != nil {
		return
	}

	if cmder, e := execute(path.Join(createOpts.ProjectPath, "main"), createOpts.ProjectPath); e != nil {
		err = e
		return
	} else {
		cmder.Wait()
	}

	return
}

func (p *SpiritHelper) parse(gosrc string, sources []string) (err error) {
	if sources == nil || len(sources) == 0 {
		err = ErrNoURNPackageSourceFound
		return
	}

	var urns []string

	if urns = parseActorsUsingURN(
		p.conf.InputTranslators,
		p.conf.OutputTranslators,
		p.conf.Inboxes,
		p.conf.Outboxes,
		p.conf.Receivers,
		p.conf.Senders,
		p.conf.Routers,
		p.conf.Components,
		p.conf.LabelMatchers,
		p.conf.URNRewriters,
	); err != nil {
		return
	}

	for _, readerPool := range p.conf.ReaderPools {
		urns = append(urns, parseActorUsingURN(readerPool.ActorConfig)...)
		if readerPool.Reader != nil {
			urns = append(urns, parseActorUsingURN(*readerPool.Reader)...)
		}
	}

	for _, writerPool := range p.conf.WriterPools {
		urns = append(urns, parseActorUsingURN(writerPool.ActorConfig)...)
		if writerPool.Writer != nil {
			urns = append(urns, parseActorUsingURN(*writerPool.Writer)...)
		}
	}

	p.RefURNs = urns

	if p.RefPackages, err = urnsToPackages(gosrc, urns, sources...); err != nil {
		return
	}

	return
}

func parseActorsUsingURN(confs ...[]spirit.ActorConfig) (urns []string) {
	for _, conf := range confs {
		for _, c := range conf {
			urns = append(urns, c.URN)
		}
	}
	return
}

func parseActorUsingURN(actorConfs ...spirit.ActorConfig) (urns []string) {
	for _, conf := range actorConfs {
		urns = append(urns, conf.URN)
	}
	return
}

func urnsToPackages(gosrc string, urns []string, sourceFiles ...string) (packages []Package, err error) {
	urnPkgMap := map[string]string{}

	for _, sourceFile := range sourceFiles {
		var data []byte

		if data, err = ioutil.ReadFile(sourceFile); err != nil {
			return
		}

		sourceConf := SourceConfig{}
		if err = json.Unmarshal(data, &sourceConf); err != nil {
			return
		}

		for _, urnPkg := range sourceConf.Packages {
			if oldVal, exist := urnPkgMap[urnPkg.URN]; exist {
				if oldVal != urnPkg.Pkg {
					err = fmt.Errorf("source have duplicate urn pkg, urn:%s, pkg1:%s, pkg2: %s, file: %s", urnPkg.URN, oldVal, urnPkg.Pkg, sourceFile)
					return
				}
			}
			urnPkgMap[urnPkg.URN] = urnPkg.Pkg
		}
	}

	pkgs := map[string]bool{}

	for _, urn := range urns {
		if pkg, exist := urnPkgMap[urn]; !exist {
			err = fmt.Errorf("urn of %s not exist", urn)
		} else {
			pkgs[pkg] = true
		}
	}

	for pkg, _ := range pkgs {
		packages = append(packages, Package{gosrc: gosrc, URI: pkg, Revision: ""})
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach 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.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package sql_test

import (
	"bytes"
	"database/sql"
	"fmt"
	"os"
	"sync"
	"testing"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/security/securitytest"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/testutils/sqlutils"
	"github.com/cockroachdb/cockroach/testutils/storageutils"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/randutil"
)

//go:generate ../util/leaktest/add-leaktest.sh *_test.go

func init() {
	security.SetReadFileFn(securitytest.Asset)
}

// CommandFilters provides facilities for registering "TestingCommandFilters"
// (i.e. functions to be run on every replica command).
// CommandFilters is thread-safe.
// CommandFilters also optionally does replay protection if filters need it.
type CommandFilters struct {
	sync.RWMutex
	filters []struct {
		id         int
		idempotent bool
		filter     storageutils.ReplicaCommandFilter
	}
	nextID int

	numFiltersTrackingReplays int
	replayProtection          storageutils.ReplicaCommandFilter
}

// runFilters executes the registered filters, stopping at the first one
// that returns an error.
func (c *CommandFilters) runFilters(args storageutils.FilterArgs) *roachpb.Error {

	c.RLock()
	defer c.RUnlock()

	if c.replayProtection != nil {
		return c.replayProtection(args)
	} else {
		return c.runFiltersInternal(args)
	}
}

func (c *CommandFilters) runFiltersInternal(args storageutils.FilterArgs) *roachpb.Error {
	for _, f := range c.filters {
		if pErr := f.filter(args); pErr != nil {
			return pErr
		}
	}
	return nil
}

// AppendFilter registers a filter function to run after all the previously
// registered filters.
// idempotent specifies if this filter can be safely run multiple times on the
// same command. If this property doesn't hold, CommandFilters will start
// tracking commands for replay protection, which might be expensive.
// Returns a closure that the client must run for doing cleanup when the
// filter should be deregistered.
func (c *CommandFilters) AppendFilter(
	filter storageutils.ReplicaCommandFilter, idempotent bool) func() {

	c.Lock()
	defer c.Unlock()
	id := c.nextID
	c.nextID++
	c.filters = append(c.filters, struct {
		id         int
		idempotent bool
		filter     storageutils.ReplicaCommandFilter
	}{id, idempotent, filter})

	if !idempotent {
		if c.numFiltersTrackingReplays == 0 {
			c.replayProtection =
				storageutils.WrapFilterForReplayProtection(c.runFiltersInternal)
		}
		c.numFiltersTrackingReplays++
	}

	return func() {
		c.removeFilter(id)
	}
}

// removeFilter removes a filter previously registered. Meant to be used as the
// closure returned by AppendFilter.
func (c *CommandFilters) removeFilter(id int) {
	c.Lock()
	defer c.Unlock()
	for i, f := range c.filters {
		if f.id == id {
			if !f.idempotent {
				c.numFiltersTrackingReplays--
				if c.numFiltersTrackingReplays == 0 {
					c.replayProtection = nil
				}
			}
			c.filters = append(c.filters[:i], c.filters[i+1:]...)
			return
		}
	}
	panic(fmt.Sprintf("failed to find filter with id: %d.", id))
}

// checkEndTransactionTrigger verifies that an EndTransactionRequest
// that includes intents for the SystemDB keys sets the proper trigger.
func checkEndTransactionTrigger(args storageutils.FilterArgs) *roachpb.Error {
	req, ok := args.Req.(*roachpb.EndTransactionRequest)
	if !ok {
		return nil
	}

	if !req.Commit {
		// This is a rollback: skip trigger verification.
		return nil
	}

	modifiedSpanTrigger := req.InternalCommitTrigger.GetModifiedSpanTrigger()
	modifiedSystemConfigSpan := modifiedSpanTrigger != nil && modifiedSpanTrigger.SystemConfigSpan

	var hasSystemKey bool
	for _, span := range req.IntentSpans {
		keyAddr, err := keys.Addr(span.Key)
		if err != nil {
			return err
		}
		if bytes.Compare(keyAddr, keys.SystemConfigSpan.Key) >= 0 &&
			bytes.Compare(keyAddr, keys.SystemConfigSpan.EndKey) < 0 {
			hasSystemKey = true
			break
		}
	}
	// If the transaction in question has intents in the system span, then
	// modifiedSystemConfigSpan should always be true. However, it is possible
	// for modifiedSystemConfigSpan to be set, even though no system keys are
	// present. This can occur with certain conditional DDL statements (e.g.
	// "CREATE TABLE IF NOT EXISTS"), which set the SystemConfigTrigger
	// aggressively but may not actually end up changing the system DB depending
	// on the current state.
	// For more information, see the related comment at the beginning of
	// planner.makePlan().
	if hasSystemKey && !modifiedSystemConfigSpan {
		return roachpb.NewError(util.Errorf("EndTransaction hasSystemKey=%t, but hasSystemConfigTrigger=%t",
			hasSystemKey, modifiedSystemConfigSpan))
	}

	return nil
}

type testServer struct {
	server.TestServer
	cleanupFns []func()
}

func createTestServerContext() (*server.Context, *CommandFilters) {
	ctx := server.NewTestContext()
	var cmdFilters CommandFilters
	cmdFilters.AppendFilter(checkEndTransactionTrigger, true)
	ctx.TestingKnobs.StoreTestingKnobs.TestingCommandFilter = cmdFilters.runFilters
	return ctx, &cmdFilters
}

// The context used should probably come from createTestServerContext.
func setupTestServerWithContext(t *testing.T, ctx *server.Context) *testServer {
	s := &testServer{TestServer: server.TestServer{Ctx: ctx}}
	if err := s.Start(); err != nil {
		t.Fatal(err)
	}
	return s
}

func setup(t *testing.T) (*testServer, *sql.DB, *client.DB) {
	return setupWithContext(t, server.NewTestContext())
}

func setupWithContext(t *testing.T, ctx *server.Context) (*testServer, *sql.DB, *client.DB) {
	s := setupTestServerWithContext(t, ctx)

	// SQL requests use security.RootUser which has ALL permissions on everything.
	url, cleanupFn := sqlutils.PGUrl(t, &s.TestServer, security.RootUser, "setupWithContext")
	sqlDB, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}
	s.cleanupFns = append(s.cleanupFns, cleanupFn)

	return s, sqlDB, s.DB()
}

func cleanupTestServer(s *testServer) {
	s.Stop()
	for _, fn := range s.cleanupFns {
		fn()
	}
}

func cleanup(s *testServer, db *sql.DB) {
	_ = db.Close()
	cleanupTestServer(s)
}

func TestMain(m *testing.M) {
	randutil.SeedForTests()
	os.Exit(m.Run())
}
<commit_msg>fix compile error<commit_after>// Copyright 2015 The Cockroach 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.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package sql_test

import (
	"bytes"
	"database/sql"
	"fmt"
	"os"
	"sync"
	"testing"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/security/securitytest"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/testutils/sqlutils"
	"github.com/cockroachdb/cockroach/testutils/storageutils"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/randutil"
)

//go:generate ../util/leaktest/add-leaktest.sh *_test.go

func init() {
	security.SetReadFileFn(securitytest.Asset)
}

// CommandFilters provides facilities for registering "TestingCommandFilters"
// (i.e. functions to be run on every replica command).
// CommandFilters is thread-safe.
// CommandFilters also optionally does replay protection if filters need it.
type CommandFilters struct {
	sync.RWMutex
	filters []struct {
		id         int
		idempotent bool
		filter     storageutils.ReplicaCommandFilter
	}
	nextID int

	numFiltersTrackingReplays int
	replayProtection          storageutils.ReplicaCommandFilter
}

// runFilters executes the registered filters, stopping at the first one
// that returns an error.
func (c *CommandFilters) runFilters(args storageutils.FilterArgs) *roachpb.Error {

	c.RLock()
	defer c.RUnlock()

	if c.replayProtection != nil {
		return c.replayProtection(args)
	} else {
		return c.runFiltersInternal(args)
	}
}

func (c *CommandFilters) runFiltersInternal(args storageutils.FilterArgs) *roachpb.Error {
	for _, f := range c.filters {
		if pErr := f.filter(args); pErr != nil {
			return pErr
		}
	}
	return nil
}

// AppendFilter registers a filter function to run after all the previously
// registered filters.
// idempotent specifies if this filter can be safely run multiple times on the
// same command. If this property doesn't hold, CommandFilters will start
// tracking commands for replay protection, which might be expensive.
// Returns a closure that the client must run for doing cleanup when the
// filter should be deregistered.
func (c *CommandFilters) AppendFilter(
	filter storageutils.ReplicaCommandFilter, idempotent bool) func() {

	c.Lock()
	defer c.Unlock()
	id := c.nextID
	c.nextID++
	c.filters = append(c.filters, struct {
		id         int
		idempotent bool
		filter     storageutils.ReplicaCommandFilter
	}{id, idempotent, filter})

	if !idempotent {
		if c.numFiltersTrackingReplays == 0 {
			c.replayProtection =
				storageutils.WrapFilterForReplayProtection(c.runFiltersInternal)
		}
		c.numFiltersTrackingReplays++
	}

	return func() {
		c.removeFilter(id)
	}
}

// removeFilter removes a filter previously registered. Meant to be used as the
// closure returned by AppendFilter.
func (c *CommandFilters) removeFilter(id int) {
	c.Lock()
	defer c.Unlock()
	for i, f := range c.filters {
		if f.id == id {
			if !f.idempotent {
				c.numFiltersTrackingReplays--
				if c.numFiltersTrackingReplays == 0 {
					c.replayProtection = nil
				}
			}
			c.filters = append(c.filters[:i], c.filters[i+1:]...)
			return
		}
	}
	panic(fmt.Sprintf("failed to find filter with id: %d.", id))
}

// checkEndTransactionTrigger verifies that an EndTransactionRequest
// that includes intents for the SystemDB keys sets the proper trigger.
func checkEndTransactionTrigger(args storageutils.FilterArgs) *roachpb.Error {
	req, ok := args.Req.(*roachpb.EndTransactionRequest)
	if !ok {
		return nil
	}

	if !req.Commit {
		// This is a rollback: skip trigger verification.
		return nil
	}

	modifiedSpanTrigger := req.InternalCommitTrigger.GetModifiedSpanTrigger()
	modifiedSystemConfigSpan := modifiedSpanTrigger != nil && modifiedSpanTrigger.SystemConfigSpan

	var hasSystemKey bool
	for _, span := range req.IntentSpans {
		keyAddr, err := keys.Addr(span.Key)
		if err != nil {
			return roachpb.NewError(err)
		}
		if bytes.Compare(keyAddr, keys.SystemConfigSpan.Key) >= 0 &&
			bytes.Compare(keyAddr, keys.SystemConfigSpan.EndKey) < 0 {
			hasSystemKey = true
			break
		}
	}
	// If the transaction in question has intents in the system span, then
	// modifiedSystemConfigSpan should always be true. However, it is possible
	// for modifiedSystemConfigSpan to be set, even though no system keys are
	// present. This can occur with certain conditional DDL statements (e.g.
	// "CREATE TABLE IF NOT EXISTS"), which set the SystemConfigTrigger
	// aggressively but may not actually end up changing the system DB depending
	// on the current state.
	// For more information, see the related comment at the beginning of
	// planner.makePlan().
	if hasSystemKey && !modifiedSystemConfigSpan {
		return roachpb.NewError(util.Errorf("EndTransaction hasSystemKey=%t, but hasSystemConfigTrigger=%t",
			hasSystemKey, modifiedSystemConfigSpan))
	}

	return nil
}

type testServer struct {
	server.TestServer
	cleanupFns []func()
}

func createTestServerContext() (*server.Context, *CommandFilters) {
	ctx := server.NewTestContext()
	var cmdFilters CommandFilters
	cmdFilters.AppendFilter(checkEndTransactionTrigger, true)
	ctx.TestingKnobs.StoreTestingKnobs.TestingCommandFilter = cmdFilters.runFilters
	return ctx, &cmdFilters
}

// The context used should probably come from createTestServerContext.
func setupTestServerWithContext(t *testing.T, ctx *server.Context) *testServer {
	s := &testServer{TestServer: server.TestServer{Ctx: ctx}}
	if err := s.Start(); err != nil {
		t.Fatal(err)
	}
	return s
}

func setup(t *testing.T) (*testServer, *sql.DB, *client.DB) {
	return setupWithContext(t, server.NewTestContext())
}

func setupWithContext(t *testing.T, ctx *server.Context) (*testServer, *sql.DB, *client.DB) {
	s := setupTestServerWithContext(t, ctx)

	// SQL requests use security.RootUser which has ALL permissions on everything.
	url, cleanupFn := sqlutils.PGUrl(t, &s.TestServer, security.RootUser, "setupWithContext")
	sqlDB, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}
	s.cleanupFns = append(s.cleanupFns, cleanupFn)

	return s, sqlDB, s.DB()
}

func cleanupTestServer(s *testServer) {
	s.Stop()
	for _, fn := range s.cleanupFns {
		fn()
	}
}

func cleanup(s *testServer, db *sql.DB) {
	_ = db.Close()
	cleanupTestServer(s)
}

func TestMain(m *testing.M) {
	randutil.SeedForTests()
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>package octorrent

import (
	"bufio"
	"bytes"
	"math"
	"reflect"
	"strconv"
	"testing"
)

func getDecoder(d []byte) decoder {
	buf := bytes.NewBuffer(d)

	return decoder{
		*bufio.NewReader(buf),
	}
}

func TestReadInt(t *testing.T) {
	i, o := "412e51", int64(412)
	dec := getDecoder([]byte(i))

	num, err := dec.readInt()
	if err != nil {
		t.Errorf("readInt returned error: %v", err)
	}

	if num != o {
		t.Errorf("readInt did return %i, but should return %i", num, o)
	}
}

func TestReadIntUntil(t *testing.T) {
	i, o := "142", int64(14)
	dec := getDecoder([]byte(i))

	num, err := dec.readIntUntil('2')
	if err != nil {
		t.Errorf("readIntUntil returned error: %v", err)
	}

	if num != o {
		t.Errorf("readIntUntil did return %i, but should return %i", num, o)
	}

	// Check broken reader
	dec = getDecoder([]byte(""))
	num, err = dec.readIntUntil('t')
	if err == nil {
		t.Errorf("readIntUntil with empty reader returned no error, but should")
	}

	dec = getDecoder([]byte("test"))
	num, err = dec.readIntUntil('e')
	if err == nil {
		t.Errorf("readIntUntil with no digits returned no error, but should")
	}

	// Check unsigned int
	var ui uint64
	ui = math.MaxInt64 + 1
	i = "i" + strconv.FormatUint(ui, 10) + "e"
	dec = getDecoder([]byte(i))
	num, err = dec.readIntUntil('e')
	if err == nil {
		t.Errorf("readIntUntil with unsigned int returned no error, but should")
	}
}

func TestReadString(t *testing.T) {
	const i, o = "4:testfoobar", "test"
	dec := getDecoder([]byte(i))

	str, err := dec.readString()
	if err != nil {
		t.Errorf("readString returned error: %v", err)
	}

	if str != o {
		t.Errorf("readString returned %s, but should return %s", str, o)
	}
}

func TestReadList(t *testing.T) {
	const i = "4:testi4ee"
	o := []interface{}{"test", int64(4)}
	dec := getDecoder([]byte(i))

	list, err := dec.readList()
	if err != nil {
		t.Errorf("readList returned error: %v", err)
	}

	if !reflect.DeepEqual(list, o) {
		t.Errorf("readList returned %v, but should return %v", list, o)
	}
}

func TestIsEnd(t *testing.T) {
	i, o := "4:test", false
	dec := getDecoder([]byte(i))
	end, err := dec.isEnd()
	if err != nil {
		t.Errorf("isEnd returned error: %v", err)
	}
	if end != o {
		t.Errorf("isEnd returned %v, but should return %v", end, o)
	}

	i, o = "e4:test", true
	dec = getDecoder([]byte(i))
	end, err = dec.isEnd()
	if err != nil {
		t.Errorf("isEnd returned error: %v", err)
	}
	if end != o {
		t.Errorf("isEnd returned %v, but should return %v", end, o)
	}
}

func TestReadDictionary(t *testing.T) {
	const i = "4:testi1337e3:foo3:bare"
	o := map[string]interface{}{
		"test": int64(1337),
		"foo":  "bar",
	}
	dec := getDecoder([]byte(i))

	dict, err := dec.readDictionary()
	if err != nil {
		t.Errorf("readDictionary returned error: %v", err)
	}

	if !reflect.DeepEqual(dict, o) {
		t.Errorf("readDictionary returned %v, but should %v", dict, o)
	}
}

func TestDecode(t *testing.T) {
	const i = "d4:testli4ei3ee3:foo3:bare"
	o := map[string]interface{}{
		"test": []interface{}{
			int64(4),
			int64(3),
		},
		"foo": "bar",
	}
	buf := bytes.NewBufferString(i)

	res, err := Decode(buf)
	if err != nil {
		t.Errorf("Decode returned error: %v", err)
	}

	if !reflect.DeepEqual(res, o) {
		t.Errorf("Decode returned %v, but should return %v", res, o)
	}
}
<commit_msg>ADd test for readString<commit_after>package octorrent

import (
	"bufio"
	"bytes"
	"math"
	"reflect"
	"strconv"
	"testing"
)

func getDecoder(d []byte) decoder {
	buf := bytes.NewBuffer(d)

	return decoder{
		*bufio.NewReader(buf),
	}
}

func TestReadInt(t *testing.T) {
	i, o := "412e51", int64(412)
	dec := getDecoder([]byte(i))

	num, err := dec.readInt()
	if err != nil {
		t.Errorf("readInt returned error: %v", err)
	}

	if num != o {
		t.Errorf("readInt did return %i, but should return %i", num, o)
	}
}

func TestReadIntUntil(t *testing.T) {
	i, o := "142", int64(14)
	dec := getDecoder([]byte(i))

	num, err := dec.readIntUntil('2')
	if err != nil {
		t.Errorf("readIntUntil returned error: %v", err)
	}

	if num != o {
		t.Errorf("readIntUntil did return %i, but should return %i", num, o)
	}

	// Check broken reader
	dec = getDecoder([]byte(""))
	num, err = dec.readIntUntil('t')
	if err == nil {
		t.Errorf("readIntUntil with empty reader returned no error, but should")
	}

	dec = getDecoder([]byte("test"))
	num, err = dec.readIntUntil('e')
	if err == nil {
		t.Errorf("readIntUntil with no digits returned no error, but should")
	}

	// Check unsigned int
	var ui uint64
	ui = math.MaxInt64 + 1
	i = "i" + strconv.FormatUint(ui, 10) + "e"
	dec = getDecoder([]byte(i))
	num, err = dec.readIntUntil('e')
	if err == nil {
		t.Errorf("readIntUntil with unsigned int returned no error, but should")
	}
}

func TestReadString(t *testing.T) {
	i, o := "4:testfoobar", "test"
	dec := getDecoder([]byte(i))
	if str, err := dec.readString(); err != nil {
		t.Errorf("readString returned error: %v", err)
	} else if str != o {
		t.Errorf("readString returned %s, but should return %s", str, o)
	}

	// Test broken reader
	dec = getDecoder([]byte(""))
	if _, err := dec.readString(); err == nil {
		t.Errorf("readSstring with empty string returned no error, but should")
	}

	// Test wrong format
	dec = getDecoder([]byte("4:foo"))
	if _, err := dec.readString(); err == nil {
		t.Errorf("readString with wrong formatted string returned no error, but should")
	}
}

func TestReadList(t *testing.T) {
	const i = "4:testi4ee"
	o := []interface{}{"test", int64(4)}
	dec := getDecoder([]byte(i))

	list, err := dec.readList()
	if err != nil {
		t.Errorf("readList returned error: %v", err)
	}

	if !reflect.DeepEqual(list, o) {
		t.Errorf("readList returned %v, but should return %v", list, o)
	}
}

func TestIsEnd(t *testing.T) {
	i, o := "4:test", false
	dec := getDecoder([]byte(i))
	end, err := dec.isEnd()
	if err != nil {
		t.Errorf("isEnd returned error: %v", err)
	}
	if end != o {
		t.Errorf("isEnd returned %v, but should return %v", end, o)
	}

	i, o = "e4:test", true
	dec = getDecoder([]byte(i))
	end, err = dec.isEnd()
	if err != nil {
		t.Errorf("isEnd returned error: %v", err)
	}
	if end != o {
		t.Errorf("isEnd returned %v, but should return %v", end, o)
	}
}

func TestReadDictionary(t *testing.T) {
	const i = "4:testi1337e3:foo3:bare"
	o := map[string]interface{}{
		"test": int64(1337),
		"foo":  "bar",
	}
	dec := getDecoder([]byte(i))

	dict, err := dec.readDictionary()
	if err != nil {
		t.Errorf("readDictionary returned error: %v", err)
	}

	if !reflect.DeepEqual(dict, o) {
		t.Errorf("readDictionary returned %v, but should %v", dict, o)
	}
}

func TestDecode(t *testing.T) {
	const i = "d4:testli4ei3ee3:foo3:bare"
	o := map[string]interface{}{
		"test": []interface{}{
			int64(4),
			int64(3),
		},
		"foo": "bar",
	}
	buf := bytes.NewBufferString(i)

	res, err := Decode(buf)
	if err != nil {
		t.Errorf("Decode returned error: %v", err)
	}

	if !reflect.DeepEqual(res, o) {
		t.Errorf("Decode returned %v, but should return %v", res, o)
	}
}
<|endoftext|>


<commit_before>package raft

import (
	"fmt"
	"net"
	"time"
)

const (
	maxHeartbeatBackoff = 256
	maxFailureScale     = 8
	failureWait         = 10 * time.Millisecond
)

type followerReplication struct {
	peer     net.Addr
	inflight *inflight

	stopCh    chan uint64
	triggerCh chan struct{}

	currentTerm uint64
	matchIndex  uint64
	nextIndex   uint64

	lastContact time.Time
	failures    uint64
}

// replicate is a long running routine that is used to manage
// the process of replicating logs to our followers
func (r *Raft) replicate(s *followerReplication) {
	// Start an async heartbeating routing
	stopHeartbeat := make(chan struct{})
	defer close(stopHeartbeat)
	r.goFunc(func() { r.heartbeat(s, stopHeartbeat) })

	shouldStop := false
	for !shouldStop {
		select {
		case maxIndex := <-s.stopCh:
			// Make a best effort to replicate up to this index
			if maxIndex > 0 {
				r.replicateTo(s, maxIndex)
			}
			return
		case <-s.triggerCh:
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		case <-randomTimeout(r.conf.CommitTimeout):
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		}
	}
}

// replicateTo is used to replicate the logs up to a given last index.
// If the follower log is behind, we take care to bring them up to date
func (r *Raft) replicateTo(s *followerReplication, lastIndex uint64) (shouldStop bool) {
	// Create the base request
	var l Log
	var req AppendEntriesRequest
	var resp AppendEntriesResponse
	var maxIndex uint64
START:
	req = AppendEntriesRequest{
		Term:              s.currentTerm,
		Leader:            r.trans.EncodePeer(r.localAddr),
		LeaderCommitIndex: r.getCommitIndex(),
	}

	// Get the previous log entry based on the nextIndex.
	// Guard for the first index, since there is no 0 log entry
	// Guard against the previous index being a snapshot as well
	if s.nextIndex == 1 {
		req.PrevLogEntry = 0
		req.PrevLogTerm = 0

	} else if (s.nextIndex - 1) == r.getLastSnapshotIndex() {
		req.PrevLogEntry = r.getLastSnapshotIndex()
		req.PrevLogTerm = r.getLastSnapshotTerm()

	} else {
		if err := r.logs.GetLog(s.nextIndex-1, &l); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v",
				s.nextIndex-1, err)
			return
		}

		// Set the previous index and term (0 if nextIndex is 1)
		req.PrevLogEntry = l.Index
		req.PrevLogTerm = l.Term
	}

	// Append up to MaxAppendEntries or up to the lastIndex
	req.Entries = make([]*Log, 0, r.conf.MaxAppendEntries)
	maxIndex = min(s.nextIndex+uint64(r.conf.MaxAppendEntries)-1, lastIndex)
	for i := s.nextIndex; i <= maxIndex; i++ {
		oldLog := new(Log)
		if err := r.logs.GetLog(i, oldLog); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v", i, err)
			return
		}
		req.Entries = append(req.Entries, oldLog)
	}

	// Make the RPC call
	if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
		r.logger.Printf("[ERR] raft: Failed to AppendEntries to %v: %v", s.peer, err)
		s.failures++
		time.Sleep(backoff(failureWait, s.failures, maxFailureScale))
		return
	}
	s.failures = 0

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		return true
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Update the s based on success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.nextIndex, maxIndex, s.peer)

		// Update the indexes
		s.matchIndex = maxIndex
		s.nextIndex = maxIndex + 1
	} else {
		r.logger.Printf("[WARN] raft: AppendEntries to %v rejected, sending older logs", s.peer)
		s.nextIndex = max(min(s.nextIndex-1, resp.LastLog+1), 1)
		s.matchIndex = s.nextIndex - 1
	}

CHECK_MORE:
	// Check if there are more logs to replicate
	if s.nextIndex <= lastIndex {
		goto START
	}
	return

	// SEND_SNAP is used when we fail to get a log, usually because the follower
	// is too far behind, and we must ship a snapshot down instead
SEND_SNAP:
	stop, err := r.sendLatestSnapshot(s)

	// Check if we should stop
	if stop {
		return true
	}

	// Check for an error
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to send snapshot to %v: %v", s.peer, err)
		return
	}

	// Check if there is more to replicate
	goto CHECK_MORE
}

// sendLatestSnapshot is used to send the latest snapshot we have
// down to our follower
func (r *Raft) sendLatestSnapshot(s *followerReplication) (bool, error) {
	// Get the snapshots
	snapshots, err := r.snapshots.List()
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to list snapshots: %v", err)
		return false, err
	}

	// Check we have at least a single snapshot
	if len(snapshots) == 0 {
		return false, fmt.Errorf("no snapshots found")
	}

	// Open the most recent snapshot
	snapId := snapshots[0].ID
	meta, snapshot, err := r.snapshots.Open(snapId)
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to open snapshot %v: %v", snapId, err)
		return false, err
	}
	defer snapshot.Close()

	// Setup the request
	req := InstallSnapshotRequest{
		Term:         s.currentTerm,
		Leader:       r.trans.EncodePeer(r.localAddr),
		LastLogIndex: meta.Index,
		LastLogTerm:  meta.Term,
		Peers:        meta.Peers,
		Size:         meta.Size,
	}

	// Make the call
	var resp InstallSnapshotResponse
	if err := r.trans.InstallSnapshot(s.peer, &req, &resp, snapshot); err != nil {
		r.logger.Printf("[ERR] raft: Failed to install snapshot %v: %v", snapId, err)
		return false, err
	}

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		return true, nil
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Check for success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.matchIndex+1, meta.Index, s.peer)

		// Update the indexes
		s.matchIndex = meta.Index
		s.nextIndex = s.matchIndex + 1
	} else {
		r.logger.Printf("[WARN] raft: InstallSnapshot to %v rejected", s.peer)
	}
	return false, nil
}

// hearbeat is used to periodically invoke AppendEntries on a peer
// to ensure they don't time out. This is done async of replicate(),
// since that routine could potentially be blocked on disk IO
func (r *Raft) heartbeat(s *followerReplication, stopCh chan struct{}) {
	// Use an exponential backoff for retries
	backoff := 1

	req := AppendEntriesRequest{
		Term:   s.currentTerm,
		Leader: r.trans.EncodePeer(r.localAddr),
	}
	var resp AppendEntriesResponse
	for {
		select {
		case <-randomTimeout((r.conf.HeartbeatTimeout / 4) * time.Duration(backoff)):
			if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
				r.logger.Printf("[ERR] raft: Failed to heartbeat to %v: %v", s.peer, err)
				backoff *= 2
				if backoff > maxHeartbeatBackoff {
					backoff = maxHeartbeatBackoff
				}
			} else {
				s.lastContact = time.Now()
				backoff = 1
			}

		case <-stopCh:
			return
		}
	}
}
<commit_msg>Improve log message if follower has newer term<commit_after>package raft

import (
	"fmt"
	"net"
	"time"
)

const (
	maxHeartbeatBackoff = 256
	maxFailureScale     = 8
	failureWait         = 10 * time.Millisecond
)

type followerReplication struct {
	peer     net.Addr
	inflight *inflight

	stopCh    chan uint64
	triggerCh chan struct{}

	currentTerm uint64
	matchIndex  uint64
	nextIndex   uint64

	lastContact time.Time
	failures    uint64
}

// replicate is a long running routine that is used to manage
// the process of replicating logs to our followers
func (r *Raft) replicate(s *followerReplication) {
	// Start an async heartbeating routing
	stopHeartbeat := make(chan struct{})
	defer close(stopHeartbeat)
	r.goFunc(func() { r.heartbeat(s, stopHeartbeat) })

	shouldStop := false
	for !shouldStop {
		select {
		case maxIndex := <-s.stopCh:
			// Make a best effort to replicate up to this index
			if maxIndex > 0 {
				r.replicateTo(s, maxIndex)
			}
			return
		case <-s.triggerCh:
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		case <-randomTimeout(r.conf.CommitTimeout):
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		}
	}
}

// replicateTo is used to replicate the logs up to a given last index.
// If the follower log is behind, we take care to bring them up to date
func (r *Raft) replicateTo(s *followerReplication, lastIndex uint64) (shouldStop bool) {
	// Create the base request
	var l Log
	var req AppendEntriesRequest
	var resp AppendEntriesResponse
	var maxIndex uint64
START:
	req = AppendEntriesRequest{
		Term:              s.currentTerm,
		Leader:            r.trans.EncodePeer(r.localAddr),
		LeaderCommitIndex: r.getCommitIndex(),
	}

	// Get the previous log entry based on the nextIndex.
	// Guard for the first index, since there is no 0 log entry
	// Guard against the previous index being a snapshot as well
	if s.nextIndex == 1 {
		req.PrevLogEntry = 0
		req.PrevLogTerm = 0

	} else if (s.nextIndex - 1) == r.getLastSnapshotIndex() {
		req.PrevLogEntry = r.getLastSnapshotIndex()
		req.PrevLogTerm = r.getLastSnapshotTerm()

	} else {
		if err := r.logs.GetLog(s.nextIndex-1, &l); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v",
				s.nextIndex-1, err)
			return
		}

		// Set the previous index and term (0 if nextIndex is 1)
		req.PrevLogEntry = l.Index
		req.PrevLogTerm = l.Term
	}

	// Append up to MaxAppendEntries or up to the lastIndex
	req.Entries = make([]*Log, 0, r.conf.MaxAppendEntries)
	maxIndex = min(s.nextIndex+uint64(r.conf.MaxAppendEntries)-1, lastIndex)
	for i := s.nextIndex; i <= maxIndex; i++ {
		oldLog := new(Log)
		if err := r.logs.GetLog(i, oldLog); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v", i, err)
			return
		}
		req.Entries = append(req.Entries, oldLog)
	}

	// Make the RPC call
	if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
		r.logger.Printf("[ERR] raft: Failed to AppendEntries to %v: %v", s.peer, err)
		s.failures++
		time.Sleep(backoff(failureWait, s.failures, maxFailureScale))
		return
	}
	s.failures = 0

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Update the s based on success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.nextIndex, maxIndex, s.peer)

		// Update the indexes
		s.matchIndex = maxIndex
		s.nextIndex = maxIndex + 1
	} else {
		r.logger.Printf("[WARN] raft: AppendEntries to %v rejected, sending older logs", s.peer)
		s.nextIndex = max(min(s.nextIndex-1, resp.LastLog+1), 1)
		s.matchIndex = s.nextIndex - 1
	}

CHECK_MORE:
	// Check if there are more logs to replicate
	if s.nextIndex <= lastIndex {
		goto START
	}
	return

	// SEND_SNAP is used when we fail to get a log, usually because the follower
	// is too far behind, and we must ship a snapshot down instead
SEND_SNAP:
	stop, err := r.sendLatestSnapshot(s)

	// Check if we should stop
	if stop {
		return true
	}

	// Check for an error
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to send snapshot to %v: %v", s.peer, err)
		return
	}

	// Check if there is more to replicate
	goto CHECK_MORE
}

// sendLatestSnapshot is used to send the latest snapshot we have
// down to our follower
func (r *Raft) sendLatestSnapshot(s *followerReplication) (bool, error) {
	// Get the snapshots
	snapshots, err := r.snapshots.List()
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to list snapshots: %v", err)
		return false, err
	}

	// Check we have at least a single snapshot
	if len(snapshots) == 0 {
		return false, fmt.Errorf("no snapshots found")
	}

	// Open the most recent snapshot
	snapId := snapshots[0].ID
	meta, snapshot, err := r.snapshots.Open(snapId)
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to open snapshot %v: %v", snapId, err)
		return false, err
	}
	defer snapshot.Close()

	// Setup the request
	req := InstallSnapshotRequest{
		Term:         s.currentTerm,
		Leader:       r.trans.EncodePeer(r.localAddr),
		LastLogIndex: meta.Index,
		LastLogTerm:  meta.Term,
		Peers:        meta.Peers,
		Size:         meta.Size,
	}

	// Make the call
	var resp InstallSnapshotResponse
	if err := r.trans.InstallSnapshot(s.peer, &req, &resp, snapshot); err != nil {
		r.logger.Printf("[ERR] raft: Failed to install snapshot %v: %v", snapId, err)
		return false, err
	}

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true, nil
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Check for success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.matchIndex+1, meta.Index, s.peer)

		// Update the indexes
		s.matchIndex = meta.Index
		s.nextIndex = s.matchIndex + 1
	} else {
		r.logger.Printf("[WARN] raft: InstallSnapshot to %v rejected", s.peer)
	}
	return false, nil
}

// hearbeat is used to periodically invoke AppendEntries on a peer
// to ensure they don't time out. This is done async of replicate(),
// since that routine could potentially be blocked on disk IO
func (r *Raft) heartbeat(s *followerReplication, stopCh chan struct{}) {
	// Use an exponential backoff for retries
	backoff := 1

	req := AppendEntriesRequest{
		Term:   s.currentTerm,
		Leader: r.trans.EncodePeer(r.localAddr),
	}
	var resp AppendEntriesResponse
	for {
		select {
		case <-randomTimeout((r.conf.HeartbeatTimeout / 4) * time.Duration(backoff)):
			if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
				r.logger.Printf("[ERR] raft: Failed to heartbeat to %v: %v", s.peer, err)
				backoff *= 2
				if backoff > maxHeartbeatBackoff {
					backoff = maxHeartbeatBackoff
				}
			} else {
				s.lastContact = time.Now()
				backoff = 1
			}

		case <-stopCh:
			return
		}
	}
}
<|endoftext|>


<commit_before>package raft

import (
	"fmt"
	"github.com/armon/go-metrics"
	"net"
	"time"
)

const (
	maxFailureScale = 12
	failureWait     = 10 * time.Millisecond
)

type followerReplication struct {
	peer     net.Addr
	inflight *inflight

	stopCh    chan uint64
	triggerCh chan struct{}

	currentTerm     uint64
	matchIndex      uint64
	nextIndex       uint64
	lastCommitIndex uint64

	lastContact time.Time
	failures    uint64
}

// replicate is a long running routine that is used to manage
// the process of replicating logs to our followers
func (r *Raft) replicate(s *followerReplication) {
	// Start an async heartbeating routing
	stopHeartbeat := make(chan struct{})
	defer close(stopHeartbeat)
	r.goFunc(func() { r.heartbeat(s, stopHeartbeat) })

	shouldStop := false
	for !shouldStop {
		select {
		case maxIndex := <-s.stopCh:
			// Make a best effort to replicate up to this index
			if maxIndex > 0 {
				r.replicateTo(s, maxIndex)
			}
			return
		case <-s.triggerCh:
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		case <-randomTimeout(r.conf.CommitTimeout):
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		}
	}
}

// replicateTo is used to replicate the logs up to a given last index.
// If the follower log is behind, we take care to bring them up to date
func (r *Raft) replicateTo(s *followerReplication, lastIndex uint64) (shouldStop bool) {
	// Create the base request
	var l Log
	var req AppendEntriesRequest
	var resp AppendEntriesResponse
	var maxIndex uint64
	var start time.Time
START:
	req = AppendEntriesRequest{
		Term:              s.currentTerm,
		Leader:            r.trans.EncodePeer(r.localAddr),
		LeaderCommitIndex: r.getCommitIndex(),
	}

	// Get the previous log entry based on the nextIndex.
	// Guard for the first index, since there is no 0 log entry
	// Guard against the previous index being a snapshot as well
	if s.nextIndex == 1 {
		req.PrevLogEntry = 0
		req.PrevLogTerm = 0

	} else if (s.nextIndex - 1) == r.getLastSnapshotIndex() {
		req.PrevLogEntry = r.getLastSnapshotIndex()
		req.PrevLogTerm = r.getLastSnapshotTerm()

	} else {
		if err := r.logs.GetLog(s.nextIndex-1, &l); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v",
				s.nextIndex-1, err)
			return
		}

		// Set the previous index and term (0 if nextIndex is 1)
		req.PrevLogEntry = l.Index
		req.PrevLogTerm = l.Term
	}

	// Append up to MaxAppendEntries or up to the lastIndex
	req.Entries = make([]*Log, 0, r.conf.MaxAppendEntries)
	maxIndex = min(s.nextIndex+uint64(r.conf.MaxAppendEntries)-1, lastIndex)
	for i := s.nextIndex; i <= maxIndex; i++ {
		oldLog := new(Log)
		if err := r.logs.GetLog(i, oldLog); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v", i, err)
			return
		}
		req.Entries = append(req.Entries, oldLog)
	}

	// Skip the RPC call if there is nothing to do
	if len(req.Entries) == 0 && req.LeaderCommitIndex == s.lastCommitIndex {
		return
	}

	// Make the RPC call
	start = time.Now()
	if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
		r.logger.Printf("[ERR] raft: Failed to AppendEntries to %v: %v", s.peer, err)
		s.failures++
		select {
		case <-time.After(backoff(failureWait, s.failures, maxFailureScale)):
		case <-s.stopCh:
		}
		return
	}
	s.failures = 0
	metrics.MeasureSince([]string{"raft", "replication", "appendEntries", s.peer.String()}, start)

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Update the s based on success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.nextIndex, maxIndex, s.peer)

		// Update the indexes
		s.matchIndex = maxIndex
		s.nextIndex = maxIndex + 1
		s.lastCommitIndex = req.LeaderCommitIndex
	} else {
		r.logger.Printf("[WARN] raft: AppendEntries to %v rejected, sending older logs", s.peer)
		s.nextIndex = max(min(s.nextIndex-1, resp.LastLog+1), 1)
		s.matchIndex = s.nextIndex - 1
	}

CHECK_MORE:
	// Check if there are more logs to replicate
	if s.nextIndex <= lastIndex {
		goto START
	}
	return

	// SEND_SNAP is used when we fail to get a log, usually because the follower
	// is too far behind, and we must ship a snapshot down instead
SEND_SNAP:
	stop, err := r.sendLatestSnapshot(s)

	// Check if we should stop
	if stop {
		return true
	}

	// Check for an error
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to send snapshot to %v: %v", s.peer, err)
		return
	}

	// Check if there is more to replicate
	goto CHECK_MORE
}

// sendLatestSnapshot is used to send the latest snapshot we have
// down to our follower
func (r *Raft) sendLatestSnapshot(s *followerReplication) (bool, error) {
	// Get the snapshots
	snapshots, err := r.snapshots.List()
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to list snapshots: %v", err)
		return false, err
	}

	// Check we have at least a single snapshot
	if len(snapshots) == 0 {
		return false, fmt.Errorf("no snapshots found")
	}

	// Open the most recent snapshot
	snapId := snapshots[0].ID
	meta, snapshot, err := r.snapshots.Open(snapId)
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to open snapshot %v: %v", snapId, err)
		return false, err
	}
	defer snapshot.Close()

	// Setup the request
	req := InstallSnapshotRequest{
		Term:         s.currentTerm,
		Leader:       r.trans.EncodePeer(r.localAddr),
		LastLogIndex: meta.Index,
		LastLogTerm:  meta.Term,
		Peers:        meta.Peers,
		Size:         meta.Size,
	}

	// Make the call
	start := time.Now()
	var resp InstallSnapshotResponse
	if err := r.trans.InstallSnapshot(s.peer, &req, &resp, snapshot); err != nil {
		r.logger.Printf("[ERR] raft: Failed to install snapshot %v: %v", snapId, err)
		return false, err
	}
	metrics.MeasureSince([]string{"raft", "replication", "installSnapshot", s.peer.String()}, start)

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true, nil
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Check for success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.matchIndex+1, meta.Index, s.peer)

		// Update the indexes
		s.matchIndex = meta.Index
		s.nextIndex = s.matchIndex + 1
	} else {
		r.logger.Printf("[WARN] raft: InstallSnapshot to %v rejected", s.peer)
	}
	return false, nil
}

// hearbeat is used to periodically invoke AppendEntries on a peer
// to ensure they don't time out. This is done async of replicate(),
// since that routine could potentially be blocked on disk IO
func (r *Raft) heartbeat(s *followerReplication, stopCh chan struct{}) {
	var failures uint64
	req := AppendEntriesRequest{
		Term:   s.currentTerm,
		Leader: r.trans.EncodePeer(r.localAddr),
	}
	var resp AppendEntriesResponse
	for {
		select {
		case <-randomTimeout(r.conf.HeartbeatTimeout / 8):
			start := time.Now()
			if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
				r.logger.Printf("[ERR] raft: Failed to heartbeat to %v: %v", s.peer, err)
				failures++
				select {
				case <-time.After(backoff(failureWait, failures, maxFailureScale)):
				case <-stopCh:
				}
			} else {
				s.lastContact = time.Now()
				failures = 0
				metrics.MeasureSince([]string{"raft", "replication", "heartbeat", s.peer.String()}, start)
			}

		case <-stopCh:
			return
		}
	}
}
<commit_msg>Fix read from the wrong channel<commit_after>package raft

import (
	"fmt"
	"github.com/armon/go-metrics"
	"net"
	"time"
)

const (
	maxFailureScale = 12
	failureWait     = 10 * time.Millisecond
)

type followerReplication struct {
	peer     net.Addr
	inflight *inflight

	stopCh    chan uint64
	triggerCh chan struct{}

	currentTerm     uint64
	matchIndex      uint64
	nextIndex       uint64
	lastCommitIndex uint64

	lastContact time.Time
	failures    uint64
}

// replicate is a long running routine that is used to manage
// the process of replicating logs to our followers
func (r *Raft) replicate(s *followerReplication) {
	// Start an async heartbeating routing
	stopHeartbeat := make(chan struct{})
	defer close(stopHeartbeat)
	r.goFunc(func() { r.heartbeat(s, stopHeartbeat) })

	shouldStop := false
	for !shouldStop {
		select {
		case maxIndex := <-s.stopCh:
			// Make a best effort to replicate up to this index
			if maxIndex > 0 {
				r.replicateTo(s, maxIndex)
			}
			return
		case <-s.triggerCh:
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		case <-randomTimeout(r.conf.CommitTimeout):
			shouldStop = r.replicateTo(s, r.getLastLogIndex())
		}
	}
}

// replicateTo is used to replicate the logs up to a given last index.
// If the follower log is behind, we take care to bring them up to date
func (r *Raft) replicateTo(s *followerReplication, lastIndex uint64) (shouldStop bool) {
	// Create the base request
	var l Log
	var req AppendEntriesRequest
	var resp AppendEntriesResponse
	var maxIndex uint64
	var start time.Time
START:
	req = AppendEntriesRequest{
		Term:              s.currentTerm,
		Leader:            r.trans.EncodePeer(r.localAddr),
		LeaderCommitIndex: r.getCommitIndex(),
	}

	// Get the previous log entry based on the nextIndex.
	// Guard for the first index, since there is no 0 log entry
	// Guard against the previous index being a snapshot as well
	if s.nextIndex == 1 {
		req.PrevLogEntry = 0
		req.PrevLogTerm = 0

	} else if (s.nextIndex - 1) == r.getLastSnapshotIndex() {
		req.PrevLogEntry = r.getLastSnapshotIndex()
		req.PrevLogTerm = r.getLastSnapshotTerm()

	} else {
		if err := r.logs.GetLog(s.nextIndex-1, &l); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v",
				s.nextIndex-1, err)
			return
		}

		// Set the previous index and term (0 if nextIndex is 1)
		req.PrevLogEntry = l.Index
		req.PrevLogTerm = l.Term
	}

	// Append up to MaxAppendEntries or up to the lastIndex
	req.Entries = make([]*Log, 0, r.conf.MaxAppendEntries)
	maxIndex = min(s.nextIndex+uint64(r.conf.MaxAppendEntries)-1, lastIndex)
	for i := s.nextIndex; i <= maxIndex; i++ {
		oldLog := new(Log)
		if err := r.logs.GetLog(i, oldLog); err != nil {
			if err == LogNotFound {
				goto SEND_SNAP
			}
			r.logger.Printf("[ERR] raft: Failed to get log at index %d: %v", i, err)
			return
		}
		req.Entries = append(req.Entries, oldLog)
	}

	// Skip the RPC call if there is nothing to do
	if len(req.Entries) == 0 && req.LeaderCommitIndex == s.lastCommitIndex {
		return
	}

	// Make the RPC call
	start = time.Now()
	if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
		r.logger.Printf("[ERR] raft: Failed to AppendEntries to %v: %v", s.peer, err)
		s.failures++
		select {
		case <-time.After(backoff(failureWait, s.failures, maxFailureScale)):
		case <-r.shutdownCh:
		}
		return
	}
	s.failures = 0
	metrics.MeasureSince([]string{"raft", "replication", "appendEntries", s.peer.String()}, start)

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Update the s based on success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.nextIndex, maxIndex, s.peer)

		// Update the indexes
		s.matchIndex = maxIndex
		s.nextIndex = maxIndex + 1
		s.lastCommitIndex = req.LeaderCommitIndex
	} else {
		r.logger.Printf("[WARN] raft: AppendEntries to %v rejected, sending older logs", s.peer)
		s.nextIndex = max(min(s.nextIndex-1, resp.LastLog+1), 1)
		s.matchIndex = s.nextIndex - 1
	}

CHECK_MORE:
	// Check if there are more logs to replicate
	if s.nextIndex <= lastIndex {
		goto START
	}
	return

	// SEND_SNAP is used when we fail to get a log, usually because the follower
	// is too far behind, and we must ship a snapshot down instead
SEND_SNAP:
	stop, err := r.sendLatestSnapshot(s)

	// Check if we should stop
	if stop {
		return true
	}

	// Check for an error
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to send snapshot to %v: %v", s.peer, err)
		return
	}

	// Check if there is more to replicate
	goto CHECK_MORE
}

// sendLatestSnapshot is used to send the latest snapshot we have
// down to our follower
func (r *Raft) sendLatestSnapshot(s *followerReplication) (bool, error) {
	// Get the snapshots
	snapshots, err := r.snapshots.List()
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to list snapshots: %v", err)
		return false, err
	}

	// Check we have at least a single snapshot
	if len(snapshots) == 0 {
		return false, fmt.Errorf("no snapshots found")
	}

	// Open the most recent snapshot
	snapId := snapshots[0].ID
	meta, snapshot, err := r.snapshots.Open(snapId)
	if err != nil {
		r.logger.Printf("[ERR] raft: Failed to open snapshot %v: %v", snapId, err)
		return false, err
	}
	defer snapshot.Close()

	// Setup the request
	req := InstallSnapshotRequest{
		Term:         s.currentTerm,
		Leader:       r.trans.EncodePeer(r.localAddr),
		LastLogIndex: meta.Index,
		LastLogTerm:  meta.Term,
		Peers:        meta.Peers,
		Size:         meta.Size,
	}

	// Make the call
	start := time.Now()
	var resp InstallSnapshotResponse
	if err := r.trans.InstallSnapshot(s.peer, &req, &resp, snapshot); err != nil {
		r.logger.Printf("[ERR] raft: Failed to install snapshot %v: %v", snapId, err)
		return false, err
	}
	metrics.MeasureSince([]string{"raft", "replication", "installSnapshot", s.peer.String()}, start)

	// Check for a newer term, stop running
	if resp.Term > req.Term {
		r.logger.Printf("[ERR] raft: peer %v has newer term, stopping replication", s.peer)
		return true, nil
	}

	// Update the last contact
	s.lastContact = time.Now()

	// Check for success
	if resp.Success {
		// Mark any inflight logs as committed
		s.inflight.CommitRange(s.matchIndex+1, meta.Index, s.peer)

		// Update the indexes
		s.matchIndex = meta.Index
		s.nextIndex = s.matchIndex + 1
	} else {
		r.logger.Printf("[WARN] raft: InstallSnapshot to %v rejected", s.peer)
	}
	return false, nil
}

// hearbeat is used to periodically invoke AppendEntries on a peer
// to ensure they don't time out. This is done async of replicate(),
// since that routine could potentially be blocked on disk IO
func (r *Raft) heartbeat(s *followerReplication, stopCh chan struct{}) {
	var failures uint64
	req := AppendEntriesRequest{
		Term:   s.currentTerm,
		Leader: r.trans.EncodePeer(r.localAddr),
	}
	var resp AppendEntriesResponse
	for {
		select {
		case <-randomTimeout(r.conf.HeartbeatTimeout / 8):
			start := time.Now()
			if err := r.trans.AppendEntries(s.peer, &req, &resp); err != nil {
				r.logger.Printf("[ERR] raft: Failed to heartbeat to %v: %v", s.peer, err)
				failures++
				select {
				case <-time.After(backoff(failureWait, failures, maxFailureScale)):
				case <-stopCh:
				}
			} else {
				s.lastContact = time.Now()
				failures = 0
				metrics.MeasureSince([]string{"raft", "replication", "heartbeat", s.peer.String()}, start)
			}

		case <-stopCh:
			return
		}
	}
}
<|endoftext|>