<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	docker "github.com/smashwilson/go-dockerclient"
)

// OutputCollector is an io.Writer that accumulates output from a specified stream in an attached
// Docker container and appends it to the appropriate field within a SubmittedJob.
type OutputCollector struct {
	context  *Context
	job      *SubmittedJob
	isStdout bool
}

// DescribeStream returns "stdout" or "stderr" to indicate which stream this collector is consuming.
func (c OutputCollector) DescribeStream() string {
	if c.isStdout {
		return "stdout"
	}
	return "stderr"
}

// Write appends bytes to the selected stream and updates the SubmittedJob.
func (c OutputCollector) Write(p []byte) (int, error) {
	log.WithFields(log.Fields{
		"length": len(p),
		"bytes":  string(p),
		"stream": c.DescribeStream(),
	}).Debug("Received output from a job")

	if c.isStdout {
		c.job.Stdout += string(p)
	} else {
		c.job.Stderr += string(p)
	}

	if err := c.context.UpdateJob(c.job); err != nil {
		return 0, err
	}

	return len(p), nil
}

// Runner is the main entry point for the job runner goroutine.
func Runner(c *Context) {
	var client *docker.Client
	var err error

	if c.DockerTLS {
		client, err = docker.NewTLSClient(c.DockerHost, c.DockerCert, c.DockerKey, c.DockerCACert)
		if err != nil {
			log.WithFields(log.Fields{
				"docker host":    c.DockerHost,
				"docker cert":    c.DockerCert,
				"docker key":     c.DockerKey,
				"docker CA cert": c.DockerCACert,
			}).Fatal("Unable to connect to Docker with TLS.")
			return
		}
	} else {
		client, err = docker.NewClient(c.DockerHost)
		if err != nil {
			log.WithFields(log.Fields{
				"docker host": c.DockerHost,
				"error":       err,
			}).Fatal("Unable to connect to Docker.")
			return
		}
	}

	for {
		select {
		case <-time.After(time.Duration(c.Poll) * time.Millisecond):
			Claim(c, client)
		}
	}
}

// Claim acquires the oldest single pending job and launches a goroutine to execute its command in
// a new container.
func Claim(c *Context, client *docker.Client) {
	job, err := c.ClaimJob()
	if err != nil {
		log.WithFields(log.Fields{"error": err}).Error("Unable to claim a job.")
		return
	}
	if job == nil {
		// Nothing to claim.
		return
	}

	go Execute(c, client, job)
}

// Execute launches a container to process the submitted job. It passes any provided stdin data
// to the container and consumes stdout and stderr, updating Mongo as it runs. Once completed, it
// acquires the job's result from its configured source and marks the job as finished.
func Execute(c *Context, client *docker.Client, job *SubmittedJob) {
	defaultFields := log.Fields{
		"jid":     job.JID,
		"account": job.Account,
	}

	// Logging utility messages.
	debug := func(message string) {
		log.WithFields(defaultFields).Debug(message)
	}
	reportErr := func(message string, err error) {
		fs := log.Fields{}
		for k, v := range defaultFields {
			fs[k] = v
		}
		fs["err"] = err
		log.WithFields(fs).Error(message)
	}
	checkErr := func(message string, err error) bool {
		if err == nil {
			debug(fmt.Sprintf("%s: ok", message))
			return false
		}

		reportErr(fmt.Sprintf("%s: ERROR", message), err)
		return true
	}

	log.WithFields(defaultFields).Info("Launching a job.")

	job.StartedAt = StoreTime(time.Now())
	if err := c.UpdateJob(job); err != nil {
		reportErr("Unable to update the job's start timestamp.", err)
	}

	container, err := client.CreateContainer(docker.CreateContainerOptions{
		Name: job.ContainerName(),
		Config: &docker.Config{
			Image:     c.Image,
			Cmd:       []string{"/bin/bash", "-c", job.Command},
			OpenStdin: true,
			StdinOnce: true,
		},
	})
	if checkErr("Created the job's container", err) {
		return
	}

	// Include container information in this job's logging messages.
	defaultFields["container id"] = container.ID
	defaultFields["container name"] = container.Name

	// Prepare the input and output streams.
	stdin := bytes.NewReader(job.Stdin)
	stdout := OutputCollector{
		context:  c,
		job:      job,
		isStdout: true,
	}
	stderr := OutputCollector{
		context:  c,
		job:      job,
		isStdout: false,
	}

	go func() {
		err = client.AttachToContainer(docker.AttachToContainerOptions{
			Container:    container.ID,
			Stream:       true,
			InputStream:  stdin,
			OutputStream: stdout,
			ErrorStream:  stderr,
			Stdin:        true,
			Stdout:       true,
			Stderr:       true,
		})
		checkErr("Attached to the container", err)
	}()

	// Start the created container.
	err = client.StartContainer(container.ID, &docker.HostConfig{})
	if checkErr("Started the container", err) {
		return
	}

	status, err := client.WaitContainer(container.ID)
	if checkErr("Waited for the container to complete", err) {
		return
	}

	job.FinishedAt = StoreTime(time.Now())
	if status == 0 {
		// Successful termination.
		job.Status = StatusDone
	} else {
		// Something went wrong.
		job.Status = StatusError
	}

	// Extract the result from the job.
	if job.ResultSource == "stdout" {
		job.Result = []byte(job.Stdout)
		debug("Acquired job result from stdout: ok")
	} else if strings.HasPrefix(job.ResultSource, "file:") {
		resultPath := job.ResultSource[len("file:"):len(job.ResultSource)]

		var resultBuffer bytes.Buffer
		err = client.CopyFromContainer(docker.CopyFromContainerOptions{
			Container:    container.ID,
			Resource:     resultPath,
			OutputStream: &resultBuffer,
		})
		if checkErr(fmt.Sprintf("Acquired the job's result from the file [%s]", resultPath), err) {
			job.Status = StatusError
		} else {
			// CopyFromContainer returns the file contents as a tarball.
			var content bytes.Buffer
			r := bytes.NewReader(resultBuffer.Bytes())
			tr := tar.NewReader(r)

			for {
				_, err := tr.Next()
				if err == io.EOF {
					break
				}
				if err != nil {
					reportErr("Read tar-encoded content: ERROR", err)
					job.Status = StatusError
					break
				}

				if _, err = io.Copy(&content, tr); err != nil {
					reportErr("Copy decoded content: ERROR", err)
					job.Status = StatusError
					break
				}
			}

			job.Result = content.Bytes()
		}
	}

	err = client.RemoveContainer(docker.RemoveContainerOptions{ID: container.ID})
	checkErr("Removed the container", err)

	err = c.UpdateJob(job)
	if checkErr("Updated the job's status", err) {
		return
	}

	log.WithFields(log.Fields{"jid": job.JID}).Info("Job complete.")
}
<commit_msg>Validate jobs on dequeue, too.<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	docker "github.com/smashwilson/go-dockerclient"
)

// OutputCollector is an io.Writer that accumulates output from a specified stream in an attached
// Docker container and appends it to the appropriate field within a SubmittedJob.
type OutputCollector struct {
	context  *Context
	job      *SubmittedJob
	isStdout bool
}

// DescribeStream returns "stdout" or "stderr" to indicate which stream this collector is consuming.
func (c OutputCollector) DescribeStream() string {
	if c.isStdout {
		return "stdout"
	}
	return "stderr"
}

// Write appends bytes to the selected stream and updates the SubmittedJob.
func (c OutputCollector) Write(p []byte) (int, error) {
	log.WithFields(log.Fields{
		"length": len(p),
		"bytes":  string(p),
		"stream": c.DescribeStream(),
	}).Debug("Received output from a job")

	if c.isStdout {
		c.job.Stdout += string(p)
	} else {
		c.job.Stderr += string(p)
	}

	if err := c.context.UpdateJob(c.job); err != nil {
		return 0, err
	}

	return len(p), nil
}

// Runner is the main entry point for the job runner goroutine.
func Runner(c *Context) {
	var client *docker.Client
	var err error

	if c.DockerTLS {
		client, err = docker.NewTLSClient(c.DockerHost, c.DockerCert, c.DockerKey, c.DockerCACert)
		if err != nil {
			log.WithFields(log.Fields{
				"docker host":    c.DockerHost,
				"docker cert":    c.DockerCert,
				"docker key":     c.DockerKey,
				"docker CA cert": c.DockerCACert,
			}).Fatal("Unable to connect to Docker with TLS.")
			return
		}
	} else {
		client, err = docker.NewClient(c.DockerHost)
		if err != nil {
			log.WithFields(log.Fields{
				"docker host": c.DockerHost,
				"error":       err,
			}).Fatal("Unable to connect to Docker.")
			return
		}
	}

	for {
		select {
		case <-time.After(time.Duration(c.Poll) * time.Millisecond):
			Claim(c, client)
		}
	}
}

// Claim acquires the oldest single pending job and launches a goroutine to execute its command in
// a new container.
func Claim(c *Context, client *docker.Client) {
	job, err := c.ClaimJob()
	if err != nil {
		log.WithFields(log.Fields{"error": err}).Error("Unable to claim a job.")
		return
	}
	if job == nil {
		// Nothing to claim.
		return
	}
	if err := job.Validate(); err != nil {
		fields := log.Fields{
			"jid":     job.JID,
			"account": job.Account,
			"error":   err,
		}

		log.WithFields(fields).Error("Invalid job in queue.")

		job.Status = StatusError
		if err := c.UpdateJob(job); err != nil {
			fields["error"] = err
			log.WithFields(fields).Error("Unable to update job status.")
		}

		return
	}

	go Execute(c, client, job)
}

// Execute launches a container to process the submitted job. It passes any provided stdin data
// to the container and consumes stdout and stderr, updating Mongo as it runs. Once completed, it
// acquires the job's result from its configured source and marks the job as finished.
func Execute(c *Context, client *docker.Client, job *SubmittedJob) {
	defaultFields := log.Fields{
		"jid":     job.JID,
		"account": job.Account,
	}

	// Logging utility messages.
	debug := func(message string) {
		log.WithFields(defaultFields).Debug(message)
	}
	reportErr := func(message string, err error) {
		fs := log.Fields{}
		for k, v := range defaultFields {
			fs[k] = v
		}
		fs["err"] = err
		log.WithFields(fs).Error(message)
	}
	checkErr := func(message string, err error) bool {
		if err == nil {
			debug(fmt.Sprintf("%s: ok", message))
			return false
		}

		reportErr(fmt.Sprintf("%s: ERROR", message), err)
		return true
	}

	log.WithFields(defaultFields).Info("Launching a job.")

	job.StartedAt = StoreTime(time.Now())
	if err := c.UpdateJob(job); err != nil {
		reportErr("Unable to update the job's start timestamp.", err)
	}

	container, err := client.CreateContainer(docker.CreateContainerOptions{
		Name: job.ContainerName(),
		Config: &docker.Config{
			Image:     c.Image,
			Cmd:       []string{"/bin/bash", "-c", job.Command},
			OpenStdin: true,
			StdinOnce: true,
		},
	})
	if checkErr("Created the job's container", err) {
		return
	}

	// Include container information in this job's logging messages.
	defaultFields["container id"] = container.ID
	defaultFields["container name"] = container.Name

	// Prepare the input and output streams.
	stdin := bytes.NewReader(job.Stdin)
	stdout := OutputCollector{
		context:  c,
		job:      job,
		isStdout: true,
	}
	stderr := OutputCollector{
		context:  c,
		job:      job,
		isStdout: false,
	}

	go func() {
		err = client.AttachToContainer(docker.AttachToContainerOptions{
			Container:    container.ID,
			Stream:       true,
			InputStream:  stdin,
			OutputStream: stdout,
			ErrorStream:  stderr,
			Stdin:        true,
			Stdout:       true,
			Stderr:       true,
		})
		checkErr("Attached to the container", err)
	}()

	// Start the created container.
	err = client.StartContainer(container.ID, &docker.HostConfig{})
	if checkErr("Started the container", err) {
		return
	}

	status, err := client.WaitContainer(container.ID)
	if checkErr("Waited for the container to complete", err) {
		return
	}

	job.FinishedAt = StoreTime(time.Now())
	if status == 0 {
		// Successful termination.
		job.Status = StatusDone
	} else {
		// Something went wrong.
		job.Status = StatusError
	}

	// Extract the result from the job.
	if job.ResultSource == "stdout" {
		job.Result = []byte(job.Stdout)
		debug("Acquired job result from stdout: ok")
	} else if strings.HasPrefix(job.ResultSource, "file:") {
		resultPath := job.ResultSource[len("file:"):len(job.ResultSource)]

		var resultBuffer bytes.Buffer
		err = client.CopyFromContainer(docker.CopyFromContainerOptions{
			Container:    container.ID,
			Resource:     resultPath,
			OutputStream: &resultBuffer,
		})
		if checkErr(fmt.Sprintf("Acquired the job's result from the file [%s]", resultPath), err) {
			job.Status = StatusError
		} else {
			// CopyFromContainer returns the file contents as a tarball.
			var content bytes.Buffer
			r := bytes.NewReader(resultBuffer.Bytes())
			tr := tar.NewReader(r)

			for {
				_, err := tr.Next()
				if err == io.EOF {
					break
				}
				if err != nil {
					reportErr("Read tar-encoded content: ERROR", err)
					job.Status = StatusError
					break
				}

				if _, err = io.Copy(&content, tr); err != nil {
					reportErr("Copy decoded content: ERROR", err)
					job.Status = StatusError
					break
				}
			}

			job.Result = content.Bytes()
		}
	}

	err = client.RemoveContainer(docker.RemoveContainerOptions{ID: container.ID})
	checkErr("Removed the container", err)

	err = c.UpdateJob(job)
	if checkErr("Updated the job's status", err) {
		return
	}

	log.WithFields(log.Fields{"jid": job.JID}).Info("Job complete.")
}
<|endoftext|>


<commit_before>package netengine

import (
	"net"
	"time"
	"bytes"
)

type conntion struct {
	ID			int
	C			*NetEngine
	SendChan	chan []byte
	Con			net.Conn
	Notify		NetNotify
	MaxBufLen	int
	RecvBufLen	int

	ReadTimeout		time.Duration
	WriteTimeout	time.Duration
	IsStart   bool
	Send      SendFunc
}

func (con *conntion) get_send() SendFunc {
	writeTimeout := con.WriteTimeout

	r := func(d []byte) error {
		netcon := con.Con

		for {
			if writeTimeout != 0 {
				netcon.SetWriteDeadline(time.Now().Add(writeTimeout))
			}

			n, err := netcon.Write(d)
			if err != nil {
				netcon.Close()
				return err
			}
			if n == len(d) {
				break
			} else {
				d = d[n:]
			}
		}
		return nil
	}
	return r
}

func (con *conntion) start() {
	if con.IsStart {
		return
	}
	con.IsStart = true
	go con.conntion_run()
}
func (con *conntion) close() {
	con.Con.Close()
}
func (con *conntion) conntion_run() {
	/*
	退出流程，write失败，buffer limit 都会关闭链接，然后recv失败，发生del请求，
	del处理会关闭sendchan ，这个时候才会退出send的routine
	只有这样才不会引起崩溃，卡住等问题。
	*/
	go con.conntion_recv()

	con.conntion_write()

	con.Notify.OnClosed(con.ID)
}

func (con *conntion) conntion_recv() {
	net_con := con.Con
	defer net_con.Close()
	defer func() {
		con.C.del_conntion_chan <- con.ID
	}()

	send_fun := con.Send

	buflen := int(con.RecvBufLen)
	all_buf := make([]byte, buflen)

	notify := con.Notify
	readTimeout := con.ReadTimeout

	valid_begin_pos := 0
	valid_end_pos := 0

recv_loop:
	for {
		if readTimeout != 0 {
			net_con.SetReadDeadline(time.Now().Add(readTimeout))
		}

		if valid_end_pos >= len(all_buf) {
			if valid_begin_pos > 2048 {
				copy(all_buf, all_buf[valid_begin_pos:valid_end_pos])
			} else {
				all_len := len(all_buf)
				org_buf := all_buf
				all_buf = make([]byte, all_len*2)
				copy(all_buf, org_buf[valid_begin_pos:valid_end_pos])
			}
			valid_end_pos = valid_end_pos - valid_begin_pos
			valid_begin_pos = 0
		}

		currecv := all_buf[valid_end_pos:len(all_buf)]
		n, err := net_con.Read(currecv)
		if err != nil {
			break
		}
		valid_end_pos = valid_end_pos + n

		for {
			curdata := all_buf[valid_begin_pos:valid_end_pos]
			r := notify.OnRecv(con.ID, curdata, send_fun)

			if r < 0 {
				break recv_loop
			}
			if r > 0 {
				if r > len(curdata) {
					break recv_loop
				}
				valid_begin_pos = valid_begin_pos + r
			} else {
				break
			}
			if valid_begin_pos == valid_end_pos {
				valid_begin_pos = 0
				valid_end_pos = 0
				break
			}
		}
	}
}
func (con *conntion) send_data(data []byte) {
	con.SendChan <- data
}

// 不要轻易修改这个函数的退出方式，会引起崩溃，卡住等问题
func (con *conntion) conntion_write() {
	data_chan := make(chan []byte)
	defer close(data_chan)

	go con.conntion_write_net(data_chan)

	notify := con.Notify

	datalen := 0
	data := make([][]byte, 0, 1)
	maxBufLen := int(con.MaxBufLen)

	timer := time.NewTicker(time.Hour)
	defer timer.Stop()

for_loop:
	for {
		if len(data) > 0 {
			s := data[0]
			select {
			case msg, ok := <-con.SendChan:
				if !ok {
					break for_loop
				}
				datalen += len(msg)
				data = append(data, msg)
				if datalen > maxBufLen {
					notify.OnBufferLimit(con.ID)
					break for_loop
				}
			case data_chan <- s:
				data = data[1:]
				datalen = datalen - len(s)
			case <-timer.C:
			}
		} else {
			select {
			case msg, ok := <-con.SendChan:
				if !ok {
					break for_loop
				}
				datalen += len(msg)
				data = append(data, msg)
				if datalen > maxBufLen {
					notify.OnBufferLimit(con.ID)
					break for_loop
				}
			case <-timer.C:
			}
		}
	}

	con.Con.Close()

	// 尽快清掉内存
	data = nil
	// 消耗数据，等待退出
	for _ = range con.SendChan {
	}
}
func (con *conntion) conntion_write_net(data_chan chan []byte) {
	send := con.Send

	var buf bytes.Buffer
	add_first := false
	add_buf := false

	for first_data := range data_chan {
		add_first = false
		add_buf = false

		if len(first_data) > 60*1024 {
			goto send_pos
		}

send_loop:
		for {
			select {
				case cur_data := <- data_chan :
					add_buf = true

					if !add_first {
						add_first = true
						buf.Write(first_data)
					}

					buf.Write(cur_data)

					if buf.Len() > 60*1024 {
						break send_loop
					}

				default:
					// 没有数据了，直接去发送
					break send_loop
			}
		}

send_pos:
		if !add_buf {
			// 只有一个包
			err := send(first_data)
			if err != nil {
				break
			}
		} else {
			data := buf.Bytes()
			err := send(data)
			if err != nil {
				break
			}

			buf.Truncate(0)
		}
	}

	// 不要轻易修改这个函数的退出方式，会引起崩溃，卡住等问题
	// 消耗数据，等待退出
	for _ = range data_chan {
	}
}
<commit_msg>修正一个关闭连接时可能不退出循环的bug，CPU100%<commit_after>package netengine

import (
	"net"
	"time"
	"bytes"
)

type conntion struct {
	ID			int
	C			*NetEngine
	SendChan	chan []byte
	Con			net.Conn
	Notify		NetNotify
	MaxBufLen	int
	RecvBufLen	int

	ReadTimeout		time.Duration
	WriteTimeout	time.Duration
	IsStart   bool
	Send      SendFunc
}

func (con *conntion) get_send() SendFunc {
	writeTimeout := con.WriteTimeout

	r := func(d []byte) error {
		netcon := con.Con

		for {
			if writeTimeout != 0 {
				netcon.SetWriteDeadline(time.Now().Add(writeTimeout))
			}

			n, err := netcon.Write(d)
			if err != nil {
				netcon.Close()
				return err
			}
			if n == len(d) {
				break
			} else {
				d = d[n:]
			}
		}
		return nil
	}
	return r
}

func (con *conntion) start() {
	if con.IsStart {
		return
	}
	con.IsStart = true
	go con.conntion_run()
}
func (con *conntion) close() {
	con.Con.Close()
}
func (con *conntion) conntion_run() {
	/*
	退出流程，write失败，buffer limit 都会关闭链接，然后recv失败，发生del请求，
	del处理会关闭sendchan ，这个时候才会退出send的routine
	只有这样才不会引起崩溃，卡住等问题。
	*/
	go con.conntion_recv()

	con.conntion_write()

	con.Notify.OnClosed(con.ID)
}

func (con *conntion) conntion_recv() {
	net_con := con.Con
	defer net_con.Close()
	defer func() {
		con.C.del_conntion_chan <- con.ID
	}()

	send_fun := con.Send

	buflen := int(con.RecvBufLen)
	all_buf := make([]byte, buflen)

	notify := con.Notify
	readTimeout := con.ReadTimeout

	valid_begin_pos := 0
	valid_end_pos := 0

recv_loop:
	for {
		if readTimeout != 0 {
			net_con.SetReadDeadline(time.Now().Add(readTimeout))
		}

		if valid_end_pos >= len(all_buf) {
			if valid_begin_pos > 2048 {
				copy(all_buf, all_buf[valid_begin_pos:valid_end_pos])
			} else {
				all_len := len(all_buf)
				org_buf := all_buf
				all_buf = make([]byte, all_len*2)
				copy(all_buf, org_buf[valid_begin_pos:valid_end_pos])
			}
			valid_end_pos = valid_end_pos - valid_begin_pos
			valid_begin_pos = 0
		}

		currecv := all_buf[valid_end_pos:len(all_buf)]
		n, err := net_con.Read(currecv)
		if err != nil {
			break
		}
		valid_end_pos = valid_end_pos + n

		for {
			curdata := all_buf[valid_begin_pos:valid_end_pos]
			r := notify.OnRecv(con.ID, curdata, send_fun)

			if r < 0 {
				break recv_loop
			}
			if r > 0 {
				if r > len(curdata) {
					break recv_loop
				}
				valid_begin_pos = valid_begin_pos + r
			} else {
				break
			}
			if valid_begin_pos == valid_end_pos {
				valid_begin_pos = 0
				valid_end_pos = 0
				break
			}
		}
	}
}
func (con *conntion) send_data(data []byte) {
	con.SendChan <- data
}

// 不要轻易修改这个函数的退出方式，会引起崩溃，卡住等问题
func (con *conntion) conntion_write() {
	data_chan := make(chan []byte)
	defer close(data_chan)

	go con.conntion_write_net(data_chan)

	notify := con.Notify

	datalen := 0
	data := make([][]byte, 0, 1)
	maxBufLen := int(con.MaxBufLen)

	timer := time.NewTicker(time.Hour)
	defer timer.Stop()

for_loop:
	for {
		if len(data) > 0 {
			s := data[0]
			select {
			case msg, ok := <-con.SendChan:
				if !ok {
					break for_loop
				}
				datalen += len(msg)
				data = append(data, msg)
				if datalen > maxBufLen {
					notify.OnBufferLimit(con.ID)
					break for_loop
				}
			case data_chan <- s:
				data = data[1:]
				datalen = datalen - len(s)
			case <-timer.C:
			}
		} else {
			select {
			case msg, ok := <-con.SendChan:
				if !ok {
					break for_loop
				}
				datalen += len(msg)
				data = append(data, msg)
				if datalen > maxBufLen {
					notify.OnBufferLimit(con.ID)
					break for_loop
				}
			case <-timer.C:
			}
		}
	}

	con.Con.Close()

	// 尽快清掉内存
	data = nil
	// 消耗数据，等待退出
	for _ = range con.SendChan {
	}
}
func (con *conntion) conntion_write_net(data_chan chan []byte) {
	send := con.Send

	var buf bytes.Buffer
	add_first := false
	add_buf := false

	for first_data := range data_chan {
		add_first = false
		add_buf = false

		if len(first_data) > 60*1024 {
			goto send_pos
		}

send_loop:
		for {
			select {
				case cur_data,ok := <- data_chan :
					if !ok {
						// data_chan is closed
						break send_loop
					}

					add_buf = true

					if !add_first {
						add_first = true
						buf.Write(first_data)
					}

					buf.Write(cur_data)

					if buf.Len() > 60*1024 {
						break send_loop
					}

				default:
					// 没有数据了，直接去发送
					break send_loop
			}
		}

send_pos:
		if !add_buf {
			// 只有一个包
			err := send(first_data)
			if err != nil {
				break
			}
		} else {
			data := buf.Bytes()
			err := send(data)
			if err != nil {
				break
			}

			buf.Truncate(0)
		}
	}

	// 不要轻易修改这个函数的退出方式，会引起崩溃，卡住等问题
	// 消耗数据，等待退出
	for _ = range data_chan {
	}
}
<|endoftext|>


<commit_before>package main

import s "syscall"
import "fmt"
import "unsafe"

const (
    MAX_AUDIT_MESSAGE_LENGTH = 8970
    AUDIT_GET                = 1000
    AUDIT_LIST               = 1002
    AUDIT_LIST_RULES         = 1013
)

type NetlinkAuditRequest struct {
    Header s.NlMsghdr
    Data   []byte //string
}

type AuditReply struct {

    RepHeader s.NlMsghdr
    Message   NetlinkAuditRequest
}

func (rr *NetlinkAuditRequest) toWireFormat() []byte {
    b := make([]byte, uint32(s.NLMSG_HDRLEN + MAX_AUDIT_MESSAGE_LENGTH))
    *(*uint32)(unsafe.Pointer(&b[0:4][0])) = rr.Header.Len
    *(*uint16)(unsafe.Pointer(&b[4:6][0])) = rr.Header.Type
    *(*uint16)(unsafe.Pointer(&b[6:8][0])) = rr.Header.Flags
    *(*uint32)(unsafe.Pointer(&b[8:12][0])) = rr.Header.Seq
    *(*uint32)(unsafe.Pointer(&b[12:16][0])) = rr.Header.Pid
    //b[16] = byte(rr.Data.Family)
    return b
}

func newNetlinkAuditRequest(proto, seq, family int) []byte {
    rr := &NetlinkAuditRequest{}
    rr.Header.Len = uint32(s.NLMSG_HDRLEN + MAX_AUDIT_MESSAGE_LENGTH) //+ s.SizeofRtGenmsg)
    rr.Header.Type = uint16(proto)
    rr.Header.Flags = s.NLM_F_REQUEST//s.MSG_PEEK|s.MSG_DONTWAIT//s.NLM_F_REQUEST //| s.NLM_F_ACK
    rr.Header.Seq = uint32(seq)
    //rr.Data.Family = uint8(family)
    return rr.toWireFormat()
}

func  newNetlinkAuditReply() []byte {
    rr := &NetlinkAuditRequest{}
    return rr.toWireFormat()
}

func nlmAlignOf(msglen int) int {
    return (msglen + s.NLMSG_ALIGNTO - 1) & ^(s.NLMSG_ALIGNTO - 1)
}

func ParseAuditNetlinkMessage(b []byte) ([]NetlinkAuditRequest, error) {
    var msgs []NetlinkAuditRequest
    for len(b) >= s.NLMSG_HDRLEN {
        h, dbuf, dlen, err := netlinkMessageHeaderAndData(b)
        if err != nil {
            fmt.Println("Error in parse audit")
            return nil, err
        }
        m := NetlinkAuditRequest{Header: *h, Data: dbuf[:int(h.Len)-s.NLMSG_HDRLEN]}
        msgs = append(msgs, m)
        b = b[dlen:]
    }
    return msgs, nil
}

func netlinkMessageHeaderAndData(b []byte) (*s.NlMsghdr, []byte, int, error) {
    h := (*s.NlMsghdr)(unsafe.Pointer(&b[0]))
    //fmt.Println(int(s.NLMSG_HDRLEN));
    if int(h.Len) < s.NLMSG_HDRLEN || int(h.Len) > len(b) {
        fmt.Println("Error Here")
        fmt.Println(s.NLMSG_HDRLEN, h.Len, h.Len, len(b))
        return nil, nil, 0, s.EINVAL
    }
    return h, b[s.NLMSG_HDRLEN:], nlmAlignOf(int(h.Len)), nil
}

func main() {

    // Create a netlink socket
    sock, err := s.Socket(s.AF_NETLINK, s.SOCK_RAW, s.NETLINK_AUDIT)
    if err != nil {
        fmt.Println("scoket error:")
        fmt.Println(err)
        return
    }
    defer s.Close(sock)

    lsa := &s.SockaddrNetlink{Family: s.AF_NETLINK}
    if err := s.Bind(sock, lsa); err != nil {
        fmt.Println("bind error:")
        fmt.Println(err)
        return
    }
    fmt.Println("**** Starting Netlink")

    wb := newNetlinkAuditRequest(1005, 1, s.NETLINK_AUDIT)

    //TODO: try sendmsg here
    if err := s.Sendto(sock, wb, 0, lsa); err != nil {
        fmt.Println("sending error: ", err)
        return
    }
    var tab []byte

done:
    for {

        // r := make([]byte, s.Getpagesize())
        r := newNetlinkAuditReply();

        //nr, _, err := s.Recvfrom(sock, r, s.MSG_PEEK|s.MSG_DONTWAIT)
        nr, _, err := s.Recvfrom(sock, r, 0)

        //r = r[s.NLMSG_HDRLEN:nr]
        if err != nil {
            fmt.Println("rec err: ", err)
            return
        }
        if nr < s.NLMSG_HDRLEN {
            fmt.Println(s.EINVAL)
            fmt.Println("nr < nlmsg hdrlen")
            return
        }

        r = r[:nr]
        tab = append(tab, r...)
        msgs, err := ParseAuditNetlinkMessage(r)

        if err != nil {
            fmt.Println("parse error:", err)
            return
        }

        for _, m := range msgs {

            fmt.Println(string(m.Data))

            lsa, err := s.Getsockname(sock)
            if err != nil {
                fmt.Println("get socket name error:", err)
                return
            }

            switch v := lsa.(type) {
            case *s.SockaddrNetlink:
                if m.Header.Seq != 1 || m.Header.Pid != v.Pid {
                    fmt.Println("case")
                    fmt.Println(s.EINVAL)
                    return
                }
            default:
                fmt.Println(s.EINVAL)
                fmt.Println("default case")
                return
            }
            if m.Header.Type == s.NLMSG_DONE {
                break done
            }

            if m.Header.Type == s.NLMSG_ERROR {
                fmt.Println(s.EINVAL)
                fmt.Println("somewhere below")
                return
            }

        }

        fmt.Println(tab[:]);
        //a,_ := s.ParseNetlinkMessage(tab)
        //fmt.Println(a)
    }
    fmt.Println(tab);
}
<commit_msg>FIx conflict<commit_after>package main

import s "syscall"
import "fmt"
import "unsafe"

const (
    MAX_AUDIT_MESSAGE_LENGTH = 8970
    AUDIT_GET                = 1000
    AUDIT_LIST               = 1002
    AUDIT_LIST_RULES         = 1013
    AUDIT_BITMASK_SIZE       = 64
    AUDIT_MAX_FIELDS         = 64
)

type NetlinkAuditRequest struct {
    Header s.NlMsghdr
    Data   []byte //string
}

/*
struct audit_reply {
387         int                      type;
388         int                      len;
389         struct nlmsghdr         *nlh;
390         struct audit_message     msg;
391
392         /* Using a union to compress this structure since only one of
393          * the following should be valid for any packet.
394         union {
395         struct audit_status     *status;
396         struct audit_rule_data  *ruledata;
397         struct audit_login      *login;
398         const char              *message;
399         struct nlmsgerr         *error;
400         struct audit_sig_info   *signal_info;
401         struct daemon_conf      *conf;
*/

type audit_rule_data struct {
    flags uint32 /* AUDIT_PER_{TASK,CALL}, AUDIT_PREPEND */
    action uint32/* AUDIT_NEVER, AUDIT_POSSIBLE, AUDIT_ALWAYS */
    field_count uint32
    mask [AUDIT_BITMASK_SIZE]uint32 /* syscall(s) affected */
    fields [AUDIT_MAX_FIELDS]uint32
    values [AUDIT_MAX_FIELDS]uint32
    fieldflags [AUDIT_MAX_FIELDS]uint32
    buflen  uint32/* total length of string fields */
    buf string; /* string fields buffer */
};

type AuditReply struct {
    Len   uint32
    Type  uint16
    RepHeader s.NlMsghdr
    msg   NetlinkAuditRequest
    rule  audit_rule_data
}

func (rr *NetlinkAuditRequest) toWireFormat() []byte {
    b := make([]byte, uint32(s.NLMSG_HDRLEN + MAX_AUDIT_MESSAGE_LENGTH))
    *(*uint32)(unsafe.Pointer(&b[0:4][0])) = rr.Header.Len
    *(*uint16)(unsafe.Pointer(&b[4:6][0])) = rr.Header.Type
    *(*uint16)(unsafe.Pointer(&b[6:8][0])) = rr.Header.Flags
    *(*uint32)(unsafe.Pointer(&b[8:12][0])) = rr.Header.Seq
    *(*uint32)(unsafe.Pointer(&b[12:16][0])) = rr.Header.Pid
    //b[16] = byte(rr.Data.Family)
    return b
}

func newNetlinkAuditRequest(proto, seq, family int) []byte {
    rr := &NetlinkAuditRequest{}
    rr.Header.Len = uint32(s.NLMSG_HDRLEN + MAX_AUDIT_MESSAGE_LENGTH) //+ s.SizeofRtGenmsg)
    rr.Header.Type = uint16(proto)
    rr.Header.Flags = s.NLM_F_REQUEST//|s.NLM_F_ACK//s.MSG_PEEK|s.MSG_DONTWAIT//s.NLM_F_REQUEST //| s.NLM_F_ACK
    rr.Header.Seq = uint32(seq)
    //rr.Data.Family = uint8(family)
    return rr.toWireFormat()
}

func  newNetlinkAuditReply() []byte {
    rr := &AuditReply{}
    return rr.msg.toWireFormat()
}

func nlmAlignOf(msglen int) int {
    return (msglen + s.NLMSG_ALIGNTO - 1) & ^(s.NLMSG_ALIGNTO - 1)
}

func ParseAuditNetlinkMessage(b []byte) ([]NetlinkAuditRequest, error) {
    var msgs []NetlinkAuditRequest
    for len(b) >= s.NLMSG_HDRLEN {
        h, dbuf, dlen, err := netlinkMessageHeaderAndData(b)
        if err != nil {
            fmt.Println("Error in parse audit")
            return nil, err
        }
        m := NetlinkAuditRequest{Header: *h, Data: dbuf[:int(h.Len)-s.NLMSG_HDRLEN]}
        msgs = append(msgs, m)
        b = b[dlen:]
    }
    return msgs, nil
}

func netlinkMessageHeaderAndData(b []byte) (*s.NlMsghdr, []byte, int, error) {
    h := (*s.NlMsghdr)(unsafe.Pointer(&b[0]))
    //fmt.Println(int(s.NLMSG_HDRLEN));
    if int(h.Len) < s.NLMSG_HDRLEN || int(h.Len) > len(b) {
        fmt.Println("Error Here")
        fmt.Println(s.NLMSG_HDRLEN, h.Len, h.Len, len(b))
        return nil, nil, 0, s.EINVAL
    }
    return h, b[s.NLMSG_HDRLEN:], nlmAlignOf(int(h.Len)), nil
}

func main() {

    // Create a netlink socket
    sock, err := s.Socket(s.AF_NETLINK, s.SOCK_RAW, s.NETLINK_AUDIT)
    if err != nil {
        fmt.Println("scoket error:")
        fmt.Println(err)
        return
    }
    defer s.Close(sock)

    lsa := &s.SockaddrNetlink{Family: s.AF_NETLINK}
    if err := s.Bind(sock, lsa); err != nil {
        fmt.Println("bind error:")
        fmt.Println(err)
        return
    }
    fmt.Println("**** Starting Netlink")

    wb := newNetlinkAuditRequest(1005, 1, s.NETLINK_AUDIT)

    //TODO: try sendmsg here
    if err := s.Sendto(sock, wb, 0, lsa); err != nil {
        fmt.Println("sending error: ", err)
        return
    }
    var tab []byte

done:
    for {

        // r := make([]byte, s.Getpagesize())
        r := newNetlinkAuditReply();

        //nr, _, err := s.Recvfrom(sock, r, s.MSG_PEEK|s.MSG_DONTWAIT)
        nr, _, err := s.Recvfrom(sock, r, 0)

        //r = r[s.NLMSG_HDRLEN:nr]
        fmt.Println(nr);
        if err != nil {
            fmt.Println("rec err: ", err)
            return
        }
        if nr < s.NLMSG_HDRLEN {
            fmt.Println(s.EINVAL)
            fmt.Println("nr < nlmsg hdrlen")
            return
        }

        r = r[:nr]
        tab = append(tab, r...)
        msgs, err := ParseAuditNetlinkMessage(r)

        if err != nil {
            fmt.Println("parse error:", err)
            return
        }

        for _, m := range msgs {

            fmt.Println(string(m.Data))

            lsa, err := s.Getsockname(sock)
            if err != nil {
                fmt.Println("get socket name error:", err)
                return
            }

            switch v := lsa.(type) {
            case *s.SockaddrNetlink:
                if m.Header.Seq != 1 || m.Header.Pid != v.Pid {
                    fmt.Println("case")
                    fmt.Println(s.EINVAL)
                    return
                }
            default:
                fmt.Println(s.EINVAL)
                fmt.Println("default case")
                return
            }
            if m.Header.Type == s.NLMSG_DONE {
                break done
            }

            if m.Header.Type == s.NLMSG_ERROR {
                fmt.Println(s.EINVAL)
                fmt.Println("somewhere below")
                return
            }

        }

        fmt.Println(tab[:]);
        //a,_ := s.ParseNetlinkMessage(tab)
        //fmt.Println(a)
    }
    fmt.Println(tab);
}
<|endoftext|>


<commit_before>// Copyright 2016 Netflix, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package batched

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"io"
	"math/rand"
	"net"
	"sync/atomic"
	"time"

	"github.com/netflix/rend/binprot"
	"github.com/netflix/rend/common"
	"github.com/netflix/rend/metrics"
)

type conn struct {
	rw           *bufio.ReadWriter
	rand         *rand.Rand
	batchDelay   time.Duration
	batchSize    uint32
	maxBatchSize *uint32
	reqchan      chan request
	expand       chan struct{}
}

func newConn(c net.Conn, batchDelay time.Duration, batchSize uint32, readerSize, writerSize int, expand chan struct{}) conn {
	println("NEWCONN")

	r := bufio.NewReaderSize(c, readerSize)
	w := bufio.NewWriterSize(c, writerSize)

	nc := conn{
		rw:           bufio.NewReadWriter(r, w),
		rand:         rand.New(rand.NewSource(randSeed())),
		batchDelay:   batchDelay,
		batchSize:    batchSize,
		maxBatchSize: new(uint32),
		reqchan:      make(chan request),
		expand:       expand,
	}

	go nc.loop()

	return nc
}

func (c conn) loop() {
	var req request
	var reqs []request
	var timedout bool

	for {
		select {
		case req = <-c.reqchan:
			timedout = false
			// queue up the request
			reqs = append(reqs, req)

		case <-time.After(c.batchDelay * time.Microsecond):
			//println("TIMEDOUT")
			timedout = true
		}
		//println(timedout)
		//fmt.Printf("%#v\n", reqs)

		// After 1 millisecond we want to get the requests that do exist moving along
		// Or, if there's enough to batch together, send them off
		if (timedout && len(reqs) > 0) || len(reqs) >= int(c.batchSize) {
			// store the batch size if it's greater than the maxBatchSize
			for {
				max := atomic.LoadUint32(c.maxBatchSize)
				if uint32(len(reqs)) < max || atomic.CompareAndSwapUint32(c.maxBatchSize, max, uint32(len(reqs))) {
					break
				}
			}

			// If we hit the max size, notify the monitor to add a new connection
			if len(reqs) >= int(c.batchSize) {
				println("NOTIFYING WE NEED TO EXPAND")
				select {
				case c.expand <- struct{}{}:
					println("SUCCESS")
				default:
					println("FAIL")
				}
			}

			c.do(reqs)
			reqs = nil
		}

		// block until a request comes in if there's a timeout earlier so this doesn't constantly spin
		if timedout {
			//println("BLOCKING")
			req = <-c.reqchan
			reqs = append(reqs, req)
			timedout = false
		}
	}
}

func (c conn) do(reqs []request) error {
	//println("DO")

	// batch and perform requests
	buf, responses, channels := c.batchIntoBuffer(reqs)

	//fmt.Printf("%#v\n%#v\n", buf, responses)

	// Write out the whole buffer
	c.rw.Write(buf)
	if err := c.rw.Flush(); err != nil {
		oops(err, responses)
		return err
	}

	//println("flushed")

	// read in all of the responses
	for len(responses) > 0 {
		//println(len(responses))
		//println("READING HEADER")
		resHeader, err := binprot.ReadResponseHeader(c.rw)
		if err != nil {
			oops(err, responses)
			return err
		}
		defer binprot.PutResponseHeader(resHeader)

		//println("GOT HEADER")

		err = binprot.DecodeError(resHeader)
		if err != nil {
			//println("ERROR NOT NIL")
			n, ioerr := c.rw.Discard(int(resHeader.TotalBodyLength))
			metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))

			if ioerr != nil {
				oops(ioerr, responses)
				return err
			}

			if err != common.ErrKeyNotFound && err != common.ErrKeyExists && err != common.ErrItemNotStored {
				println("UH OH", err.Error())
			}

			// TODO: these are not necessarily all misses, check the error
			if rh, ok := responses[resHeader.OpaqueToken]; ok {
				//println("SENDING MISS")
				// this is an application-level error and should be treated as such
				rh.reschan <- response{
					err: nil,
					gr: common.GetEResponse{
						Miss:   true,
						Quiet:  rh.quiet,
						Opaque: rh.opaque,
						Key:    rh.key,
					},
				}
				//println("MISS SENT")

				delete(responses, resHeader.OpaqueToken)
				continue

			} else {
				// FATAL ERROR!!!! no idea what to do here though...
				panic("FATAL ERROR")
			}
		}

		// if reading information (and not just a response header) from the remote
		// process, do some extra parsing
		if resHeader.Opcode == binprot.OpcodeGet ||
			resHeader.Opcode == binprot.OpcodeGetQ ||
			resHeader.Opcode == binprot.OpcodeGat ||
			resHeader.Opcode == binprot.OpcodeGetE ||
			resHeader.Opcode == binprot.OpcodeGetEQ {

			b := make([]byte, 4)
			n, err := io.ReadAtLeast(c.rw, b, 4)
			metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
			if err != nil {
				oops(err, responses)
				return err
			}
			serverFlags := binary.BigEndian.Uint32(b)

			var serverExp uint32
			if resHeader.Opcode == binprot.OpcodeGetE || resHeader.Opcode == binprot.OpcodeGetEQ {
				n, err = io.ReadAtLeast(c.rw, b, 4)
				metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
				if err != nil {
					oops(err, responses)
					return err
				}
				serverExp = binary.BigEndian.Uint32(b)
			}

			// total body - key - extra
			dataLen := resHeader.TotalBodyLength - uint32(resHeader.KeyLength) - uint32(resHeader.ExtraLength)
			buf := make([]byte, dataLen)

			// Read in value
			n, err = io.ReadAtLeast(c.rw, buf, int(dataLen))
			metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
			if err != nil {
				oops(err, responses)
				return err
			}

			if rh, ok := responses[resHeader.OpaqueToken]; ok {
				// send the response back on the channel assigned to this token
				rh.reschan <- response{
					err: nil,
					gr: common.GetEResponse{
						Key:     rh.key,
						Data:    buf,
						Flags:   serverFlags,
						Exptime: serverExp,
						Opaque:  rh.opaque,
						Quiet:   rh.quiet,
					},
				}

				delete(responses, resHeader.OpaqueToken)

			} else {
				// we are out of sync here, something is really wrong. We got the wrong opaque
				// value for the set of requests we thing we sent. We can't fix this, so we have
				// to close the connection
				panic("FATAL ERROR 2")
			}

		} else {
			// Non-get repsonses
			// Discard the message for non-get responses
			n, err := c.rw.Discard(int(resHeader.TotalBodyLength))
			metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
			if err != nil {
				oops(err, responses)
				return err
			}

			if rh, ok := responses[resHeader.OpaqueToken]; ok {
				// send the response back on the channel assigned to this token
				rh.reschan <- response{}
				delete(responses, resHeader.OpaqueToken)

			} else {
				// we are out of sync here, something is really wrong. We got the wrong opaque
				// value for the set of requests we thing we sent. We can't fix this, so we have
				// to close the connection
				panic("FATAL ERROR 3")
			}
		}
	}

	//println("CLOSING CHANNELS")
	for _, c := range channels {
		close(c)
	}

	return nil
}

// Need a struct here to hold all the metadata about the request.
// Responses don't have all the metadata, so it needs to be kept to the side
type reshandle struct {
	key     []byte
	opaque  uint32
	quiet   bool
	reschan chan response
}

// If something goes seriously wrong (i.e. an I/O error on this connection) then send back the
// error to all of the connections waiting on responses from this connection.
func oops(err error, res map[uint32]reshandle) {
	println("OOPS", err.Error())
	for _, c := range res {
		c.reschan <- response{err: err}
	}
}

// The opaque value in the requests is related to the base value returned by its index in the array
// meaning the first request sent out will have the opaque value <base>, the second <base+1>, and so on
// The base is a random uint32 value. This is done to prevent possible confusion between requests,
// where one request would get another's data. This would be really bad if we sent things like user
// information to the wrong place.
func (c conn) batchIntoBuffer(reqs []request) ([]byte, map[uint32]reshandle, []chan response) {
	// get random base value
	// serialize all requests into a buffer
	// while serializing, collect channels
	// return everything

	// Get base opaque value
	opaque := uint32(c.rand.Int31())
	buf := new(bytes.Buffer)
	responses := make(map[uint32]reshandle)
	channels := make([]chan response, 0, len(reqs))

	for _, req := range reqs {

		// maintain the "sequence number"
		opaque++

		// build slice of channels for closing later
		channels = append(channels, req.reschan)

		// serialize the requests into the buffer. There's no error handling because writes to a bytes.Buffer
		// never return an error.
		switch req.reqtype {
		case common.RequestSet:
			cmd := req.req.(common.SetRequest)
			binprot.WriteSetCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestAdd:
			cmd := req.req.(common.SetRequest)
			binprot.WriteAddCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestReplace:
			cmd := req.req.(common.SetRequest)
			binprot.WriteReplaceCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestAppend:
			cmd := req.req.(common.SetRequest)
			binprot.WriteAppendCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestPrepend:
			cmd := req.req.(common.SetRequest)
			binprot.WritePrependCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestDelete:
			cmd := req.req.(common.DeleteRequest)
			binprot.WriteDeleteCmd(buf, cmd.Key, opaque)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestTouch:
			cmd := req.req.(common.TouchRequest)
			binprot.WriteTouchCmd(buf, cmd.Key, cmd.Exptime, opaque)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestGat:
			cmd := req.req.(common.GATRequest)
			binprot.WriteGATCmd(buf, cmd.Key, cmd.Exptime, opaque)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestGet:
			cmd := req.req.(common.GetRequest)

			for idx := range cmd.Keys {
				binprot.WriteGetCmd(buf, cmd.Keys[idx], opaque)
				responses[opaque] = reshandle{
					key:     cmd.Keys[idx],
					opaque:  cmd.Opaques[idx],
					quiet:   cmd.Quiet[idx],
					reschan: req.reschan,
				}
				opaque++
			}

		case common.RequestGetE:
			cmd := req.req.(common.GetRequest)

			for idx := range cmd.Keys {
				binprot.WriteGetECmd(buf, cmd.Keys[idx], opaque)
				responses[opaque] = reshandle{
					key:     cmd.Keys[idx],
					opaque:  cmd.Opaques[idx],
					quiet:   cmd.Quiet[idx],
					reschan: req.reschan,
				}
				opaque++
			}
		}
	}

	//fmt.Printf("%#v\n", responses)

	return buf.Bytes(), responses, channels
}
<commit_msg>Splitting writing and reading side for each poole connection to support large batches<commit_after>// Copyright 2016 Netflix, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package batched

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"io"
	"math/rand"
	"net"
	"sync/atomic"
	"time"

	"github.com/netflix/rend/binprot"
	"github.com/netflix/rend/common"
	"github.com/netflix/rend/metrics"
)

type conn struct {
	rw           *bufio.ReadWriter
	rand         *rand.Rand
	batchDelay   time.Duration
	batchSize    uint32
	maxBatchSize *uint32
	reqchan      chan request
	batchchan    chan batch
	expand       chan struct{}
}

type batch struct {
	responses map[uint32]reshandle
	channels  []chan response
}

func newConn(c net.Conn, batchDelay time.Duration, batchSize uint32, readerSize, writerSize int, expand chan struct{}) conn {
	println("NEWCONN")

	r := bufio.NewReaderSize(c, readerSize)
	w := bufio.NewWriterSize(c, writerSize)

	nc := conn{
		rw:           bufio.NewReadWriter(r, w),
		rand:         rand.New(rand.NewSource(randSeed())),
		batchDelay:   batchDelay,
		batchSize:    batchSize,
		maxBatchSize: new(uint32),
		reqchan:      make(chan request),
		batchchan:    make(chan batch, 1),
		expand:       expand,
	}

	go nc.batcher()
	go nc.reader()

	return nc
}

func (c conn) batcher() {
	var req request
	var reqs []request
	var timedout bool

	for {
		//println("LOOP")
		select {
		case req = <-c.reqchan:
			timedout = false
			// queue up the request
			reqs = append(reqs, req)

		case <-time.After(c.batchDelay * time.Microsecond):
			//println("TIMEDOUT")
			timedout = true
		}
		//println(timedout)
		//fmt.Printf("%#v\n", reqs)

		// After 1 millisecond we want to get the requests that do exist moving along
		// Or, if there's enough to batch together, send them off
		if (timedout && len(reqs) > 0) || len(reqs) >= int(c.batchSize) {
			// store the batch size if it's greater than the maxBatchSize
			for {
				max := atomic.LoadUint32(c.maxBatchSize)
				if uint32(len(reqs)) < max || atomic.CompareAndSwapUint32(c.maxBatchSize, max, uint32(len(reqs))) {
					break
				}
			}

			// If we hit the max size, notify the monitor to add a new connection
			if len(reqs) >= int(c.batchSize) {
				println("NOTIFYING WE NEED TO EXPAND")
				select {
				case c.expand <- struct{}{}:
					println("SUCCESS")
				default:
					println("FAIL")
				}
			}

			c.do(reqs)
			reqs = nil
		}

		// block until a request comes in if there's a timeout earlier so this doesn't constantly spin
		if timedout {
			//println("BLOCKING")
			req = <-c.reqchan
			reqs = append(reqs, req)
			timedout = false
		}
	}
}

func (c conn) do(reqs []request) {
	//println("DO")

	// batch and perform requests
	buf, responses, channels := c.batchIntoBuffer(reqs)

	//fmt.Printf("%#v\n%#v\n", buf, responses)

	// send the batch before the write because the write may block for a long time until
	// some responses can be read from memcached.
	c.batchchan <- batch{
		responses: responses,
		channels:  channels,
	}

	// Write out the whole buffer
	c.rw.Write(buf)
	if err := c.rw.Flush(); err != nil {
		oops(err, responses)
	}

	//println("flushed")
}

func (c conn) reader() {
	for {

		batch := <-c.batchchan

		// read in all of the responses
		for len(batch.responses) > 0 {
			//println(len(responses))
			//println("READING HEADER")
			resHeader, err := binprot.ReadResponseHeader(c.rw)
			if err != nil {
				oops(err, batch.responses)
			}
			defer binprot.PutResponseHeader(resHeader)

			//println("GOT HEADER")

			err = binprot.DecodeError(resHeader)
			if err != nil {
				//println("ERROR NOT NIL")
				n, ioerr := c.rw.Discard(int(resHeader.TotalBodyLength))
				metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))

				if ioerr != nil {
					oops(ioerr, batch.responses)
				}

				if err != common.ErrKeyNotFound && err != common.ErrKeyExists && err != common.ErrItemNotStored {
					println("UH OH", err.Error())
				}

				// TODO: these are not necessarily all misses, check the error
				if rh, ok := batch.responses[resHeader.OpaqueToken]; ok {
					//println("SENDING MISS")
					// this is an application-level error and should be treated as such
					rh.reschan <- response{
						err: nil,
						gr: common.GetEResponse{
							Miss:   true,
							Quiet:  rh.quiet,
							Opaque: rh.opaque,
							Key:    rh.key,
						},
					}
					//println("MISS SENT")

					delete(batch.responses, resHeader.OpaqueToken)
					continue

				} else {
					// FATAL ERROR!!!! no idea what to do here though...
					panic("FATAL ERROR")
				}
			}

			// if reading information (and not just a response header) from the remote
			// process, do some extra parsing
			if resHeader.Opcode == binprot.OpcodeGet ||
				resHeader.Opcode == binprot.OpcodeGetQ ||
				resHeader.Opcode == binprot.OpcodeGat ||
				resHeader.Opcode == binprot.OpcodeGetE ||
				resHeader.Opcode == binprot.OpcodeGetEQ {

				b := make([]byte, 4)
				n, err := io.ReadAtLeast(c.rw, b, 4)
				metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
				if err != nil {
					oops(err, batch.responses)
				}
				serverFlags := binary.BigEndian.Uint32(b)

				var serverExp uint32
				if resHeader.Opcode == binprot.OpcodeGetE || resHeader.Opcode == binprot.OpcodeGetEQ {
					n, err = io.ReadAtLeast(c.rw, b, 4)
					metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
					if err != nil {
						oops(err, batch.responses)
					}
					serverExp = binary.BigEndian.Uint32(b)
				}

				// total body - key - extra
				dataLen := resHeader.TotalBodyLength - uint32(resHeader.KeyLength) - uint32(resHeader.ExtraLength)
				buf := make([]byte, dataLen)

				// Read in value
				n, err = io.ReadAtLeast(c.rw, buf, int(dataLen))
				metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
				if err != nil {
					oops(err, batch.responses)
				}

				if rh, ok := batch.responses[resHeader.OpaqueToken]; ok {
					// send the response back on the channel assigned to this token
					rh.reschan <- response{
						err: nil,
						gr: common.GetEResponse{
							Key:     rh.key,
							Data:    buf,
							Flags:   serverFlags,
							Exptime: serverExp,
							Opaque:  rh.opaque,
							Quiet:   rh.quiet,
						},
					}

					delete(batch.responses, resHeader.OpaqueToken)

				} else {
					// we are out of sync here, something is really wrong. We got the wrong opaque
					// value for the set of requests we thing we sent. We can't fix this, so we have
					// to close the connection
					panic("FATAL ERROR 2")
				}

			} else {
				// Non-get repsonses
				// Discard the message for non-get responses
				n, err := c.rw.Discard(int(resHeader.TotalBodyLength))
				metrics.IncCounterBy(common.MetricBytesReadLocal, uint64(n))
				if err != nil {
					oops(err, batch.responses)
				}

				if rh, ok := batch.responses[resHeader.OpaqueToken]; ok {
					// send the response back on the channel assigned to this token
					rh.reschan <- response{}
					delete(batch.responses, resHeader.OpaqueToken)

				} else {
					// we are out of sync here, something is really wrong. We got the wrong opaque
					// value for the set of requests we thing we sent. We can't fix this, so we have
					// to close the connection
					panic("FATAL ERROR 3")
				}
			}
		}

		//println("CLOSING CHANNELS")
		for _, c := range batch.channels {
			close(c)
		}
	}
}

// Need a struct here to hold all the metadata about the request.
// Responses don't have all the metadata, so it needs to be kept to the side
type reshandle struct {
	key     []byte
	opaque  uint32
	quiet   bool
	reschan chan response
}

// If something goes seriously wrong (i.e. an I/O error on this connection) then send back the
// error to all of the connections waiting on responses from this connection.
func oops(err error, res map[uint32]reshandle) {
	println("OOPS", err.Error())
	for _, c := range res {
		c.reschan <- response{err: err}
	}
}

// The opaque value in the requests is related to the base value returned by its index in the array
// meaning the first request sent out will have the opaque value <base>, the second <base+1>, and so on
// The base is a random uint32 value. This is done to prevent possible confusion between requests,
// where one request would get another's data. This would be really bad if we sent things like user
// information to the wrong place.
func (c conn) batchIntoBuffer(reqs []request) ([]byte, map[uint32]reshandle, []chan response) {
	// get random base value
	// serialize all requests into a buffer
	// while serializing, collect channels
	// return everything

	// Get base opaque value
	opaque := uint32(c.rand.Int31())
	buf := new(bytes.Buffer)
	responses := make(map[uint32]reshandle)
	channels := make([]chan response, 0, len(reqs))

	for _, req := range reqs {

		// maintain the "sequence number"
		opaque++

		// build slice of channels for closing later
		channels = append(channels, req.reschan)

		// serialize the requests into the buffer. There's no error handling because writes to a bytes.Buffer
		// never return an error.
		switch req.reqtype {
		case common.RequestSet:
			cmd := req.req.(common.SetRequest)
			binprot.WriteSetCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestAdd:
			cmd := req.req.(common.SetRequest)
			binprot.WriteAddCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestReplace:
			cmd := req.req.(common.SetRequest)
			binprot.WriteReplaceCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestAppend:
			cmd := req.req.(common.SetRequest)
			binprot.WriteAppendCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestPrepend:
			cmd := req.req.(common.SetRequest)
			binprot.WritePrependCmd(buf, cmd.Key, cmd.Flags, cmd.Exptime, uint32(len(cmd.Data)), opaque)
			buf.Write(cmd.Data)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestDelete:
			cmd := req.req.(common.DeleteRequest)
			binprot.WriteDeleteCmd(buf, cmd.Key, opaque)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestTouch:
			cmd := req.req.(common.TouchRequest)
			binprot.WriteTouchCmd(buf, cmd.Key, cmd.Exptime, opaque)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestGat:
			cmd := req.req.(common.GATRequest)
			binprot.WriteGATCmd(buf, cmd.Key, cmd.Exptime, opaque)
			responses[opaque] = reshandle{
				key:     cmd.Key,
				opaque:  cmd.Opaque,
				quiet:   cmd.Quiet,
				reschan: req.reschan,
			}

		case common.RequestGet:
			cmd := req.req.(common.GetRequest)

			for idx := range cmd.Keys {
				binprot.WriteGetCmd(buf, cmd.Keys[idx], opaque)
				responses[opaque] = reshandle{
					key:     cmd.Keys[idx],
					opaque:  cmd.Opaques[idx],
					quiet:   cmd.Quiet[idx],
					reschan: req.reschan,
				}
				opaque++
			}

		case common.RequestGetE:
			cmd := req.req.(common.GetRequest)

			for idx := range cmd.Keys {
				binprot.WriteGetECmd(buf, cmd.Keys[idx], opaque)
				responses[opaque] = reshandle{
					key:     cmd.Keys[idx],
					opaque:  cmd.Opaques[idx],
					quiet:   cmd.Quiet[idx],
					reschan: req.reschan,
				}
				opaque++
			}
		}
	}

	//fmt.Printf("%#v\n", responses)
	//println(buf.Len())
	//fmt.Println(buf.Bytes())

	return buf.Bytes(), responses, channels
}
<|endoftext|>


<commit_before>package oortstore

import (
	"bytes"
	"log"
	"os"
	"sync"
	"time"

	"github.com/gholt/ring"
	"github.com/gholt/valuestore"
	"github.com/pandemicsyn/oort/oort"
	"github.com/pandemicsyn/oort/rediscache"
	"github.com/spaolacci/murmur3"
	"gopkg.in/gholt/brimtime.v1"
)

type OortStore struct {
	sync.RWMutex
	vs      valuestore.ValueStore
	t       *ring.TCPMsgRing
	o       *oort.Server
	c       *Config
	stopped bool
}

type Config struct {
	Debug   bool
	Profile bool
}

func New(oort *oort.Server, config *Config) *OortStore {
	s := &OortStore{}
	s.o = oort
	s.c = config
	if s.c.Debug {
		log.Println("Ring entries:")
		ring := s.o.Ring()
		for k, _ := range ring.Nodes() {
			log.Println(ring.Nodes()[k].ID(), ring.Nodes()[k].Addresses())
		}
	}
	log.Println("LocalID appears to be:", s.o.GetLocalID())
	s.t = ring.NewTCPMsgRing(nil)
	s.t.SetRing(s.o.Ring())
	l := log.New(os.Stdout, "DebugStore ", log.LstdFlags)
	s.o.ValueStoreConfig.MsgRing = s.t
	s.o.ValueStoreConfig.LogDebug = l.Printf
	var err error
	s.vs, err = valuestore.New(&s.o.ValueStoreConfig)
	if err != nil {
		panic(err)
	}
	s.vs.EnableAll()
	go func() {
		s.t.Listen()
		log.Println("Listen() returned, shutdown?")
	}()
	go func() {
		tcpMsgRingStats := s.t.Stats(false)
		for !tcpMsgRingStats.Shutdown {
			time.Sleep(time.Minute)
			tcpMsgRingStats = s.t.Stats(false)
			log.Printf("%v\n", tcpMsgRingStats)
			log.Printf("%s\n", s.vs.Stats(false))
		}
	}()
	return s
}

func (vsc *OortStore) UpdateRing(ring ring.Ring) {
	vsc.Lock()
	vsc.t.SetRing(ring)
	vsc.Unlock()
	log.Println("Oortstore updated tcp msg ring.")
}

func (vsc *OortStore) Get(key []byte, value []byte) []byte {
	keyA, keyB := murmur3.Sum128(key)
	var err error
	_, value, err = vsc.vs.Read(keyA, keyB, value)
	if err != nil {
		log.Printf("Get: %#v %s\n", string(key), err)
	}
	return value
}

func (vsc *OortStore) Set(key []byte, value []byte) {
	if bytes.Equal(key, rediscache.BYTES_SHUTDOWN) && bytes.Equal(value, rediscache.BYTES_NOW) {
		vsc.vs.DisableAll()
		vsc.vs.Flush()
		log.Println(vsc.vs.Stats(true))
		os.Exit(0)
		return
	}
	keyA, keyB := murmur3.Sum128(key)
	_, err := vsc.vs.Write(keyA, keyB, brimtime.TimeToUnixMicro(time.Now()), value)
	if err != nil {
		panic(err)
	}
}

func (vsc *OortStore) Stop() {
	vsc.Lock()
	if vsc.stopped {
		vsc.Unlock()
		return
	}
	vsc.vs.DisableAll()
	vsc.vs.Flush()
	vsc.t.Shutdown()
	vsc.Unlock()
	log.Println(vsc.vs.Stats(true))
	log.Println("Ortstore stop complete")
}

func (vsc *OortStore) Stats() []byte {
	return []byte(vsc.vs.Stats(true).String())
}
<commit_msg>Add ability to actually start gholt valuestore remotely<commit_after>package oortstore

import (
	"bytes"
	"log"
	"os"
	"sync"
	"time"

	"github.com/gholt/ring"
	"github.com/gholt/valuestore"
	"github.com/pandemicsyn/oort/oort"
	"github.com/pandemicsyn/oort/rediscache"
	"github.com/spaolacci/murmur3"
	"gopkg.in/gholt/brimtime.v1"
)

type OortStore struct {
	sync.RWMutex
	vs      valuestore.ValueStore
	t       *ring.TCPMsgRing
	o       *oort.Server
	c       *Config
	stopped bool
}

type Config struct {
	Debug   bool
	Profile bool
}

func New(oort *oort.Server, config *Config) *OortStore {
	s := &OortStore{}
	s.o = oort
	s.c = config
	if s.c.Debug {
		log.Println("Ring entries:")
		ring := s.o.Ring()
		for k, _ := range ring.Nodes() {
			log.Println(ring.Nodes()[k].ID(), ring.Nodes()[k].Addresses())
		}
	}
	log.Println("LocalID appears to be:", s.o.GetLocalID())
	s.t = ring.NewTCPMsgRing(nil)
	s.t.SetRing(s.o.Ring())
	l := log.New(os.Stdout, "DebugStore ", log.LstdFlags)
	s.o.ValueStoreConfig.MsgRing = s.t
	s.o.ValueStoreConfig.LogDebug = l.Printf
	s.start()
	return s
}

func (vsc *OortStore) start() {
	var err error
	vsc.vs, err = valuestore.New(&vsc.o.ValueStoreConfig)
	if err != nil {
		panic(err)
	}
	vsc.vs.EnableAll()
	go func() {
		vsc.t.Listen()
		log.Println("Listen() returned, shutdown?")
	}()
	go func() {
		tcpMsgRingStats := vsc.t.Stats(false)
		for !tcpMsgRingStats.Shutdown {
			time.Sleep(time.Minute)
			tcpMsgRingStats = vsc.t.Stats(false)
			log.Printf("%v\n", tcpMsgRingStats)
			log.Printf("%s\n", vsc.vs.Stats(false))
		}
	}()
}

func (vsc *OortStore) UpdateRing(ring ring.Ring) {
	vsc.Lock()
	vsc.t.SetRing(ring)
	vsc.Unlock()
	log.Println("Oortstore updated tcp msg ring.")
}

func (vsc *OortStore) Get(key []byte, value []byte) []byte {
	keyA, keyB := murmur3.Sum128(key)
	var err error
	_, value, err = vsc.vs.Read(keyA, keyB, value)
	if err != nil {
		log.Printf("Get: %#v %s\n", string(key), err)
	}
	return value
}

func (vsc *OortStore) Set(key []byte, value []byte) {
	if bytes.Equal(key, rediscache.BYTES_SHUTDOWN) && bytes.Equal(value, rediscache.BYTES_NOW) {
		vsc.vs.DisableAll()
		vsc.vs.Flush()
		log.Println(vsc.vs.Stats(true))
		os.Exit(0)
		return
	}
	keyA, keyB := murmur3.Sum128(key)
	_, err := vsc.vs.Write(keyA, keyB, brimtime.TimeToUnixMicro(time.Now()), value)
	if err != nil {
		panic(err)
	}
}

func (vsc *OortStore) Start() {
	vsc.Lock()
	if !vsc.stopped {
		vsc.Unlock()
		return
	}
	vsc.start()
	vsc.stopped = false
	vsc.Unlock()
	log.Println(vsc.vs.Stats(true))
	log.Println("Ortstore stop complete")
}

func (vsc *OortStore) Stop() {
	vsc.Lock()
	if vsc.stopped {
		vsc.Unlock()
		return
	}
	vsc.vs.DisableAll()
	vsc.vs.Flush()
	vsc.t.Shutdown()
	vsc.stopped = true
	vsc.Unlock()
	log.Println(vsc.vs.Stats(true))
	log.Println("Ortstore stop complete")
}

func (vsc *OortStore) Stats() []byte {
	return []byte(vsc.vs.Stats(true).String())
}
<|endoftext|>


<commit_before>package models

type StagingRequestFromCC struct {
	AppId                          string                `json:"app_id"`
	TaskId                         string                `json:"task_id"`
	Stack                          string                `json:"stack"`
	AppBitsDownloadUri             string                `json:"app_bits_download_uri"`
	BuildArtifactsCacheDownloadUri string                `json:"build_artifacts_cache_download_uri,omitempty"`
	FileDescriptors                int                   `json:"file_descriptors"`
	MemoryMB                       int                   `json:"memory_mb"`
	DiskMB                         int                   `json:"disk_mb"`
	Buildpacks                     []Buildpack           `json:"buildpacks"`
	Environment                    []EnvironmentVariable `json:"environment"`
}

type Buildpack struct {
	Name string `json:"name"`
	Key  string `json:"key"`
	Url  string `json:"url"`
}

type EnvironmentVariable struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

type StagingInfo struct {
	BuildpackKey      string `yaml:"-" json:"buildpack_key,omitempty"`
	DetectedBuildpack string `yaml:"detected_buildpack" json:"detected_buildpack"`
	StartCommand      string `yaml:"start_command" json:"-"`
}

type StagingResponseForCC struct {
	AppId             string `json:"app_id,omitempty"`
	TaskId            string `json:"task_id,omitempty"`
	BuildpackKey      string `json:"buildpack_key,omitempty"`
	DetectedBuildpack string `json:"detected_buildpack,omitempty"`
	Error             string `json:"error,omitempty"`
}

type StagingTaskAnnotation struct {
	AppId  string `json:"app_id"`
	TaskId string `json:"task_id"`
}
<commit_msg>add start command<commit_after>package models

type StagingRequestFromCC struct {
	AppId                          string                `json:"app_id"`
	TaskId                         string                `json:"task_id"`
	Stack                          string                `json:"stack"`
	AppBitsDownloadUri             string                `json:"app_bits_download_uri"`
	BuildArtifactsCacheDownloadUri string                `json:"build_artifacts_cache_download_uri,omitempty"`
	FileDescriptors                int                   `json:"file_descriptors"`
	MemoryMB                       int                   `json:"memory_mb"`
	DiskMB                         int                   `json:"disk_mb"`
	Buildpacks                     []Buildpack           `json:"buildpacks"`
	Environment                    []EnvironmentVariable `json:"environment"`
}

type Buildpack struct {
	Name string `json:"name"`
	Key  string `json:"key"`
	Url  string `json:"url"`
}

type EnvironmentVariable struct {
	Key   string `json:"key"`
	Value string `json:"value"`
}

type StagingInfo struct {
	BuildpackKey      string `yaml:"-" json:"buildpack_key,omitempty"`
	DetectedBuildpack string `yaml:"detected_buildpack" json:"detected_buildpack"`
	StartCommand      string `yaml:"start_command" json:"start_command"`
}

type StagingResponseForCC struct {
	AppId             string `json:"app_id,omitempty"`
	TaskId            string `json:"task_id,omitempty"`
	BuildpackKey      string `json:"buildpack_key,omitempty"`
	DetectedBuildpack string `json:"detected_buildpack,omitempty"`
	StartCommand      string `json:"start_command,omitempty"`
	Error             string `json:"error,omitempty"`
}

type StagingTaskAnnotation struct {
	AppId  string `json:"app_id"`
	TaskId string `json:"task_id"`
}
<|endoftext|>


<commit_before>package lfshttp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/git-lfs/git-lfs/config"
	"github.com/git-lfs/git-lfs/subprocess"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/rubyist/tracerx"
)

type SSHResolver interface {
	Resolve(Endpoint, string) (sshAuthResponse, error)
}

func withSSHCache(ssh SSHResolver) SSHResolver {
	return &sshCache{
		endpoints: make(map[string]*sshAuthResponse),
		ssh:       ssh,
	}
}

type sshCache struct {
	endpoints map[string]*sshAuthResponse
	ssh       SSHResolver
}

func (c *sshCache) Resolve(e Endpoint, method string) (sshAuthResponse, error) {
	if len(e.SshUserAndHost) == 0 {
		return sshAuthResponse{}, nil
	}

	key := strings.Join([]string{e.SshUserAndHost, e.SshPort, e.SshPath, method}, "//")
	if res, ok := c.endpoints[key]; ok {
		if _, expired := res.IsExpiredWithin(5 * time.Second); !expired {
			tracerx.Printf("ssh cache: %s git-lfs-authenticate %s %s",
				e.SshUserAndHost, e.SshPath, endpointOperation(e, method))
			return *res, nil
		} else {
			tracerx.Printf("ssh cache expired: %s git-lfs-authenticate %s %s",
				e.SshUserAndHost, e.SshPath, endpointOperation(e, method))
		}
	}

	res, err := c.ssh.Resolve(e, method)
	if err == nil {
		c.endpoints[key] = &res
	}
	return res, err
}

type sshAuthResponse struct {
	Message   string            `json:"-"`
	Href      string            `json:"href"`
	Header    map[string]string `json:"header"`
	ExpiresAt time.Time         `json:"expires_at"`
	ExpiresIn int               `json:"expires_in"`

	createdAt time.Time
}

func (r *sshAuthResponse) IsExpiredWithin(d time.Duration) (time.Time, bool) {
	return tools.IsExpiredAtOrIn(r.createdAt, d, r.ExpiresAt,
		time.Duration(r.ExpiresIn)*time.Second)
}

type sshAuthClient struct {
	os  config.Environment
	git config.Environment
}

func (c *sshAuthClient) Resolve(e Endpoint, method string) (sshAuthResponse, error) {
	res := sshAuthResponse{}
	if len(e.SshUserAndHost) == 0 {
		return res, nil
	}

	exe, args := sshGetLFSExeAndArgs(c.os, c.git, e, method)
	cmd := exec.Command(exe, args...)

	// Save stdout and stderr in separate buffers
	var outbuf, errbuf bytes.Buffer
	cmd.Stdin = os.Stdin
	cmd.Stdout = &outbuf
	cmd.Stderr = &errbuf

	now := time.Now()

	// Execute command
	err := cmd.Start()
	if err == nil {
		err = cmd.Wait()
	}

	// Processing result
	if err != nil {
		res.Message = strings.TrimSpace(errbuf.String())
	} else {
		err = json.Unmarshal(outbuf.Bytes(), &res)
		if res.ExpiresIn == 0 && res.ExpiresAt.IsZero() {
			ttl := c.git.Int("lfs.defaulttokenttl", 0)
			if ttl < 0 {
				ttl = 0
			}
			res.ExpiresIn = ttl
		}
		res.createdAt = now
	}

	return res, err
}

func sshFormatArgs(cmd string, args []string, needShell bool) (string, []string) {
	if !needShell {
		return cmd, args
	}

	return subprocess.FormatForShellQuotedArgs(cmd, args)
}

func sshGetLFSExeAndArgs(osEnv config.Environment, gitEnv config.Environment, e Endpoint, method string) (string, []string) {
	exe, args, needShell := sshGetExeAndArgs(osEnv, gitEnv, e)
	operation := endpointOperation(e, method)
	args = append(args, fmt.Sprintf("git-lfs-authenticate %s %s", e.SshPath, operation))
	exe, args = sshFormatArgs(exe, args, needShell)
	tracerx.Printf("run_command: %s %s", exe, strings.Join(args, " "))
	return exe, args
}

// Parse command, and if it looks like a valid command, return the ssh binary
// name, the command to run, and whether we need a shell.  If not, return
// existing as the ssh binary name.
func sshParseShellCommand(command string, existing string) (ssh string, cmd string, needShell bool) {
	ssh = existing
	if cmdArgs := tools.QuotedFields(command); len(cmdArgs) > 0 {
		needShell = true
		ssh = cmdArgs[0]
		cmd = command
	}
	return
}

// Return the executable name for ssh on this machine and the base args
// Base args includes port settings, user/host, everything pre the command to execute
func sshGetExeAndArgs(osEnv config.Environment, gitEnv config.Environment, e Endpoint) (exe string, baseargs []string, needShell bool) {
	var cmd string

	isPlink := false
	isTortoise := false

	ssh, _ := osEnv.Get("GIT_SSH")
	sshCmd, _ := osEnv.Get("GIT_SSH_COMMAND")
	ssh, cmd, needShell = sshParseShellCommand(sshCmd, ssh)

	if ssh == "" {
		sshCmd, _ := gitEnv.Get("core.sshcommand")
		ssh, cmd, needShell = sshParseShellCommand(sshCmd, defaultSSHCmd)
	}

	if cmd == "" {
		cmd = ssh
	}

	basessh := filepath.Base(ssh)

	if basessh != defaultSSHCmd {
		// Strip extension for easier comparison
		if ext := filepath.Ext(basessh); len(ext) > 0 {
			basessh = basessh[:len(basessh)-len(ext)]
		}
		isPlink = strings.EqualFold(basessh, "plink")
		isTortoise = strings.EqualFold(basessh, "tortoiseplink")
	}

	args := make([]string, 0, 7)

	if isTortoise {
		// TortoisePlink requires the -batch argument to behave like ssh/plink
		args = append(args, "-batch")
	}

	if len(e.SshPort) > 0 {
		if isPlink || isTortoise {
			args = append(args, "-P")
		} else {
			args = append(args, "-p")
		}
		args = append(args, e.SshPort)
	}

	if sep, ok := sshSeparators[basessh]; ok {
		// inserts a separator between cli -options and host/cmd commands
		// example: $ ssh -p 12345 -- user@host.com git-lfs-authenticate ...
		args = append(args, sep, e.SshUserAndHost)
	} else {
		// no prefix supported, strip leading - off host to prevent cmd like:
		// $ git config lfs.url ssh://-proxycmd=whatever
		// $ plink -P 12345 -proxycmd=foo git-lfs-authenticate ...
		//
		// Instead, it'll attempt this, and eventually return an error
		// $ plink -P 12345 proxycmd=foo git-lfs-authenticate ...
		args = append(args, sshOptPrefixRE.ReplaceAllString(e.SshUserAndHost, ""))
	}

	return cmd, args, needShell
}

const defaultSSHCmd = "ssh"

var (
	sshOptPrefixRE = regexp.MustCompile(`\A\-+`)
	sshSeparators  = map[string]string{
		"ssh":          "--",
		"lfs-ssh-echo": "--", // used in lfs integration tests only
	}
)
<commit_msg>ssh: redirect stdin from /dev/null<commit_after>package lfshttp

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/git-lfs/git-lfs/config"
	"github.com/git-lfs/git-lfs/subprocess"
	"github.com/git-lfs/git-lfs/tools"
	"github.com/rubyist/tracerx"
)

type SSHResolver interface {
	Resolve(Endpoint, string) (sshAuthResponse, error)
}

func withSSHCache(ssh SSHResolver) SSHResolver {
	return &sshCache{
		endpoints: make(map[string]*sshAuthResponse),
		ssh:       ssh,
	}
}

type sshCache struct {
	endpoints map[string]*sshAuthResponse
	ssh       SSHResolver
}

func (c *sshCache) Resolve(e Endpoint, method string) (sshAuthResponse, error) {
	if len(e.SshUserAndHost) == 0 {
		return sshAuthResponse{}, nil
	}

	key := strings.Join([]string{e.SshUserAndHost, e.SshPort, e.SshPath, method}, "//")
	if res, ok := c.endpoints[key]; ok {
		if _, expired := res.IsExpiredWithin(5 * time.Second); !expired {
			tracerx.Printf("ssh cache: %s git-lfs-authenticate %s %s",
				e.SshUserAndHost, e.SshPath, endpointOperation(e, method))
			return *res, nil
		} else {
			tracerx.Printf("ssh cache expired: %s git-lfs-authenticate %s %s",
				e.SshUserAndHost, e.SshPath, endpointOperation(e, method))
		}
	}

	res, err := c.ssh.Resolve(e, method)
	if err == nil {
		c.endpoints[key] = &res
	}
	return res, err
}

type sshAuthResponse struct {
	Message   string            `json:"-"`
	Href      string            `json:"href"`
	Header    map[string]string `json:"header"`
	ExpiresAt time.Time         `json:"expires_at"`
	ExpiresIn int               `json:"expires_in"`

	createdAt time.Time
}

func (r *sshAuthResponse) IsExpiredWithin(d time.Duration) (time.Time, bool) {
	return tools.IsExpiredAtOrIn(r.createdAt, d, r.ExpiresAt,
		time.Duration(r.ExpiresIn)*time.Second)
}

type sshAuthClient struct {
	os  config.Environment
	git config.Environment
}

func (c *sshAuthClient) Resolve(e Endpoint, method string) (sshAuthResponse, error) {
	res := sshAuthResponse{}
	if len(e.SshUserAndHost) == 0 {
		return res, nil
	}

	exe, args := sshGetLFSExeAndArgs(c.os, c.git, e, method)
	cmd := exec.Command(exe, args...)

	// Save stdout and stderr in separate buffers
	var outbuf, errbuf bytes.Buffer
	cmd.Stdout = &outbuf
	cmd.Stderr = &errbuf

	now := time.Now()

	// Execute command
	err := cmd.Start()
	if err == nil {
		err = cmd.Wait()
	}

	// Processing result
	if err != nil {
		res.Message = strings.TrimSpace(errbuf.String())
	} else {
		err = json.Unmarshal(outbuf.Bytes(), &res)
		if res.ExpiresIn == 0 && res.ExpiresAt.IsZero() {
			ttl := c.git.Int("lfs.defaulttokenttl", 0)
			if ttl < 0 {
				ttl = 0
			}
			res.ExpiresIn = ttl
		}
		res.createdAt = now
	}

	return res, err
}

func sshFormatArgs(cmd string, args []string, needShell bool) (string, []string) {
	if !needShell {
		return cmd, args
	}

	return subprocess.FormatForShellQuotedArgs(cmd, args)
}

func sshGetLFSExeAndArgs(osEnv config.Environment, gitEnv config.Environment, e Endpoint, method string) (string, []string) {
	exe, args, needShell := sshGetExeAndArgs(osEnv, gitEnv, e)
	operation := endpointOperation(e, method)
	args = append(args, fmt.Sprintf("git-lfs-authenticate %s %s", e.SshPath, operation))
	exe, args = sshFormatArgs(exe, args, needShell)
	tracerx.Printf("run_command: %s %s", exe, strings.Join(args, " "))
	return exe, args
}

// Parse command, and if it looks like a valid command, return the ssh binary
// name, the command to run, and whether we need a shell.  If not, return
// existing as the ssh binary name.
func sshParseShellCommand(command string, existing string) (ssh string, cmd string, needShell bool) {
	ssh = existing
	if cmdArgs := tools.QuotedFields(command); len(cmdArgs) > 0 {
		needShell = true
		ssh = cmdArgs[0]
		cmd = command
	}
	return
}

// Return the executable name for ssh on this machine and the base args
// Base args includes port settings, user/host, everything pre the command to execute
func sshGetExeAndArgs(osEnv config.Environment, gitEnv config.Environment, e Endpoint) (exe string, baseargs []string, needShell bool) {
	var cmd string

	isPlink := false
	isTortoise := false

	ssh, _ := osEnv.Get("GIT_SSH")
	sshCmd, _ := osEnv.Get("GIT_SSH_COMMAND")
	ssh, cmd, needShell = sshParseShellCommand(sshCmd, ssh)

	if ssh == "" {
		sshCmd, _ := gitEnv.Get("core.sshcommand")
		ssh, cmd, needShell = sshParseShellCommand(sshCmd, defaultSSHCmd)
	}

	if cmd == "" {
		cmd = ssh
	}

	basessh := filepath.Base(ssh)

	if basessh != defaultSSHCmd {
		// Strip extension for easier comparison
		if ext := filepath.Ext(basessh); len(ext) > 0 {
			basessh = basessh[:len(basessh)-len(ext)]
		}
		isPlink = strings.EqualFold(basessh, "plink")
		isTortoise = strings.EqualFold(basessh, "tortoiseplink")
	}

	args := make([]string, 0, 7)

	if isTortoise {
		// TortoisePlink requires the -batch argument to behave like ssh/plink
		args = append(args, "-batch")
	}

	if len(e.SshPort) > 0 {
		if isPlink || isTortoise {
			args = append(args, "-P")
		} else {
			args = append(args, "-p")
		}
		args = append(args, e.SshPort)
	}

	if sep, ok := sshSeparators[basessh]; ok {
		// inserts a separator between cli -options and host/cmd commands
		// example: $ ssh -p 12345 -- user@host.com git-lfs-authenticate ...
		args = append(args, sep, e.SshUserAndHost)
	} else {
		// no prefix supported, strip leading - off host to prevent cmd like:
		// $ git config lfs.url ssh://-proxycmd=whatever
		// $ plink -P 12345 -proxycmd=foo git-lfs-authenticate ...
		//
		// Instead, it'll attempt this, and eventually return an error
		// $ plink -P 12345 proxycmd=foo git-lfs-authenticate ...
		args = append(args, sshOptPrefixRE.ReplaceAllString(e.SshUserAndHost, ""))
	}

	return cmd, args, needShell
}

const defaultSSHCmd = "ssh"

var (
	sshOptPrefixRE = regexp.MustCompile(`\A\-+`)
	sshSeparators  = map[string]string{
		"ssh":          "--",
		"lfs-ssh-echo": "--", // used in lfs integration tests only
	}
)
<|endoftext|>


<commit_before>// Copyright 2016 The Periph Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

// Specification
//
// Motorola never published a proper specification.
// http://electronics.stackexchange.com/questions/30096/spi-specifications
// http://www.nxp.com/files/microcontrollers/doc/data_sheet/M68HC11E.pdf page 120
// http://www.st.com/content/ccc/resource/technical/document/technical_note/58/17/ad/50/fa/c9/48/07/DM00054618.pdf/files/DM00054618.pdf/jcr:content/translations/en.DM00054618.pdf

package bitbang

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"periph.io/x/periph/conn"
	"periph.io/x/periph/conn/gpio"
	"periph.io/x/periph/conn/physic"
	"periph.io/x/periph/conn/spi"
	"periph.io/x/periph/host/cpu"
)

// NewSPI returns an spi.PortCloser that communicates SPI over 3 or 4 pins.
//
// BUG(maruel): Completely untested.
//
// cs can be nil.
func NewSPI(clk, mosi gpio.PinOut, miso gpio.PinIn, cs gpio.PinOut) (*SPI, error) {
	if err := clk.Out(gpio.High); err != nil {
		return nil, err
	}
	if err := mosi.Out(gpio.High); err != nil {
		return nil, err
	}
	if miso != nil {
		if err := miso.In(gpio.PullUp, gpio.NoEdge); err != nil {
			return nil, err
		}
	}
	if cs != nil {
		// Low means active.
		if err := cs.Out(gpio.High); err != nil {
			return nil, err
		}
	}
	return &SPI{spiConn: spiConn{sck: clk, sdi: miso, sdo: mosi, csn: cs}}, nil
}

// SPI represents a SPI master port implemented as bit-banging on 3 or 4 GPIO
// pins.
type SPI struct {
	spiConn spiConn
}

func (s *SPI) String() string {
	return fmt.Sprintf("bitbang/spi(%s, %s, %s, %s)", s.spiConn.sck, s.spiConn.sdi, s.spiConn.sdo, s.spiConn.csn)
}

// Close implements spi.PortCloser.
func (s *SPI) Close() error {
	return nil
}

// Connect implements spi.PortCloser.
func (s *SPI) Connect(f physic.Frequency, mode spi.Mode, bits int) (spi.Conn, error) {
	if f < 0 {
		return nil, errors.New("bitbang-spi: invalid frequency")
	}
	if mode&spi.HalfDuplex == spi.HalfDuplex {
		return nil, errors.New("bitbang-spi: half-duplex mode not supported")
	}
	if mode&spi.LSBFirst == spi.LSBFirst {
		return nil, errors.New("bitbang-spi: LSBFirst mode not supported")
	}
	if mode >= 0x20 {
		return nil, fmt.Errorf("bitbang-spi: unhandled mode %d(%s)", mode, mode.String())
	}
	s.spiConn.mu.Lock()
	defer s.spiConn.mu.Unlock()
	s.spiConn.freqDev = f
	if s.spiConn.freqDev != 0 && (s.spiConn.freqPort == 0 || s.spiConn.freqDev < s.spiConn.freqPort) {
		s.spiConn.halfCycle = f.Duration() / 2
	}
	s.spiConn.mode = mode
	s.spiConn.bits = bits
	s.spiConn.readAfterClockPulse = mode&spi.Mode1 == spi.Mode1

	// Set clock idle polarity, ensuring an idle clock to start
	s.spiConn.clockIdle = gpio.Level(mode&spi.Mode2 == spi.Mode2)
	if err := s.spiConn.sck.Out(s.spiConn.clockIdle); err != nil {
		return nil, fmt.Errorf("bitbang-spi: failed to idle clock: %v", err)
	}

	return &s.spiConn, nil
}

// LimitSpeed implements spi.PortCloser.
func (s *SPI) LimitSpeed(f physic.Frequency) error {
	if f <= 0 {
		return errors.New("bitbang-spi: invalid frequency")
	}
	s.spiConn.mu.Lock()
	defer s.spiConn.mu.Unlock()
	s.spiConn.freqPort = f
	if s.spiConn.freqDev == 0 || s.spiConn.freqPort < s.spiConn.freqDev {
		s.spiConn.halfCycle = f.Duration() / 2
	}
	return nil
}

// CLK implements spi.Pins.
func (s *SPI) CLK() gpio.PinOut {
	return s.spiConn.sck
}

// MOSI implements spi.Pins.
func (s *SPI) MOSI() gpio.PinOut {
	return s.spiConn.sdo
}

// MISO implements spi.Pins.
func (s *SPI) MISO() gpio.PinIn {
	return s.spiConn.sdi
}

// CS implements spi.Pins.
func (s *SPI) CS() gpio.PinOut {
	return s.spiConn.csn
}

//

// spiConn implements spi.Conn.
type spiConn struct {
	// Immutable.
	sck gpio.PinOut // Clock
	sdi gpio.PinIn  // MISO
	sdo gpio.PinOut // MOSI
	csn gpio.PinOut // CS

	// Mutable.
	mu                  sync.Mutex
	freqPort            physic.Frequency
	freqDev             physic.Frequency
	clockIdle           gpio.Level
	readAfterClockPulse bool
	mode                spi.Mode
	bits                int
	halfCycle           time.Duration
}

func (s *spiConn) String() string {
	return fmt.Sprintf("bitbang/spi(%s, %s, %s, %s)", s.sck, s.sdi, s.sdo, s.csn)
}

// Duplex implements spi.Conn.
func (s *spiConn) Duplex() conn.Duplex {
	// Maybe implement bitbanging SPI only in half mode?
	return conn.Full
}

// Tx implements spi.Conn.
//
// BUG(maruel): Implement mode (HalfDuplex and LSBFirst remain to be done).
// BUG(maruel): Implement bits.
// BUG(maruel): Test if read works.
func (s *spiConn) Tx(w, r []byte) (err error) {
	if len(r) != 0 && len(w) != len(r) {
		return errors.New("bitbang-spi: write and read buffers must be the same length")
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if err = s.assertCS(); err != nil {
		return fmt.Errorf("bitbang-spi: failed to assert chip-select: %v", err)
	}

	for i := uint(0); i < uint(len(w)*8); i++ {
		if err = s.sdo.Out(w[i/8]&(1<<(i%8)) != 0); err != nil {
			return fmt.Errorf("bitbang-spi: failed to send bit %d of word %d: %v", i%8, i/8, err)
		}

		s.sleepHalfCycle()
		if err = s.sck.Out(!s.clockIdle); err != nil {
			return fmt.Errorf("bitbang-spi: failed to assert clock: %v", err)
		}
		s.sleepHalfCycle()

		if s.readAfterClockPulse {
			if err = s.sck.Out(s.clockIdle); err != nil {
				return fmt.Errorf("bitbang-spi: failed to idle clock: %v", err)
			}
			s.sleepHalfCycle()
		}

		if len(r) != 0 {
			if s.sdi.Read() == gpio.High {
				r[i/8] |= 1 << (i % 8)
			}
		}

		if !s.readAfterClockPulse {
			if err = s.sck.Out(s.clockIdle); err != nil {
				return fmt.Errorf("bitbang-spi: failed to idle clock: %v", err)
			}
		}
	}
	if err = s.unassertCS(); err != nil {
		return fmt.Errorf("bitbang-spi: failed to unassert chip-select: %v", err)
	}

	return nil
}

// TxPackets implements spi.Conn.
func (s *spiConn) TxPackets(p []spi.Packet) error {
	return errors.New("bitbang-spi: not implemented")
}

// Write implements io.Writer.
func (s *spiConn) Write(d []byte) (int, error) {
	if err := s.Tx(d, nil); err != nil {
		return 0, err
	}
	return len(d), nil
}

// CLK implements spi.Pins.
func (s *spiConn) CLK() gpio.PinOut {
	return s.sck
}

// MOSI implements spi.Pins.
func (s *spiConn) MOSI() gpio.PinOut {
	return s.sdo
}

// MISO implements spi.Pins.
func (s *spiConn) MISO() gpio.PinIn {
	return s.sdi
}

// CS implements spi.Pins.
func (s *spiConn) CS() gpio.PinOut {
	return s.csn
}

//

// sleep does a busy loop to act as fast as possible.
func (s *spiConn) sleepHalfCycle() {
	cpu.Nanospin(s.halfCycle)
}

func (s *spiConn) assertCS() error {
	if s.csn == nil || s.mode&spi.NoCS == spi.NoCS {
		return nil
	}
	if err := s.csn.Out(gpio.Low); err != nil {
		return err
	}
	s.sleepHalfCycle()
	return nil
}

func (s *spiConn) unassertCS() error {
	if s.csn == nil || s.mode&spi.NoCS == spi.NoCS {
		return nil
	}
	if err := s.csn.Out(gpio.High); err != nil {
		return err
	}
	s.sleepHalfCycle()
	return nil
}

var _ spi.Conn = &spiConn{}
var _ spi.PortCloser = &SPI{}
<commit_msg>bitbang/spi: disable CS before manipulating CLK (#276)<commit_after>// Copyright 2016 The Periph Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

// Specification
//
// Motorola never published a proper specification.
// http://electronics.stackexchange.com/questions/30096/spi-specifications
// http://www.nxp.com/files/microcontrollers/doc/data_sheet/M68HC11E.pdf page 120
// http://www.st.com/content/ccc/resource/technical/document/technical_note/58/17/ad/50/fa/c9/48/07/DM00054618.pdf/files/DM00054618.pdf/jcr:content/translations/en.DM00054618.pdf

package bitbang

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"periph.io/x/periph/conn"
	"periph.io/x/periph/conn/gpio"
	"periph.io/x/periph/conn/physic"
	"periph.io/x/periph/conn/spi"
	"periph.io/x/periph/host/cpu"
)

// NewSPI returns an spi.PortCloser that communicates SPI over 3 or 4 pins.
//
// BUG(maruel): Completely untested.
//
// cs can be nil.
func NewSPI(clk, mosi gpio.PinOut, miso gpio.PinIn, cs gpio.PinOut) (*SPI, error) {
	if cs != nil {
		if err := cs.Out(gpio.High); err != nil {
			return nil, err
		}
	}
	if err := clk.Out(gpio.High); err != nil {
		return nil, err
	}
	if err := mosi.Out(gpio.High); err != nil {
		return nil, err
	}
	if miso != nil {
		if err := miso.In(gpio.PullUp, gpio.NoEdge); err != nil {
			return nil, err
		}
	}
	return &SPI{spiConn: spiConn{sck: clk, sdi: miso, sdo: mosi, csn: cs}}, nil
}

// SPI represents a SPI master port implemented as bit-banging on 3 or 4 GPIO
// pins.
type SPI struct {
	spiConn spiConn
}

func (s *SPI) String() string {
	return fmt.Sprintf("bitbang/spi(%s, %s, %s, %s)", s.spiConn.sck, s.spiConn.sdi, s.spiConn.sdo, s.spiConn.csn)
}

// Close implements spi.PortCloser.
func (s *SPI) Close() error {
	return nil
}

// Connect implements spi.PortCloser.
func (s *SPI) Connect(f physic.Frequency, mode spi.Mode, bits int) (spi.Conn, error) {
	if f < 0 {
		return nil, errors.New("bitbang-spi: invalid frequency")
	}
	if mode&spi.HalfDuplex == spi.HalfDuplex {
		return nil, errors.New("bitbang-spi: half-duplex mode not supported")
	}
	if mode&spi.LSBFirst == spi.LSBFirst {
		return nil, errors.New("bitbang-spi: LSBFirst mode not supported")
	}
	if mode >= 0x20 {
		return nil, fmt.Errorf("bitbang-spi: unhandled mode %d(%s)", mode, mode.String())
	}
	s.spiConn.mu.Lock()
	defer s.spiConn.mu.Unlock()
	s.spiConn.freqDev = f
	if s.spiConn.freqDev != 0 && (s.spiConn.freqPort == 0 || s.spiConn.freqDev < s.spiConn.freqPort) {
		s.spiConn.halfCycle = f.Duration() / 2
	}
	s.spiConn.mode = mode
	s.spiConn.bits = bits
	s.spiConn.readAfterClockPulse = mode&spi.Mode1 == spi.Mode1

	// Set clock idle polarity, ensuring an idle clock to start
	s.spiConn.clockIdle = gpio.Level(mode&spi.Mode2 == spi.Mode2)
	if err := s.spiConn.sck.Out(s.spiConn.clockIdle); err != nil {
		return nil, fmt.Errorf("bitbang-spi: failed to idle clock: %v", err)
	}

	return &s.spiConn, nil
}

// LimitSpeed implements spi.PortCloser.
func (s *SPI) LimitSpeed(f physic.Frequency) error {
	if f <= 0 {
		return errors.New("bitbang-spi: invalid frequency")
	}
	s.spiConn.mu.Lock()
	defer s.spiConn.mu.Unlock()
	s.spiConn.freqPort = f
	if s.spiConn.freqDev == 0 || s.spiConn.freqPort < s.spiConn.freqDev {
		s.spiConn.halfCycle = f.Duration() / 2
	}
	return nil
}

// CLK implements spi.Pins.
func (s *SPI) CLK() gpio.PinOut {
	return s.spiConn.sck
}

// MOSI implements spi.Pins.
func (s *SPI) MOSI() gpio.PinOut {
	return s.spiConn.sdo
}

// MISO implements spi.Pins.
func (s *SPI) MISO() gpio.PinIn {
	return s.spiConn.sdi
}

// CS implements spi.Pins.
func (s *SPI) CS() gpio.PinOut {
	return s.spiConn.csn
}

//

// spiConn implements spi.Conn.
type spiConn struct {
	// Immutable.
	sck gpio.PinOut // Clock
	sdi gpio.PinIn  // MISO
	sdo gpio.PinOut // MOSI
	csn gpio.PinOut // CS

	// Mutable.
	mu                  sync.Mutex
	freqPort            physic.Frequency
	freqDev             physic.Frequency
	clockIdle           gpio.Level
	readAfterClockPulse bool
	mode                spi.Mode
	bits                int
	halfCycle           time.Duration
}

func (s *spiConn) String() string {
	return fmt.Sprintf("bitbang/spi(%s, %s, %s, %s)", s.sck, s.sdi, s.sdo, s.csn)
}

// Duplex implements spi.Conn.
func (s *spiConn) Duplex() conn.Duplex {
	// Maybe implement bitbanging SPI only in half mode?
	return conn.Full
}

// Tx implements spi.Conn.
//
// BUG(maruel): Implement mode (HalfDuplex and LSBFirst remain to be done).
// BUG(maruel): Implement bits.
// BUG(maruel): Test if read works.
func (s *spiConn) Tx(w, r []byte) (err error) {
	if len(r) != 0 && len(w) != len(r) {
		return errors.New("bitbang-spi: write and read buffers must be the same length")
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if err = s.assertCS(); err != nil {
		return fmt.Errorf("bitbang-spi: failed to assert chip-select: %v", err)
	}

	for i := uint(0); i < uint(len(w)*8); i++ {
		if err = s.sdo.Out(w[i/8]&(1<<(i%8)) != 0); err != nil {
			return fmt.Errorf("bitbang-spi: failed to send bit %d of word %d: %v", i%8, i/8, err)
		}

		s.sleepHalfCycle()
		if err = s.sck.Out(!s.clockIdle); err != nil {
			return fmt.Errorf("bitbang-spi: failed to assert clock: %v", err)
		}
		s.sleepHalfCycle()

		if s.readAfterClockPulse {
			if err = s.sck.Out(s.clockIdle); err != nil {
				return fmt.Errorf("bitbang-spi: failed to idle clock: %v", err)
			}
			s.sleepHalfCycle()
		}

		if len(r) != 0 {
			if s.sdi.Read() == gpio.High {
				r[i/8] |= 1 << (i % 8)
			}
		}

		if !s.readAfterClockPulse {
			if err = s.sck.Out(s.clockIdle); err != nil {
				return fmt.Errorf("bitbang-spi: failed to idle clock: %v", err)
			}
		}
	}
	if err = s.unassertCS(); err != nil {
		return fmt.Errorf("bitbang-spi: failed to unassert chip-select: %v", err)
	}

	return nil
}

// TxPackets implements spi.Conn.
func (s *spiConn) TxPackets(p []spi.Packet) error {
	return errors.New("bitbang-spi: not implemented")
}

// Write implements io.Writer.
func (s *spiConn) Write(d []byte) (int, error) {
	if err := s.Tx(d, nil); err != nil {
		return 0, err
	}
	return len(d), nil
}

// CLK implements spi.Pins.
func (s *spiConn) CLK() gpio.PinOut {
	return s.sck
}

// MOSI implements spi.Pins.
func (s *spiConn) MOSI() gpio.PinOut {
	return s.sdo
}

// MISO implements spi.Pins.
func (s *spiConn) MISO() gpio.PinIn {
	return s.sdi
}

// CS implements spi.Pins.
func (s *spiConn) CS() gpio.PinOut {
	return s.csn
}

//

// sleep does a busy loop to act as fast as possible.
func (s *spiConn) sleepHalfCycle() {
	cpu.Nanospin(s.halfCycle)
}

func (s *spiConn) assertCS() error {
	if s.csn == nil || s.mode&spi.NoCS == spi.NoCS {
		return nil
	}
	if err := s.csn.Out(gpio.Low); err != nil {
		return err
	}
	s.sleepHalfCycle()
	return nil
}

func (s *spiConn) unassertCS() error {
	if s.csn == nil || s.mode&spi.NoCS == spi.NoCS {
		return nil
	}
	if err := s.csn.Out(gpio.High); err != nil {
		return err
	}
	s.sleepHalfCycle()
	return nil
}

var _ spi.Conn = &spiConn{}
var _ spi.PortCloser = &SPI{}
<|endoftext|>