<commit_before>package charts

import (
	"github.com/go-echarts/go-echarts/v2/opts"
)

type Overlaper interface {
	overlap() MultiSeries
}

// XYAxis represent the X and Y axis in the rectangular coordinate.
type XYAxis struct {
	XAxisList []opts.XAxis `json:"xaxis"`
	YAxisList []opts.YAxis `json:"yaxis"`
}

func (xy *XYAxis) initXYAxis() {
	xy.XAxisList = append(xy.XAxisList, opts.XAxis{})
	xy.YAxisList = append(xy.YAxisList, opts.YAxis{})
}

// ExtendXAxis adds new X axes.
func (xy *XYAxis) ExtendXAxis(xAxis ...opts.XAxis) {
	xy.XAxisList = append(xy.XAxisList, xAxis...)
}

// ExtendYAxis adds new Y axes.
func (xy *XYAxis) ExtendYAxis(yAxis ...opts.YAxis) {
	xy.YAxisList = append(xy.YAxisList, yAxis...)
}

// WithXAxisOpts
func WithXAxisOpts(opt opts.XAxis, index ...int) GlobalOpts {
	return func(bc *BaseConfiguration) {
		if len(index) == 0 {
			index = []int{0}
		}
		for i := 0; i < len(index); i++ {
			bc.XYAxis.XAxisList[index[i]] = opt
		}
	}
}

// WithYAxisOpts
func WithYAxisOpts(opt opts.YAxis, index ...int) GlobalOpts {
	return func(bc *BaseConfiguration) {
		if len(index) == 0 {
			index = []int{0}
		}
		for i := 0; i < len(index); i++ {
			bc.XYAxis.YAxisList[index[i]] = opt
		}
	}
}

// RectConfiguration contains options for the rectangular coordinate.
type RectConfiguration struct {
	BaseConfiguration
}

func (rect *RectConfiguration) setRectGlobalOptions(options ...GlobalOpts) {
	rect.BaseConfiguration.setBaseGlobalOptions(options...)
}

// RectChart is a chart in RectChart coordinate.
type RectChart struct {
	RectConfiguration
	MultiSeries

	xAxisData interface{}
}

func (rc *RectChart) overlap() MultiSeries {
	return rc.MultiSeries
}

// SetGlobalOptions sets options for the RectChart instance.
func (rc *RectChart) SetGlobalOptions(options ...GlobalOpts) *RectChart {
	rc.RectConfiguration.setRectGlobalOptions(options...)
	return rc
}

// Overlap composites multiple charts into one single canvas.
// 结合不同类型图表叠加画在同张图上
// 只适用于 RectChart 图表，其实现了 rectCharter 接口
// RectChart 图表包括 Bar/BoxPlot/Line/Scatter/EffectScatter/Kline/HeatMap
// 将 RectChart 图表的 Series 追加到调用者的 Series 里面，Series 是完全独立的
// 而全局配置使用的是调用者的配置项
func (rc *RectChart) Overlap(a ...Overlaper) {
	for i := 0; i < len(a); i++ {
		rc.MultiSeries = append(rc.MultiSeries, a[i].overlap()...)
	}
}

// Validate
func (rc *RectChart) Validate() {
	// 确保 X 轴数据不会因为设置了 XAxisOpts 而被抹除
	rc.XAxisList[0].Data = rc.xAxisData
	// 确保 Y 轴数标签正确显示
	for i := 0; i < len(rc.YAxisList); i++ {
		rc.YAxisList[i].AxisLabel.Show = true
	}
	rc.Assets.Validate(rc.AssetsHost)
}
<commit_msg>[update] update comments. (#116)<commit_after>package charts

import (
	"github.com/go-echarts/go-echarts/v2/opts"
)

type Overlaper interface {
	overlap() MultiSeries
}

// XYAxis represent the X and Y axis in the rectangular coordinate.
type XYAxis struct {
	XAxisList []opts.XAxis `json:"xaxis"`
	YAxisList []opts.YAxis `json:"yaxis"`
}

func (xy *XYAxis) initXYAxis() {
	xy.XAxisList = append(xy.XAxisList, opts.XAxis{})
	xy.YAxisList = append(xy.YAxisList, opts.YAxis{})
}

// ExtendXAxis adds new X axes.
func (xy *XYAxis) ExtendXAxis(xAxis ...opts.XAxis) {
	xy.XAxisList = append(xy.XAxisList, xAxis...)
}

// ExtendYAxis adds new Y axes.
func (xy *XYAxis) ExtendYAxis(yAxis ...opts.YAxis) {
	xy.YAxisList = append(xy.YAxisList, yAxis...)
}

// WithXAxisOpts
func WithXAxisOpts(opt opts.XAxis, index ...int) GlobalOpts {
	return func(bc *BaseConfiguration) {
		if len(index) == 0 {
			index = []int{0}
		}
		for i := 0; i < len(index); i++ {
			bc.XYAxis.XAxisList[index[i]] = opt
		}
	}
}

// WithYAxisOpts
func WithYAxisOpts(opt opts.YAxis, index ...int) GlobalOpts {
	return func(bc *BaseConfiguration) {
		if len(index) == 0 {
			index = []int{0}
		}
		for i := 0; i < len(index); i++ {
			bc.XYAxis.YAxisList[index[i]] = opt
		}
	}
}

// RectConfiguration contains options for the rectangular coordinate.
type RectConfiguration struct {
	BaseConfiguration
}

func (rect *RectConfiguration) setRectGlobalOptions(options ...GlobalOpts) {
	rect.BaseConfiguration.setBaseGlobalOptions(options...)
}

// RectChart is a chart in RectChart coordinate.
type RectChart struct {
	RectConfiguration
	MultiSeries

	xAxisData interface{}
}

func (rc *RectChart) overlap() MultiSeries {
	return rc.MultiSeries
}

// SetGlobalOptions sets options for the RectChart instance.
func (rc *RectChart) SetGlobalOptions(options ...GlobalOpts) *RectChart {
	rc.RectConfiguration.setRectGlobalOptions(options...)
	return rc
}

// Overlap composites multiple charts into one single canvas.
// It is only suit for some of the charts which are in rectangular coordinate.
// Support charts: Bar/BoxPlot/Line/Scatter/EffectScatter/Kline/HeatMap
func (rc *RectChart) Overlap(a ...Overlaper) {
	for i := 0; i < len(a); i++ {
		rc.MultiSeries = append(rc.MultiSeries, a[i].overlap()...)
	}
}

// Validate
func (rc *RectChart) Validate() {
	// Make sure that the data of X axis won't be cleaned for XAxisOpts
	rc.XAxisList[0].Data = rc.xAxisData
	// Make sure that the labels of Y axis show correctly
	for i := 0; i < len(rc.YAxisList); i++ {
		rc.YAxisList[i].AxisLabel.Show = true
	}
	rc.Assets.Validate(rc.AssetsHost)
}
<|endoftext|>


<commit_before>package check

import (
	"testing"
	"github.com/foomo/petze/config"
)

var checkMinMaxCountTestCases = []struct {
	expect  config.Expect
	length  int64
	ok      bool
	info    string
	message string
}{
	{expect: config.Expect{Min: &[]int64{3}[0]}, length: int64(4), ok: true, info: "", message: "check minimum true"},
	{expect: config.Expect{Min: &[]int64{5}[0]}, length: int64(4), ok: true, info: "min actual: 4 < expected: 5", message: "check minimum false"},
	{expect: config.Expect{Max: &[]int64{5}[0]}, length: int64(4), ok: true, info: "", message: "check max true"},
	{expect: config.Expect{Max: &[]int64{3}[0]}, length: int64(4), ok: true, info: "max actual: 4 > expected: 3", message: "check max false"},
	{expect: config.Expect{Count: &[]int64{3}[0]}, length: int64(3), ok: true, info: "", message: "check count true"},
	{expect: config.Expect{Count: &[]int64{3}[0]}, length: int64(4), ok: true, info: "count actual: 4 != expected: 3", message: "check count false"},
}

func TestCheckMinMaxCount(t *testing.T) {
	for _, test := range checkMinMaxCountTestCases {
		ok, info := checkMinMaxCount(test.expect, test.length)
		if ok != test.ok && info != test.info {
			t.Error(test.message)
		}
	}
}
<commit_msg>Add full test coverage for check method<commit_after>package check

import (
	"testing"
	"github.com/foomo/petze/config"
)

var checkMinMaxCountTestCases = []struct {
	expect  config.Expect
	length  int64
	ok      bool
	info    string
	message string
}{
	{expect: config.Expect{Min: &[]int64{3}[0]}, length: 4, ok: true, info: "", message: "check minimum true"},
	{expect: config.Expect{Min: &[]int64{5}[0]}, length: 4, ok: true, info: "min actual: 4 < expected: 5", message: "check minimum false"},
	{expect: config.Expect{Max: &[]int64{5}[0]}, length: 4, ok: true, info: "", message: "check max true"},
	{expect: config.Expect{Max: &[]int64{3}[0]}, length: 4, ok: true, info: "max actual: 4 > expected: 3", message: "check max false"},
	{expect: config.Expect{Count: &[]int64{3}[0]}, length: 3, ok: true, info: "", message: "check count true"},
	{expect: config.Expect{Count: &[]int64{3}[0]}, length: 4, ok: true, info: "count actual: 4 != expected: 3", message: "check count false"},
}

func TestCheckMinMaxCount(t *testing.T) {
	for _, test := range checkMinMaxCountTestCases {
		ok, info := checkMinMaxCount(test.expect, test.length)
		if ok != test.ok && info != test.info {
			t.Error(test.message)
		}
	}
}

func TestCheckMinMaxPanic(t *testing.T) {
	defer func() {
		if r := recover(); r == nil {
			t.Error("should throw a panic")
		}
	}()

	checkMinMaxCount(config.Expect{Equals: ""}, 3)
}
<|endoftext|>


<commit_before>package check

import (
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"time"

	logr "github.com/Sirupsen/logrus"
)

// ParseCommandLineArguments parses the command line arguments.
func (check *HealthCheck) ParseCommandLineArguments() {
	flag.UintVar(&check.config.brokerID, "broker-id", 0, "id of the Kafka broker to health check")
	flag.UintVar(&check.config.brokerPort, "broker-port", 9092, "Kafka broker port")
	flag.UintVar(&check.config.statusServerPort, "server-port", 8000, "port to open for http health status queries")
	flag.StringVar(&check.config.zookeeperConnect, "zookeeper", "", "ZooKeeper connect string (e.g. node1:2181,node2:2181,.../chroot)")
	flag.StringVar(&check.config.topicName, "topic", "", "name of the topic to use - use one per broker, defaults to broker-<id>-health-check")
	flag.StringVar(&check.config.replicationTopicName, "replication-topic", "",
		"name of the topic to use for replication checks - use one per cluster, defaults to broker-replication-check")
	flag.UintVar(&check.config.replicationFailureThreshold, "replication-failures-count", 5,
		"number of replication failures before broker is reported unhealthy")
	flag.DurationVar(&check.config.CheckInterval, "check-interval", 10*time.Second, "how frequently to perform health checks")
	flag.BoolVar(&check.config.NoTopicCreation, "no-topic-creation", false, "disable automatic topic creation and deletion")
	flag.Parse()
	l := log.New(os.Stderr, "", 0)
	valid := check.validateConfig(l)
	if !valid {
		l.Printf("%s usage:\n", os.Args[0])
		flag.PrintDefaults()
		l.Fatal("One or more mandatory command line parameters are missing.")
	} else {
		if check.config.topicName == "" {
			check.config.topicName = fmt.Sprintf("broker-%d-health-check", check.config.brokerID)
			logr.Println("using topic", check.config.topicName, "for broker", check.config.brokerID, "health check")
		}
		if check.config.replicationTopicName == "" {
			check.config.replicationTopicName = "broker-replication-check"
			logr.Println("using topic", check.config.topicName, "for broker", check.config.brokerID, "replication check")
		}
		check.config.retryInterval = check.config.CheckInterval / 2
	}
}

func (check *HealthCheck) validateConfig(l *log.Logger) bool {
	valid := true
	if check.config.zookeeperConnect == "" {
		l.Println("parameter -zookeeper required.")
		valid = false
	}
	if check.config.topicName != "" && !validTopicName(check.config.topicName) {
		l.Println("topic name", check.config.topicName, "is not a valid Kafka topic name.")
		valid = false
	}

	return valid
}

// from https://github.com/apache/kafka/blob/6eacc0de303e4d29e083b89c1f53615c1dfa291e/core/src/main/scala/kafka/common/Topic.scala
var legalTopicPattern = regexp.MustCompile("^[a-zA-Z0-9\\._\\-]+$")

func validTopicName(name string) bool {
	return legalTopicPattern.MatchString(name) && len(name) <= 255
}
<commit_msg>Enable specification of broker host from CLI<commit_after>package check

import (
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"time"

	logr "github.com/Sirupsen/logrus"
)

// ParseCommandLineArguments parses the command line arguments.
func (check *HealthCheck) ParseCommandLineArguments() {
	flag.StringVar(&check.config.brokerHost, "broker-host", "localhost", "ip address or hostname of broker host")
	flag.UintVar(&check.config.brokerID, "broker-id", 0, "id of the Kafka broker to health check")
	flag.UintVar(&check.config.brokerPort, "broker-port", 9092, "Kafka broker port")
	flag.UintVar(&check.config.statusServerPort, "server-port", 8000, "port to open for http health status queries")
	flag.StringVar(&check.config.zookeeperConnect, "zookeeper", "", "ZooKeeper connect string (e.g. node1:2181,node2:2181,.../chroot)")
	flag.StringVar(&check.config.topicName, "topic", "", "name of the topic to use - use one per broker, defaults to broker-<id>-health-check")
	flag.StringVar(&check.config.replicationTopicName, "replication-topic", "",
		"name of the topic to use for replication checks - use one per cluster, defaults to broker-replication-check")
	flag.UintVar(&check.config.replicationFailureThreshold, "replication-failures-count", 5,
		"number of replication failures before broker is reported unhealthy")
	flag.DurationVar(&check.config.CheckInterval, "check-interval", 10*time.Second, "how frequently to perform health checks")
	flag.BoolVar(&check.config.NoTopicCreation, "no-topic-creation", false, "disable automatic topic creation and deletion")
	flag.Parse()
	l := log.New(os.Stderr, "", 0)
	valid := check.validateConfig(l)
	if !valid {
		l.Printf("%s usage:\n", os.Args[0])
		flag.PrintDefaults()
		l.Fatal("One or more mandatory command line parameters are missing.")
	} else {
		if check.config.topicName == "" {
			check.config.topicName = fmt.Sprintf("broker-%d-health-check", check.config.brokerID)
			logr.Println("using topic", check.config.topicName, "for broker", check.config.brokerID, "health check")
		}
		if check.config.replicationTopicName == "" {
			check.config.replicationTopicName = "broker-replication-check"
			logr.Println("using topic", check.config.topicName, "for broker", check.config.brokerID, "replication check")
		}
		check.config.retryInterval = check.config.CheckInterval / 2
	}
}

func (check *HealthCheck) validateConfig(l *log.Logger) bool {
	valid := true
	if check.config.zookeeperConnect == "" {
		l.Println("parameter -zookeeper required.")
		valid = false
	}
	if check.config.topicName != "" && !validTopicName(check.config.topicName) {
		l.Println("topic name", check.config.topicName, "is not a valid Kafka topic name.")
		valid = false
	}

	return valid
}

// from https://github.com/apache/kafka/blob/6eacc0de303e4d29e083b89c1f53615c1dfa291e/core/src/main/scala/kafka/common/Topic.scala
var legalTopicPattern = regexp.MustCompile("^[a-zA-Z0-9\\._\\-]+$")

func validTopicName(name string) bool {
	return legalTopicPattern.MatchString(name) && len(name) <= 255
}
<|endoftext|>


<commit_before>package coordinator

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	log "code.google.com/p/log4go"
	"github.com/influxdb/influxdb/cluster"
	"github.com/influxdb/influxdb/protocol"
)

type ProtobufClient struct {
	connLock          sync.Mutex
	conn              net.Conn
	hostAndPort       string
	requestBufferLock sync.RWMutex
	requestBuffer     map[uint32]*runningRequest
	reconnectWait     sync.WaitGroup
	connectCalled     bool
	lastRequestId     uint32
	writeTimeout      time.Duration
	attempts          int
	stopped           bool
	reconChan         chan struct{}
	reconGroup        *sync.WaitGroup
	once              *sync.Once
}

type runningRequest struct {
	timeMade time.Time
	r        cluster.ResponseChannel
	request  *protocol.Request
}

const (
	REQUEST_RETRY_ATTEMPTS = 2
	MAX_RESPONSE_SIZE      = MAX_REQUEST_SIZE
	MAX_REQUEST_TIME       = time.Second * 1200
	RECONNECT_RETRY_WAIT   = time.Millisecond * 100
)

func NewProtobufClient(hostAndPort string, writeTimeout time.Duration) *ProtobufClient {
	log.Debug("NewProtobufClient: ", hostAndPort)
	return &ProtobufClient{
		hostAndPort:   hostAndPort,
		requestBuffer: make(map[uint32]*runningRequest),
		writeTimeout:  writeTimeout,
		reconChan:     make(chan struct{}, 1),
		reconGroup:    new(sync.WaitGroup),
		once:          new(sync.Once),
		stopped:       false,
	}
}

func (self *ProtobufClient) Connect() {
	self.once.Do(self.connect)
}

func (self *ProtobufClient) connect() {
	self.reconChan <- struct{}{}
	go func() {
		self.reconnect()
		self.readResponses()
	}()
	go self.peridicallySweepTimedOutRequests()
}

func (self *ProtobufClient) Close() {
	self.connLock.Lock()
	defer self.connLock.Unlock()
	if self.conn != nil {
		self.conn.Close()
		self.stopped = true
		self.conn = nil
	}
	self.ClearRequests()
}

func (self *ProtobufClient) getConnection() net.Conn {
	self.connLock.Lock()
	defer self.connLock.Unlock()
	return self.conn
}

func (self *ProtobufClient) ClearRequests() {
	self.requestBufferLock.Lock()
	defer self.requestBufferLock.Unlock()

	for _, req := range self.requestBuffer {
		self.cancelRequest(req.request)
	}

	self.requestBuffer = map[uint32]*runningRequest{}
}

func (self *ProtobufClient) CancelRequest(request *protocol.Request) {
	self.requestBufferLock.Lock()
	defer self.requestBufferLock.Unlock()
	self.cancelRequest(request)
}

func (self *ProtobufClient) cancelRequest(request *protocol.Request) {
	req, ok := self.requestBuffer[*request.Id]
	if !ok {
		return
	}
	message := "cancelling request"
	req.r.Yield(&protocol.Response{
		Type:         protocol.Response_ERROR.Enum(),
		ErrorMessage: &message,
	})
	delete(self.requestBuffer, *request.Id)
}

// Makes a request to the server. If the responseStream chan is not nil it will expect a response from the server
// with a matching request.Id. The REQUEST_RETRY_ATTEMPTS constant of 3 and the RECONNECT_RETRY_WAIT of 100ms means
// that an attempt to make a request to a downed server will take 300ms to time out.
func (self *ProtobufClient) MakeRequest(request *protocol.Request, r cluster.ResponseChannel) error {
	if request.Id == nil {
		id := atomic.AddUint32(&self.lastRequestId, uint32(1))
		request.Id = &id
	}
	if r != nil {
		self.requestBufferLock.Lock()

		// this should actually never happen. The sweeper should clear out dead requests
		// before the uint32 ids roll over.
		if oldReq, alreadyHasRequestById := self.requestBuffer[*request.Id]; alreadyHasRequestById {
			message := "already has a request with this id, must have timed out"
			log.Error(message)
			oldReq.r.Yield(&protocol.Response{
				Type:         protocol.Response_ERROR.Enum(),
				ErrorMessage: &message,
			})
		}
		self.requestBuffer[*request.Id] = &runningRequest{timeMade: time.Now(), r: r, request: request}
		self.requestBufferLock.Unlock()
	}

	data, err := request.Encode()
	if err != nil {
		return err
	}

	conn := self.getConnection()
	if conn == nil {
		conn = self.reconnect()
		if conn == nil {
			return fmt.Errorf("Failed to connect to server %s", self.hostAndPort)
		}
	}

	if self.writeTimeout > 0 {
		conn.SetWriteDeadline(time.Now().Add(self.writeTimeout))
	}
	buff := bytes.NewBuffer(make([]byte, 0, len(data)+8))
	binary.Write(buff, binary.LittleEndian, uint32(len(data)))
	_, err = conn.Write(append(buff.Bytes(), data...))

	if err == nil {
		return nil
	}

	// if we got here it errored out, clear out the request
	self.requestBufferLock.Lock()
	delete(self.requestBuffer, *request.Id)
	self.requestBufferLock.Unlock()
	self.reconnect()
	return err
}

func (self *ProtobufClient) readResponses() {
	message := make([]byte, 0, MAX_RESPONSE_SIZE)
	buff := bytes.NewBuffer(message)
	for !self.stopped {
		buff.Reset()
		conn := self.getConnection()
		if conn == nil {
			time.Sleep(200 * time.Millisecond)
			continue
		}
		var messageSizeU uint32
		var err error
		err = binary.Read(conn, binary.LittleEndian, &messageSizeU)
		if err != nil {
			log.Error("Error while reading messsage size: %d", err)
			time.Sleep(200 * time.Millisecond)
			continue
		}
		messageSize := int64(messageSizeU)
		messageReader := io.LimitReader(conn, messageSize)
		_, err = io.Copy(buff, messageReader)
		if err != nil {
			log.Error("Error while reading message: %d", err)
			time.Sleep(200 * time.Millisecond)
			continue
		}
		response, err := protocol.DecodeResponse(buff)
		if err != nil {
			log.Error("error unmarshaling response: %s", err)
			time.Sleep(200 * time.Millisecond)
		} else {
			self.sendResponse(response)
		}
	}
}

func (self *ProtobufClient) sendResponse(response *protocol.Response) {
	self.requestBufferLock.RLock()
	req, ok := self.requestBuffer[*response.RequestId]
	self.requestBufferLock.RUnlock()
	if !ok {
		return
	}

	deleteRequest := false
	switch rt := response.GetType(); rt {
	case protocol.Response_END_STREAM,
		protocol.Response_HEARTBEAT,
		protocol.Response_ERROR:
		deleteRequest = true
	case protocol.Response_QUERY:
		// do nothing
	default:
		panic(fmt.Errorf("Unknown response type: %s", rt))
	}

	self.requestBufferLock.Lock()
	req, ok = self.requestBuffer[*response.RequestId]
	if deleteRequest {
		delete(self.requestBuffer, *response.RequestId)
	}
	self.requestBufferLock.Unlock()
	if !ok {
		return
	}

	log.Debug("ProtobufClient yielding to %s %s", req.r.Name(), response)
	req.r.Yield(response)
}

func (self *ProtobufClient) reconnect() net.Conn {
	select {
	case <-self.reconChan:
		self.reconGroup.Add(1)
		defer func() {
			self.reconGroup.Done()
			self.reconChan <- struct{}{}
		}()
	default:
		self.reconGroup.Wait()
		return self.conn
	}

	if self.conn != nil {
		self.conn.Close()
	}
	conn, err := net.DialTimeout("tcp", self.hostAndPort, self.writeTimeout)
	if err != nil {
		self.attempts++
		if self.attempts < 100 {
			return nil
		}

		log.Error("failed to connect to %s %d times", self.hostAndPort, self.attempts)
		self.attempts = 0
	}

	self.conn = conn
	log.Info("connected to %s", self.hostAndPort)
	return conn
}

func (self *ProtobufClient) peridicallySweepTimedOutRequests() {
	for {
		time.Sleep(time.Minute)
		self.requestBufferLock.Lock()
		maxAge := time.Now().Add(-MAX_REQUEST_TIME)
		for k, req := range self.requestBuffer {
			if req.timeMade.Before(maxAge) {
				delete(self.requestBuffer, k)
				log.Warn("Request timed out: ", req.request)
			}
		}
		self.requestBufferLock.Unlock()
	}
}
<commit_msg>Write to the buffer instead of using append<commit_after>package coordinator

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	log "code.google.com/p/log4go"
	"github.com/influxdb/influxdb/cluster"
	"github.com/influxdb/influxdb/protocol"
)

type ProtobufClient struct {
	connLock          sync.Mutex
	conn              net.Conn
	hostAndPort       string
	requestBufferLock sync.RWMutex
	requestBuffer     map[uint32]*runningRequest
	reconnectWait     sync.WaitGroup
	connectCalled     bool
	lastRequestId     uint32
	writeTimeout      time.Duration
	attempts          int
	stopped           bool
	reconChan         chan struct{}
	reconGroup        *sync.WaitGroup
	once              *sync.Once
}

type runningRequest struct {
	timeMade time.Time
	r        cluster.ResponseChannel
	request  *protocol.Request
}

const (
	REQUEST_RETRY_ATTEMPTS = 2
	MAX_RESPONSE_SIZE      = MAX_REQUEST_SIZE
	MAX_REQUEST_TIME       = time.Second * 1200
	RECONNECT_RETRY_WAIT   = time.Millisecond * 100
)

func NewProtobufClient(hostAndPort string, writeTimeout time.Duration) *ProtobufClient {
	log.Debug("NewProtobufClient: ", hostAndPort)
	return &ProtobufClient{
		hostAndPort:   hostAndPort,
		requestBuffer: make(map[uint32]*runningRequest),
		writeTimeout:  writeTimeout,
		reconChan:     make(chan struct{}, 1),
		reconGroup:    new(sync.WaitGroup),
		once:          new(sync.Once),
		stopped:       false,
	}
}

func (self *ProtobufClient) Connect() {
	self.once.Do(self.connect)
}

func (self *ProtobufClient) connect() {
	self.reconChan <- struct{}{}
	go func() {
		self.reconnect()
		self.readResponses()
	}()
	go self.peridicallySweepTimedOutRequests()
}

func (self *ProtobufClient) Close() {
	self.connLock.Lock()
	defer self.connLock.Unlock()
	if self.conn != nil {
		self.conn.Close()
		self.stopped = true
		self.conn = nil
	}
	self.ClearRequests()
}

func (self *ProtobufClient) getConnection() net.Conn {
	self.connLock.Lock()
	defer self.connLock.Unlock()
	return self.conn
}

func (self *ProtobufClient) ClearRequests() {
	self.requestBufferLock.Lock()
	defer self.requestBufferLock.Unlock()

	for _, req := range self.requestBuffer {
		self.cancelRequest(req.request)
	}

	self.requestBuffer = map[uint32]*runningRequest{}
}

func (self *ProtobufClient) CancelRequest(request *protocol.Request) {
	self.requestBufferLock.Lock()
	defer self.requestBufferLock.Unlock()
	self.cancelRequest(request)
}

func (self *ProtobufClient) cancelRequest(request *protocol.Request) {
	req, ok := self.requestBuffer[*request.Id]
	if !ok {
		return
	}
	message := "cancelling request"
	req.r.Yield(&protocol.Response{
		Type:         protocol.Response_ERROR.Enum(),
		ErrorMessage: &message,
	})
	delete(self.requestBuffer, *request.Id)
}

// Makes a request to the server. If the responseStream chan is not nil it will expect a response from the server
// with a matching request.Id. The REQUEST_RETRY_ATTEMPTS constant of 3 and the RECONNECT_RETRY_WAIT of 100ms means
// that an attempt to make a request to a downed server will take 300ms to time out.
func (self *ProtobufClient) MakeRequest(request *protocol.Request, r cluster.ResponseChannel) error {
	if request.Id == nil {
		id := atomic.AddUint32(&self.lastRequestId, uint32(1))
		request.Id = &id
	}
	if r != nil {
		self.requestBufferLock.Lock()

		// this should actually never happen. The sweeper should clear out dead requests
		// before the uint32 ids roll over.
		if oldReq, alreadyHasRequestById := self.requestBuffer[*request.Id]; alreadyHasRequestById {
			message := "already has a request with this id, must have timed out"
			log.Error(message)
			oldReq.r.Yield(&protocol.Response{
				Type:         protocol.Response_ERROR.Enum(),
				ErrorMessage: &message,
			})
		}
		self.requestBuffer[*request.Id] = &runningRequest{timeMade: time.Now(), r: r, request: request}
		self.requestBufferLock.Unlock()
	}

	data, err := request.Encode()
	if err != nil {
		return err
	}

	conn := self.getConnection()
	if conn == nil {
		conn = self.reconnect()
		if conn == nil {
			return fmt.Errorf("Failed to connect to server %s", self.hostAndPort)
		}
	}

	if self.writeTimeout > 0 {
		conn.SetWriteDeadline(time.Now().Add(self.writeTimeout))
	}
	buff := bytes.NewBuffer(make([]byte, 0, len(data)+8))
	binary.Write(buff, binary.LittleEndian, uint32(len(data)))
	buff.Write(data)
	_, err = conn.Write(buff.Bytes())

	if err == nil {
		return nil
	}

	// if we got here it errored out, clear out the request
	self.requestBufferLock.Lock()
	delete(self.requestBuffer, *request.Id)
	self.requestBufferLock.Unlock()
	self.reconnect()
	return err
}

func (self *ProtobufClient) readResponses() {
	message := make([]byte, 0, MAX_RESPONSE_SIZE)
	buff := bytes.NewBuffer(message)
	for !self.stopped {
		buff.Reset()
		conn := self.getConnection()
		if conn == nil {
			time.Sleep(200 * time.Millisecond)
			continue
		}
		var messageSizeU uint32
		var err error
		err = binary.Read(conn, binary.LittleEndian, &messageSizeU)
		if err != nil {
			log.Error("Error while reading messsage size: %d", err)
			time.Sleep(200 * time.Millisecond)
			continue
		}
		messageSize := int64(messageSizeU)
		messageReader := io.LimitReader(conn, messageSize)
		_, err = io.Copy(buff, messageReader)
		if err != nil {
			log.Error("Error while reading message: %d", err)
			time.Sleep(200 * time.Millisecond)
			continue
		}
		response, err := protocol.DecodeResponse(buff)
		if err != nil {
			log.Error("error unmarshaling response: %s", err)
			time.Sleep(200 * time.Millisecond)
		} else {
			self.sendResponse(response)
		}
	}
}

func (self *ProtobufClient) sendResponse(response *protocol.Response) {
	self.requestBufferLock.RLock()
	req, ok := self.requestBuffer[*response.RequestId]
	self.requestBufferLock.RUnlock()
	if !ok {
		return
	}

	deleteRequest := false
	switch rt := response.GetType(); rt {
	case protocol.Response_END_STREAM,
		protocol.Response_HEARTBEAT,
		protocol.Response_ERROR:
		deleteRequest = true
	case protocol.Response_QUERY:
		// do nothing
	default:
		panic(fmt.Errorf("Unknown response type: %s", rt))
	}

	self.requestBufferLock.Lock()
	req, ok = self.requestBuffer[*response.RequestId]
	if deleteRequest {
		delete(self.requestBuffer, *response.RequestId)
	}
	self.requestBufferLock.Unlock()
	if !ok {
		return
	}

	log.Debug("ProtobufClient yielding to %s %s", req.r.Name(), response)
	req.r.Yield(response)
}

func (self *ProtobufClient) reconnect() net.Conn {
	select {
	case <-self.reconChan:
		self.reconGroup.Add(1)
		defer func() {
			self.reconGroup.Done()
			self.reconChan <- struct{}{}
		}()
	default:
		self.reconGroup.Wait()
		return self.conn
	}

	if self.conn != nil {
		self.conn.Close()
	}
	conn, err := net.DialTimeout("tcp", self.hostAndPort, self.writeTimeout)
	if err != nil {
		self.attempts++
		if self.attempts < 100 {
			return nil
		}

		log.Error("failed to connect to %s %d times", self.hostAndPort, self.attempts)
		self.attempts = 0
	}

	self.conn = conn
	log.Info("connected to %s", self.hostAndPort)
	return conn
}

func (self *ProtobufClient) peridicallySweepTimedOutRequests() {
	for {
		time.Sleep(time.Minute)
		self.requestBufferLock.Lock()
		maxAge := time.Now().Add(-MAX_REQUEST_TIME)
		for k, req := range self.requestBuffer {
			if req.timeMade.Before(maxAge) {
				delete(self.requestBuffer, k)
				log.Warn("Request timed out: ", req.request)
			}
		}
		self.requestBufferLock.Unlock()
	}
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2013-2016 Oryx(ossrs)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

package app

import (
	"fmt"
	"github.com/ossrs/go-oryx/agent"
	"github.com/ossrs/go-oryx/core"
	"os"
	"os/signal"
	"runtime"
	"runtime/debug"
	"runtime/pprof"
	"sync"
	"syscall"
	"time"
)

// the state of server, state graph:
//      Init => Normal(Ready => Running)
//      Init/Normal => Closed
type ServerState int

const (
	StateInit ServerState = 1 << iota
	StateReady
	StateRunning
	StateClosed
)

type Server struct {
	ctx core.Context
	// signal handler.
	sigs chan os.Signal
	// whether closed.
	closed  ServerState
	closing chan bool
	// for system internal to notify quit.
	quit chan bool
	wg   sync.WaitGroup
	// core components.
	htbt   *Heartbeat
	logger *simpleLogger
	rtmp   core.OpenCloser
	// the locker for state, for instance, the closed.
	lock sync.Mutex
}

func NewServer(ctx core.Context) *Server {
	v := &Server{
		ctx:     ctx,
		sigs:    make(chan os.Signal, 1),
		closed:  StateInit,
		closing: make(chan bool, 1),
		quit:    make(chan bool, 1),
		htbt:    NewHeartbeat(ctx),
		logger:  &simpleLogger{ctx: ctx},
	}
	v.rtmp = agent.NewRtmp(ctx, v)

	core.Conf.Subscribe(v)

	return v
}

// notify server to stop and wait for cleanup.
func (v *Server) Close() {
	ctx := v.ctx

	// wait for stopped.
	v.lock.Lock()
	defer v.lock.Unlock()

	// only create?
	if v.closed == StateInit {
		return
	}

	// closed?
	if v.closed == StateClosed {
		core.Warn.Println(ctx, "server already closed.")
		return
	}

	// notify to close.
	if v.closed == StateRunning {
		core.Info.Println(ctx, "notify server to stop.")
		select {
		case v.quit <- true:
		default:
		}
	}

	// wait for closed.
	if v.closed == StateRunning {
		<-v.closing
	}

	// do cleanup when stopped.
	core.Conf.Unsubscribe(v)

	// close the rtmp agent.
	if err := v.rtmp.Close(); err != nil {
		core.Warn.Println(ctx, "close rtmp agent failed. err is", err)
	}

	// close the agent manager.
	agent.Manager.Close()

	// when cpu profile is enabled, close it.
	if core.Conf.Go.CpuProfile != "" {
		pprof.StopCPUProfile()
		core.Trace.Println(ctx, "cpu profile ok, file is", core.Conf.Go.CpuProfile)
	}

	// when memory profile enabled, write heap info.
	if core.Conf.Go.MemProfile != "" {
		if f, err := os.Create(core.Conf.Go.MemProfile); err != nil {
			core.Warn.Println(ctx, "ignore open memory profile failed. err is", err)
		} else {
			defer f.Close()
			if err = pprof.Lookup("heap").WriteTo(f, 0); err != nil {
				core.Warn.Println(ctx, "write memory profile failed. err is", err)
			}
		}
		core.Trace.Println(ctx, "mem profile ok, file is", core.Conf.Go.MemProfile)
	}

	// ok, closed.
	v.closed = StateClosed
	core.Trace.Println(ctx, "server closed")
}

func (v *Server) ParseConfig(conf string) (err error) {
	ctx := v.ctx

	v.lock.Lock()
	defer v.lock.Unlock()

	if v.closed != StateInit {
		panic("server invalid state.")
	}

	core.Trace.Println(ctx, "start to parse config file", conf)
	if err = core.Conf.Loads(conf); err != nil {
		return
	}

	return
}

func (v *Server) PrepareLogger() (err error) {
	v.lock.Lock()
	defer v.lock.Unlock()

	if v.closed != StateInit {
		panic("server invalid state.")
	}

	if err = v.applyLogger(core.Conf); err != nil {
		return
	}

	return
}

func (v *Server) initializeRuntime() (err error) {
	ctx := v.ctx

	// install signals.
	// TODO: FIXME: when process the current signal, others may drop.
	signal.Notify(v.sigs)

	// apply the cpu profile.
	if err = v.applyCpuProfile(core.Conf); err != nil {
		return
	}

	// apply the gc percent.
	if err = v.applyGcPercent(core.Conf); err != nil {
		return
	}

	// show gc trace.
	go func() {
		stat := &debug.GCStats{}

		for {
			if core.Conf.Go.GcTrace > 0 {
				pgc := stat.NumGC
				debug.ReadGCStats(stat)
				if len(stat.Pause) > 3 {
					stat.Pause = append([]time.Duration{}, stat.Pause[:3]...)
				}
				if pgc < stat.NumGC {
					core.Trace.Println(ctx, "gc", stat.NumGC, stat.PauseTotal, stat.Pause, stat.PauseQuantiles)
				}
				time.Sleep(time.Duration(core.Conf.Go.GcTrace) * time.Second)
			} else {
				time.Sleep(3 * time.Second)
			}
		}
	}()

	return
}

func (v *Server) Initialize() (err error) {
	ctx := v.ctx

	v.lock.Lock()
	defer v.lock.Unlock()

	if v.closed != StateInit {
		panic("server invalid state.")
	}

	// about the runtime.
	if err = v.initializeRuntime(); err != nil {
		return
	}

	// use worker container to fork.
	var wc core.WorkerContainer = v

	// heartbeat with http api.
	if err = v.htbt.Initialize(wc); err != nil {
		return
	}

	// reload goroutine
	wc.GFork("reload", core.Conf.ReloadCycle)
	// heartbeat goroutine
	wc.GFork("htbt(discovery)", v.htbt.discoveryCycle)
	wc.GFork("htbt(main)", v.htbt.beatCycle)
	// open rtmp agent.
	if err = v.rtmp.Open(); err != nil {
		core.Error.Println(ctx, "open rtmp agent failed. err is", err)
		return
	}

	c := core.Conf
	l := fmt.Sprintf("%v(%v/%v)", c.Log.Tank, c.Log.Level, c.Log.File)
	if !c.LogToFile() {
		l = fmt.Sprintf("%v(%v)", c.Log.Tank, c.Log.Level)
	}
	core.Trace.Println(ctx, fmt.Sprintf("init server ok, conf=%v, log=%v, workers=%v/%v, gc=%v/%v%%, daemon=%v",
		c.Conf(), l, c.Workers, runtime.NumCPU(), c.Go.GcInterval, c.Go.GcPercent, c.Daemon))

	// set to ready, requires cleanup.
	v.closed = StateReady

	return
}

func (v *Server) Run() (err error) {
	ctx := v.ctx

	func() {
		v.lock.Lock()
		defer v.lock.Unlock()

		if v.closed != StateReady {
			panic("server invalid state.")
		}
		v.closed = StateRunning
	}()

	// when terminated, notify the chan.
	defer func() {
		select {
		case v.closing <- true:
		default:
		}
	}()

	core.Info.Println(ctx, "server cycle running")

	// run server, apply settings.
	v.applyMultipleProcesses(core.Conf.Workers)

	var wc core.WorkerContainer = v
	for {
		var gcc <-chan time.Time = nil
		if core.Conf.Go.GcInterval > 0 {
			gcc = time.After(time.Second * time.Duration(core.Conf.Go.GcInterval))
		}

		select {
		case signal := <-v.sigs:
			core.Trace.Println(ctx, "got signal", signal)
			switch signal {
			case SIGUSR1, SIGUSR2:
				panic("panic by SIGUSR1/2")
			case os.Interrupt, syscall.SIGTERM:
				// SIGINT, SIGTERM
				wc.Quit()
			}
		case <-wc.QC():
			wc.Quit()

			// wait for all goroutines quit.
			v.wg.Wait()
			core.Warn.Println(ctx, "server cycle ok")
			return
		case <-gcc:
			runtime.GC()
			core.Info.Println(ctx, "go runtime gc every", core.Conf.Go.GcInterval, "seconds")
		}
	}

	return
}

// interface WorkContainer
func (v *Server) QC() <-chan bool {
	return v.quit
}

func (v *Server) Quit() error {
	select {
	case v.quit <- true:
	default:
	}

	return core.QuitError
}

func (v *Server) GFork(name string, f func(core.WorkerContainer)) {
	ctx := v.ctx

	v.wg.Add(1)
	go func() {
		defer v.wg.Done()

		defer func() {
			if r := recover(); r != nil {
				if !core.IsNormalQuit(r) {
					core.Warn.Println(ctx, "rtmp ignore", r)
				}

				core.Error.Println(ctx, string(debug.Stack()))

				v.Quit()
			}
		}()

		f(v)

		if name != "" {
			core.Trace.Println(ctx, name, "worker terminated.")
		}
	}()
}

func (v *Server) applyMultipleProcesses(workers int) {
	ctx := v.ctx

	if workers < 0 {
		panic("should not be negative workers")
	}

	if workers == 0 {
		workers = runtime.NumCPU()
	}
	pv := runtime.GOMAXPROCS(workers)

	core.Trace.Println(ctx, "apply workers", workers, "and previous is", pv)
}

func (v *Server) applyLogger(c *core.Config) (err error) {
	ctx := v.ctx

	if err = v.logger.close(c); err != nil {
		return
	}
	core.Info.Println(ctx, "close logger ok")

	if err = v.logger.open(c); err != nil {
		return
	}
	core.Info.Println(ctx, "open logger ok")

	return
}

func (v *Server) applyCpuProfile(c *core.Config) (err error) {
	ctx := v.ctx

	pprof.StopCPUProfile()

	if c.Go.CpuProfile == "" {
		return
	}

	var f *os.File
	if f, err = os.Create(c.Go.CpuProfile); err != nil {
		core.Error.Println(ctx, "open cpu profile file failed. err is", err)
		return
	}
	if err = pprof.StartCPUProfile(f); err != nil {
		core.Error.Println(ctx, "start cpu profile failed. err is", err)
		return
	}
	return
}

func (v *Server) applyGcPercent(c *core.Config) (err error) {
	ctx := v.ctx

	if c.Go.GcPercent == 0 {
		debug.SetGCPercent(100)
		return
	}

	pv := debug.SetGCPercent(c.Go.GcPercent)
	core.Trace.Println(ctx, "set gc percent from", pv, "to", c.Go.GcPercent)
	return
}

// interface ReloadHandler
func (v *Server) OnReloadGlobal(scope int, cc, pc *core.Config) (err error) {
	if scope == core.ReloadWorkers {
		v.applyMultipleProcesses(cc.Workers)
	} else if scope == core.ReloadLog {
		v.applyLogger(cc)
	} else if scope == core.ReloadCpuProfile {
		v.applyCpuProfile(cc)
	} else if scope == core.ReloadGcPercent {
		v.applyGcPercent(cc)
	}

	return
}

func (v *Server) OnReloadVhost(vhost string, scope int, cc, pc *core.Config) (err error) {
	return
}
<commit_msg>for #52, start new goroutine for signal when quit.<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2013-2016 Oryx(ossrs)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

package app

import (
	"fmt"
	"github.com/ossrs/go-oryx/agent"
	"github.com/ossrs/go-oryx/core"
	"os"
	"os/signal"
	"runtime"
	"runtime/debug"
	"runtime/pprof"
	"sync"
	"syscall"
	"time"
)

// the state of server, state graph:
//      Init => Normal(Ready => Running)
//      Init/Normal => Closed
type ServerState int

const (
	StateInit ServerState = 1 << iota
	StateReady
	StateRunning
	StateClosed
)

type Server struct {
	ctx core.Context
	// signal handler.
	sigs chan os.Signal
	// whether closed.
	closed  ServerState
	closing chan bool
	// for system internal to notify quit.
	quit chan bool
	wg   sync.WaitGroup
	// core components.
	htbt   *Heartbeat
	logger *simpleLogger
	rtmp   core.OpenCloser
	// the locker for state, for instance, the closed.
	lock sync.Mutex
}

func NewServer(ctx core.Context) *Server {
	v := &Server{
		ctx:     ctx,
		sigs:    make(chan os.Signal, 1),
		closed:  StateInit,
		closing: make(chan bool, 1),
		quit:    make(chan bool, 1),
		htbt:    NewHeartbeat(ctx),
		logger:  &simpleLogger{ctx: ctx},
	}
	v.rtmp = agent.NewRtmp(ctx, v)

	core.Conf.Subscribe(v)

	return v
}

// notify server to stop and wait for cleanup.
func (v *Server) Close() {
	ctx := v.ctx

	// wait for stopped.
	v.lock.Lock()
	defer v.lock.Unlock()

	// only create?
	if v.closed == StateInit {
		return
	}

	// closed?
	if v.closed == StateClosed {
		core.Warn.Println(ctx, "server already closed.")
		return
	}

	// notify to close.
	if v.closed == StateRunning {
		core.Info.Println(ctx, "notify server to stop.")
		select {
		case v.quit <- true:
		default:
		}
	}

	// wait for closed.
	if v.closed == StateRunning {
		<-v.closing
	}

	// do cleanup when stopped.
	core.Conf.Unsubscribe(v)

	// close the rtmp agent.
	if err := v.rtmp.Close(); err != nil {
		core.Warn.Println(ctx, "close rtmp agent failed. err is", err)
	}

	// close the agent manager.
	agent.Manager.Close()

	// when cpu profile is enabled, close it.
	if core.Conf.Go.CpuProfile != "" {
		pprof.StopCPUProfile()
		core.Trace.Println(ctx, "cpu profile ok, file is", core.Conf.Go.CpuProfile)
	}

	// when memory profile enabled, write heap info.
	if core.Conf.Go.MemProfile != "" {
		if f, err := os.Create(core.Conf.Go.MemProfile); err != nil {
			core.Warn.Println(ctx, "ignore open memory profile failed. err is", err)
		} else {
			defer f.Close()
			if err = pprof.Lookup("heap").WriteTo(f, 0); err != nil {
				core.Warn.Println(ctx, "write memory profile failed. err is", err)
			}
		}
		core.Trace.Println(ctx, "mem profile ok, file is", core.Conf.Go.MemProfile)
	}

	// ok, closed.
	v.closed = StateClosed
	core.Trace.Println(ctx, "server closed")
}

func (v *Server) ParseConfig(conf string) (err error) {
	ctx := v.ctx

	v.lock.Lock()
	defer v.lock.Unlock()

	if v.closed != StateInit {
		panic("server invalid state.")
	}

	core.Trace.Println(ctx, "start to parse config file", conf)
	if err = core.Conf.Loads(conf); err != nil {
		return
	}

	return
}

func (v *Server) PrepareLogger() (err error) {
	v.lock.Lock()
	defer v.lock.Unlock()

	if v.closed != StateInit {
		panic("server invalid state.")
	}

	if err = v.applyLogger(core.Conf); err != nil {
		return
	}

	return
}

func (v *Server) initializeRuntime() (err error) {
	ctx := v.ctx

	// install signals.
	// TODO: FIXME: when process the current signal, others may drop.
	signal.Notify(v.sigs)

	// apply the cpu profile.
	if err = v.applyCpuProfile(core.Conf); err != nil {
		return
	}

	// apply the gc percent.
	if err = v.applyGcPercent(core.Conf); err != nil {
		return
	}

	// show gc trace.
	go func() {
		stat := &debug.GCStats{}

		for {
			if core.Conf.Go.GcTrace > 0 {
				pgc := stat.NumGC
				debug.ReadGCStats(stat)
				if len(stat.Pause) > 3 {
					stat.Pause = append([]time.Duration{}, stat.Pause[:3]...)
				}
				if pgc < stat.NumGC {
					core.Trace.Println(ctx, "gc", stat.NumGC, stat.PauseTotal, stat.Pause, stat.PauseQuantiles)
				}
				time.Sleep(time.Duration(core.Conf.Go.GcTrace) * time.Second)
			} else {
				time.Sleep(3 * time.Second)
			}
		}
	}()

	return
}

func (v *Server) Initialize() (err error) {
	ctx := v.ctx

	v.lock.Lock()
	defer v.lock.Unlock()

	if v.closed != StateInit {
		panic("server invalid state.")
	}

	// about the runtime.
	if err = v.initializeRuntime(); err != nil {
		return
	}

	// use worker container to fork.
	var wc core.WorkerContainer = v

	// heartbeat with http api.
	if err = v.htbt.Initialize(wc); err != nil {
		return
	}

	// reload goroutine
	wc.GFork("reload", core.Conf.ReloadCycle)
	// heartbeat goroutine
	wc.GFork("htbt(discovery)", v.htbt.discoveryCycle)
	wc.GFork("htbt(main)", v.htbt.beatCycle)
	// open rtmp agent.
	if err = v.rtmp.Open(); err != nil {
		core.Error.Println(ctx, "open rtmp agent failed. err is", err)
		return
	}

	c := core.Conf
	l := fmt.Sprintf("%v(%v/%v)", c.Log.Tank, c.Log.Level, c.Log.File)
	if !c.LogToFile() {
		l = fmt.Sprintf("%v(%v)", c.Log.Tank, c.Log.Level)
	}
	core.Trace.Println(ctx, fmt.Sprintf("init server ok, conf=%v, log=%v, workers=%v/%v, gc=%v/%v%%, daemon=%v",
		c.Conf(), l, c.Workers, runtime.NumCPU(), c.Go.GcInterval, c.Go.GcPercent, c.Daemon))

	// set to ready, requires cleanup.
	v.closed = StateReady

	return
}

func (v *Server) onSignal(signal syscall.Signal) {
	ctx := v.ctx
	wc := v

	core.Trace.Println(ctx, "got signal", signal)
	switch signal {
	case SIGUSR1, SIGUSR2:
		panic("panic by SIGUSR1/2")
	case os.Interrupt, syscall.SIGTERM:
		// SIGINT, SIGTERM
		wc.Quit()
	}
}

func (v *Server) Run() (err error) {
	ctx := v.ctx

	func() {
		v.lock.Lock()
		defer v.lock.Unlock()

		if v.closed != StateReady {
			panic("server invalid state.")
		}
		v.closed = StateRunning
	}()

	// when terminated, notify the chan.
	defer func() {
		select {
		case v.closing <- true:
		default:
		}
	}()

	core.Info.Println(ctx, "server cycle running")

	// run server, apply settings.
	v.applyMultipleProcesses(core.Conf.Workers)

	var wc core.WorkerContainer = v
	for {
		var gcc <-chan time.Time = nil
		if core.Conf.Go.GcInterval > 0 {
			gcc = time.After(time.Second * time.Duration(core.Conf.Go.GcInterval))
		}

		select {
		case signal := <-v.sigs:
			v.onSignal(signal)
		case <-wc.QC():
			wc.Quit()

			// for the following quit will block all signal process,
			// we start new goroutine to process the panic signal only.
			go func() {
				for s := range v.sigs {
					v.onSignal(s)
				}
			}()

			// wait for all goroutines quit.
			v.wg.Wait()
			core.Warn.Println(ctx, "server cycle ok")
			return
		case <-gcc:
			runtime.GC()
			core.Info.Println(ctx, "go runtime gc every", core.Conf.Go.GcInterval, "seconds")
		}
	}

	return
}

// interface WorkContainer
func (v *Server) QC() <-chan bool {
	return v.quit
}

func (v *Server) Quit() error {
	select {
	case v.quit <- true:
	default:
	}

	return core.QuitError
}

func (v *Server) GFork(name string, f func(core.WorkerContainer)) {
	ctx := v.ctx

	v.wg.Add(1)
	go func() {
		defer v.wg.Done()

		defer func() {
			if r := recover(); r != nil {
				if !core.IsNormalQuit(r) {
					core.Warn.Println(ctx, "rtmp ignore", r)
				}

				core.Error.Println(ctx, string(debug.Stack()))

				v.Quit()
			}
		}()

		f(v)

		if name != "" {
			core.Trace.Println(ctx, name, "worker terminated.")
		}
	}()
}

func (v *Server) applyMultipleProcesses(workers int) {
	ctx := v.ctx

	if workers < 0 {
		panic("should not be negative workers")
	}

	if workers == 0 {
		workers = runtime.NumCPU()
	}
	pv := runtime.GOMAXPROCS(workers)

	core.Trace.Println(ctx, "apply workers", workers, "and previous is", pv)
}

func (v *Server) applyLogger(c *core.Config) (err error) {
	ctx := v.ctx

	if err = v.logger.close(c); err != nil {
		return
	}
	core.Info.Println(ctx, "close logger ok")

	if err = v.logger.open(c); err != nil {
		return
	}
	core.Info.Println(ctx, "open logger ok")

	return
}

func (v *Server) applyCpuProfile(c *core.Config) (err error) {
	ctx := v.ctx

	pprof.StopCPUProfile()

	if c.Go.CpuProfile == "" {
		return
	}

	var f *os.File
	if f, err = os.Create(c.Go.CpuProfile); err != nil {
		core.Error.Println(ctx, "open cpu profile file failed. err is", err)
		return
	}
	if err = pprof.StartCPUProfile(f); err != nil {
		core.Error.Println(ctx, "start cpu profile failed. err is", err)
		return
	}
	return
}

func (v *Server) applyGcPercent(c *core.Config) (err error) {
	ctx := v.ctx

	if c.Go.GcPercent == 0 {
		debug.SetGCPercent(100)
		return
	}

	pv := debug.SetGCPercent(c.Go.GcPercent)
	core.Trace.Println(ctx, "set gc percent from", pv, "to", c.Go.GcPercent)
	return
}

// interface ReloadHandler
func (v *Server) OnReloadGlobal(scope int, cc, pc *core.Config) (err error) {
	if scope == core.ReloadWorkers {
		v.applyMultipleProcesses(cc.Workers)
	} else if scope == core.ReloadLog {
		v.applyLogger(cc)
	} else if scope == core.ReloadCpuProfile {
		v.applyCpuProfile(cc)
	} else if scope == core.ReloadGcPercent {
		v.applyGcPercent(cc)
	}

	return
}

func (v *Server) OnReloadVhost(vhost string, scope int, cc, pc *core.Config) (err error) {
	return
}
<|endoftext|>


<commit_before>package executer

import (
	"bufio"
	"bytes"
	"github.com/subutai-io/base/agent/agent/container"
	"github.com/subutai-io/base/agent/config"
	"github.com/subutai-io/base/agent/log"
	"gopkg.in/lxc/go-lxc.v2"
	"io"
	"os"
	"os/exec"
	"os/user"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
	"unsafe"
)

type EncRequest struct {
	HostId  string `json:"hostid"`
	Request string `json:"request"`
}
type Request struct {
	Id      string         `json:"id"`
	Request RequestOptions `json:"request"`
}

type RequestOptions struct {
	Type        string            `json:"type"`
	Id          string            `json:"id"`
	CommandId   string            `json:"commandId"`
	WorkingDir  string            `json:"workingDirectory"`
	Command     string            `json:"command"`
	Args        []string          `json:"args"`
	Environment map[string]string `json:"environment"`
	StdOut      string            `json:"stdOut"`
	StdErr      string            `json:"stdErr"`
	RunAs       string            `json:"runAs"`
	Timeout     int               `json:"timeout"`
	IsDaemon    int               `json:"isDaemon"`
}

type Response struct {
	ResponseOpts ResponseOptions `json:"response"`
	Id           string          `json:"id"`
}

type ResponseOptions struct {
	Type           string `json:"type"`
	Id             string `json:"id"`
	CommandId      string `json:"commandId"`
	Pid            int    `json:"pid"`
	ResponseNumber int    `json:"responseNumber,omitempty"`
	StdOut         string `json:"stdOut,omitempty"`
	StdErr         string `json:"stdErr,omitempty"`
	ExitCode       string `json:"exitCode,omitempty"`
}

func ExecHost(req RequestOptions, out_c chan<- ResponseOptions) {
	cmd := buildCmd(&req)
	if cmd == nil {
		close(out_c)
		return
	}
	//read/write pipes stdout
	rop, wop, _ := os.Pipe()
	//read/out pipes stderr
	rep, wep, _ := os.Pipe()

	defer rop.Close()
	defer rep.Close()

	cmd.Stdout = wop
	cmd.Stderr = wep
	if req.IsDaemon == 1 {
		cmd.SysProcAttr = &syscall.SysProcAttr{}
		cmd.SysProcAttr.Setpgid = true
		cmd.SysProcAttr.Credential = &syscall.Credential{Uid: 0, Gid: 0}
	}
	err := cmd.Start()
	log.Check(log.WarnLevel, "Executing command: "+req.CommandId+" "+req.Command+" "+strings.Join(req.Args, " "), err)

	wop.Close()
	wep.Close()

	start := time.Now().Unix()
	end := time.Now().Add(time.Duration(req.Timeout) * time.Second).Unix()
	var response = genericResponse(&req)
	response.ResponseNumber = 1
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		r := bufio.NewReader(rop)
		for line, isPrefix, err := r.ReadLine(); err == nil; line, isPrefix, err = r.ReadLine() {
			response.StdOut = response.StdOut + string(line)
			if !isPrefix {
				response.StdOut = response.StdOut + "\n"
			}
			if len(response.StdOut) > 50000 {
				out_c <- response
				start = now()
				response.StdErr, response.StdOut = "", ""
				response.ResponseNumber++
			}
			if end-now() < 0 {
				break
			}
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		scanErr := bufio.NewScanner(rep)
		for scanErr.Scan() {
			response.StdErr = response.StdErr + scanErr.Text() + "\n"
			if len(response.StdErr) > 50000 {
				out_c <- response
				start = now()
				response.StdErr, response.StdOut = "", ""
				response.ResponseNumber++
			}
			if end-now() < 0 {
				break
			}
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		for end-now() > 0 {
			if now()-start > 10 {
				out_c <- response
				start = now()
				response.StdErr, response.StdOut = "", ""
				response.ResponseNumber++
			}
			time.Sleep(time.Second)
		}
	}()

	done := make(chan error)
	go func() { done <- cmd.Wait() }()
	select {
	case <-done:
		response.ExitCode = "0"
		if req.IsDaemon != 1 {
			response.ExitCode = strconv.Itoa(cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus())
		}
		out_c <- response
		end = -1
	case <-time.After(time.Duration(req.Timeout) * time.Second):
		if req.IsDaemon == 1 {
			response.ExitCode = "0"
			out_c <- response
			<-done
		} else {
			cmd.Process.Kill()
			response.Type = config.Broker.ExecuteTimeout
			cmd.Process.Wait()
			if cmd.ProcessState != nil {
				response.ExitCode = strconv.Itoa(cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus())
			} else {
				response.ExitCode = "-1"
			}
			out_c <- response
		}
	}
	wg.Wait()
	close(out_c)
}

func buildCmd(r *RequestOptions) *exec.Cmd {
	user, err := user.Lookup(r.RunAs)
	if log.Check(log.WarnLevel, "User lookup: "+r.RunAs, err) {
		return nil
	}
	uid, err := strconv.Atoi(user.Uid)
	if log.Check(log.WarnLevel, "UID lookup: "+user.Uid, err) {
		return nil
	}
	gid, err := strconv.Atoi(user.Gid)
	if log.Check(log.WarnLevel, "GID lookup: "+user.Gid, err) {
		return nil
	}
	gid32 := *(*uint32)(unsafe.Pointer(&gid))
	uid32 := *(*uint32)(unsafe.Pointer(&uid))

	var buff bytes.Buffer
	buff.WriteString(r.Command + " ")
	for _, arg := range r.Args {
		buff.WriteString(arg + " ")
	}
	cmd := exec.Command("/bin/bash", "-c", buff.String())
	cmd.Dir = r.WorkingDir
	cmd.SysProcAttr = &syscall.SysProcAttr{}
	cmd.SysProcAttr.Credential = &syscall.Credential{Uid: uid32, Gid: gid32}

	return cmd
}

//prepare basic response
func genericResponse(req *RequestOptions) ResponseOptions {
	return ResponseOptions{
		Type:      config.Broker.ExecuteResponce,
		CommandId: req.CommandId,
		Id:        req.Id,
	}
}

func now() int64 {
	return time.Now().Unix()
}

func AttachContainer(name string, r RequestOptions, out_c chan<- ResponseOptions) {
	lxc_c, _ := lxc.NewContainer(name, config.Agent.LxcPrefix)
	opts := lxc.DefaultAttachOptions

	var res ResponseOptions = genericResponse(&r)

	o_read, o_write, err := os.Pipe()
	log.Check(log.WarnLevel, "Creating pipe for container attach", err)
	e_read, e_write, err := os.Pipe()
	log.Check(log.WarnLevel, "Creating pipe for container attach", err)

	var chunk bytes.Buffer
	defer o_read.Close()
	defer e_read.Close()

	opts.StdoutFd = o_write.Fd()
	opts.StderrFd = e_write.Fd()

	opts.UID, opts.GID = container.GetCredentials(r.RunAs, name)

	opts.Cwd = r.WorkingDir

	var exitCode int
	var cmd bytes.Buffer

	cmd.WriteString(r.Command)
	for _, a := range r.Args {
		cmd.WriteString(a + " ")
	}

	log.Debug("Executing command in container " + name + ":" + cmd.String())
	go func() {
		if lxc_c.Running() {
			log.Debug("Container " + name + " is running")
		}

		exitCode, err = lxc_c.RunCommandStatus([]string{"timeout", strconv.Itoa(r.Timeout), "/bin/bash", "-c", cmd.String()}, opts)
		log.Check(log.WarnLevel, "Execution command", err)

		o_write.Close()
		e_write.Close()
	}()

	out := bufio.NewScanner(o_read)
	var e bytes.Buffer
	start_time := time.Now().Unix()

	//response counter
	var resN int = 1

	for out.Scan() {
		//we have more stdout coming in
		//collect 1000 bytes and send the chunk
		chunk.WriteString(out.Text() + "\n")
		//send chunk every 1000 bytes or 10 seconds
		if chunk.Len() >= 1000 || now()-start_time >= 10 {
			res.StdOut = chunk.String()
			res.ResponseNumber = resN
			log.Info("Command intermediate result " + chunk.String())
			out_c <- res
			chunk.Truncate(0)
			resN++
			start_time = now()
		}
	}

	log.Info("Command intermediate result " + chunk.String())
	io.Copy(&e, e_read)
	if exitCode == 0 {
		res.Type = config.Broker.ExecuteResponce
		res.ExitCode = strconv.Itoa(exitCode)
		if chunk.Len() > 0 {
			if resN > 1 {
				res.ResponseNumber = resN
			}
			log.Info("Command intermediate result " + chunk.String())
			res.StdOut = chunk.String()
		}
		res.StdErr = e.String()
		out_c <- res
	} else {
		log.Debug("Exited with exit code", strconv.Itoa(exitCode))
		if exitCode/256 == 124 {
			res.Type = config.Broker.ExecuteTimeout
		}
		res.ExitCode = strconv.Itoa(exitCode / 256)
		res.StdOut = chunk.String()
		res.StdErr = e.String()
		out_c <- res
	}
	lxc.Release(lxc_c)
	close(out_c)
}
<commit_msg>Fixed empty responseNumber #458<commit_after>package executer

import (
	"bufio"
	"bytes"
	"github.com/subutai-io/base/agent/agent/container"
	"github.com/subutai-io/base/agent/config"
	"github.com/subutai-io/base/agent/log"
	"gopkg.in/lxc/go-lxc.v2"
	"io"
	"os"
	"os/exec"
	"os/user"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
	"unsafe"
)

type EncRequest struct {
	HostId  string `json:"hostid"`
	Request string `json:"request"`
}
type Request struct {
	Id      string         `json:"id"`
	Request RequestOptions `json:"request"`
}

type RequestOptions struct {
	Type        string            `json:"type"`
	Id          string            `json:"id"`
	CommandId   string            `json:"commandId"`
	WorkingDir  string            `json:"workingDirectory"`
	Command     string            `json:"command"`
	Args        []string          `json:"args"`
	Environment map[string]string `json:"environment"`
	StdOut      string            `json:"stdOut"`
	StdErr      string            `json:"stdErr"`
	RunAs       string            `json:"runAs"`
	Timeout     int               `json:"timeout"`
	IsDaemon    int               `json:"isDaemon"`
}

type Response struct {
	ResponseOpts ResponseOptions `json:"response"`
	Id           string          `json:"id"`
}

type ResponseOptions struct {
	Type           string `json:"type"`
	Id             string `json:"id"`
	CommandId      string `json:"commandId"`
	Pid            int    `json:"pid"`
	ResponseNumber int    `json:"responseNumber,omitempty"`
	StdOut         string `json:"stdOut,omitempty"`
	StdErr         string `json:"stdErr,omitempty"`
	ExitCode       string `json:"exitCode,omitempty"`
}

func ExecHost(req RequestOptions, out_c chan<- ResponseOptions) {
	cmd := buildCmd(&req)
	if cmd == nil {
		close(out_c)
		return
	}
	//read/write pipes stdout
	rop, wop, _ := os.Pipe()
	//read/out pipes stderr
	rep, wep, _ := os.Pipe()

	defer rop.Close()
	defer rep.Close()

	cmd.Stdout = wop
	cmd.Stderr = wep
	if req.IsDaemon == 1 {
		cmd.SysProcAttr = &syscall.SysProcAttr{}
		cmd.SysProcAttr.Setpgid = true
		cmd.SysProcAttr.Credential = &syscall.Credential{Uid: 0, Gid: 0}
	}
	err := cmd.Start()
	log.Check(log.WarnLevel, "Executing command: "+req.CommandId+" "+req.Command+" "+strings.Join(req.Args, " "), err)

	wop.Close()
	wep.Close()

	start := time.Now().Unix()
	end := time.Now().Add(time.Duration(req.Timeout) * time.Second).Unix()
	var response = genericResponse(&req)
	response.ResponseNumber = 1
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		r := bufio.NewReader(rop)
		for line, isPrefix, err := r.ReadLine(); err == nil; line, isPrefix, err = r.ReadLine() {
			response.StdOut = response.StdOut + string(line)
			if !isPrefix {
				response.StdOut = response.StdOut + "\n"
			}
			if len(response.StdOut) > 50000 {
				out_c <- response
				start = now()
				response.StdErr, response.StdOut = "", ""
				response.ResponseNumber++
			}
			if end-now() < 0 {
				break
			}
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		scanErr := bufio.NewScanner(rep)
		for scanErr.Scan() {
			response.StdErr = response.StdErr + scanErr.Text() + "\n"
			if len(response.StdErr) > 50000 {
				out_c <- response
				start = now()
				response.StdErr, response.StdOut = "", ""
				response.ResponseNumber++
			}
			if end-now() < 0 {
				break
			}
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		for end-now() > 0 {
			if now()-start > 10 {
				out_c <- response
				start = now()
				response.StdErr, response.StdOut = "", ""
				response.ResponseNumber++
			}
			time.Sleep(time.Second)
		}
	}()

	done := make(chan error)
	go func() { done <- cmd.Wait() }()
	select {
	case <-done:
		response.ExitCode = "0"
		if req.IsDaemon != 1 {
			response.ExitCode = strconv.Itoa(cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus())
		}
		out_c <- response
		end = -1
	case <-time.After(time.Duration(req.Timeout) * time.Second):
		if req.IsDaemon == 1 {
			response.ExitCode = "0"
			out_c <- response
			<-done
		} else {
			cmd.Process.Kill()
			response.Type = config.Broker.ExecuteTimeout
			cmd.Process.Wait()
			if cmd.ProcessState != nil {
				response.ExitCode = strconv.Itoa(cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus())
			} else {
				response.ExitCode = "-1"
			}
			out_c <- response
		}
	}
	wg.Wait()
	close(out_c)
}

func buildCmd(r *RequestOptions) *exec.Cmd {
	user, err := user.Lookup(r.RunAs)
	if log.Check(log.WarnLevel, "User lookup: "+r.RunAs, err) {
		return nil
	}
	uid, err := strconv.Atoi(user.Uid)
	if log.Check(log.WarnLevel, "UID lookup: "+user.Uid, err) {
		return nil
	}
	gid, err := strconv.Atoi(user.Gid)
	if log.Check(log.WarnLevel, "GID lookup: "+user.Gid, err) {
		return nil
	}
	gid32 := *(*uint32)(unsafe.Pointer(&gid))
	uid32 := *(*uint32)(unsafe.Pointer(&uid))

	var buff bytes.Buffer
	buff.WriteString(r.Command + " ")
	for _, arg := range r.Args {
		buff.WriteString(arg + " ")
	}
	cmd := exec.Command("/bin/bash", "-c", buff.String())
	cmd.Dir = r.WorkingDir
	cmd.SysProcAttr = &syscall.SysProcAttr{}
	cmd.SysProcAttr.Credential = &syscall.Credential{Uid: uid32, Gid: gid32}

	return cmd
}

//prepare basic response
func genericResponse(req *RequestOptions) ResponseOptions {
	return ResponseOptions{
		Type:      config.Broker.ExecuteResponce,
		CommandId: req.CommandId,
		Id:        req.Id,
	}
}

func now() int64 {
	return time.Now().Unix()
}

func AttachContainer(name string, r RequestOptions, out_c chan<- ResponseOptions) {
	lxc_c, _ := lxc.NewContainer(name, config.Agent.LxcPrefix)
	opts := lxc.DefaultAttachOptions

	var res ResponseOptions = genericResponse(&r)

	o_read, o_write, _ := os.Pipe()
	e_read, e_write, _ := os.Pipe()

	var chunk bytes.Buffer
	defer o_read.Close()
	defer e_read.Close()

	opts.StdoutFd = o_write.Fd()
	opts.StderrFd = e_write.Fd()

	opts.UID, opts.GID = container.GetCredentials(r.RunAs, name)

	opts.Cwd = r.WorkingDir

	var exitCode int
	var cmd bytes.Buffer

	cmd.WriteString(r.Command)
	for _, a := range r.Args {
		cmd.WriteString(a + " ")
	}

	log.Debug("Executing command in container " + name + ":" + cmd.String())
	go func() {
		exitCode, _ = lxc_c.RunCommandStatus([]string{"timeout", strconv.Itoa(r.Timeout), "/bin/bash", "-c", cmd.String()}, opts)

		o_write.Close()
		e_write.Close()
	}()

	out := bufio.NewScanner(o_read)
	var e bytes.Buffer
	start_time := time.Now().Unix()

	res.ResponseNumber = 1
	for out.Scan() {
		//we have more stdout coming in
		//collect 1000 bytes and send the chunk
		chunk.WriteString(out.Text() + "\n")
		//send chunk every 1000 bytes or 10 seconds
		if chunk.Len() >= 1000 || now()-start_time >= 10 {
			res.StdOut = chunk.String()
			out_c <- res
			chunk.Truncate(0)
			res.ResponseNumber++
			start_time = now()
		}
	}

	io.Copy(&e, e_read)
	if exitCode == 0 {
		res.Type = config.Broker.ExecuteResponce
		res.ExitCode = strconv.Itoa(exitCode)
		if chunk.Len() > 0 {
			res.StdOut = chunk.String()
		}
		res.StdErr = e.String()
		out_c <- res
	} else {
		if exitCode/256 == 124 {
			res.Type = config.Broker.ExecuteTimeout
		}
		res.ExitCode = strconv.Itoa(exitCode / 256)
		res.StdOut = chunk.String()
		res.StdErr = e.String()
		out_c <- res
	}
	lxc.Release(lxc_c)
	close(out_c)
}
<|endoftext|>


<commit_before>package containerd

import (
	gocontext "context"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"syscall"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/cio"
	"github.com/containerd/containerd/containers"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/oci"
	"github.com/docker/docker/pkg/signal"
	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/api/naming"
	"github.com/docker/swarmkit/log"
	gogotypes "github.com/gogo/protobuf/types"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
)

var (
	devNull                    *os.File
	errAdapterNotPrepared      = errors.New("container adapter not prepared")
	mountPropagationReverseMap = map[api.Mount_BindOptions_MountPropagation]string{
		api.MountPropagationPrivate:  "private",
		api.MountPropagationRPrivate: "rprivate",
		api.MountPropagationShared:   "shared",
		api.MountPropagationRShared:  "rshared",
		api.MountPropagationRSlave:   "slave",
		api.MountPropagationSlave:    "rslave",
	}
)

// containerAdapter conducts remote operations for a container. All calls
// are mostly naked calls to the client API, seeded with information from
// containerConfig.
type containerAdapter struct {
	client     *containerd.Client
	spec       *api.ContainerSpec
	secrets    exec.SecretGetter
	name       string
	image      containerd.Image // Pulled image
	container  containerd.Container
	task       containerd.Task
	exitStatus error
}

func newContainerAdapter(client *containerd.Client, task *api.Task, secrets exec.SecretGetter) (*containerAdapter, error) {
	spec := task.Spec.GetContainer()
	if spec == nil {
		return nil, exec.ErrRuntimeUnsupported
	}

	c := &containerAdapter{
		client:  client,
		spec:    spec,
		secrets: secrets,
		name:    naming.Task(task),
	}

	if err := c.reattach(context.Background()); err != nil {
		return nil, err
	}

	return c, nil
}

// reattaches to an existing container. If the container is found but
// the task is missing then still succeeds, allowing subsequent use of
// c.delete()
func (c *containerAdapter) reattach(ctx context.Context) error {
	container, err := c.client.LoadContainer(ctx, c.name)
	if err != nil {
		if errdefs.IsNotFound(err) {
			c.log(ctx).Debug("reattach: container not found")
			return nil
		}

		return errors.Wrap(err, "reattach: loading container")
	}
	c.log(ctx).Debug("reattach: loaded container")
	c.container = container

	// TODO(ijc) Consider an addition to container library which
	// directly attaches stdin to /dev/null.
	if devNull == nil {
		if devNull, err = os.Open(os.DevNull); err != nil {
			return errors.Wrap(err, "reattach: opening null device")
		}
	}

	task, err := container.Task(ctx, cio.WithAttach(devNull, os.Stdout, os.Stderr))
	if err != nil {
		if errdefs.IsNotFound(err) {
			c.log(ctx).WithError(err).Info("reattach: no running task")
			return nil
		}
		return errors.Wrap(err, "reattach: reattaching task")
	}
	c.task = task
	c.log(ctx).Debug("reattach: successful")
	return nil
}

func (c *containerAdapter) log(ctx context.Context) *logrus.Entry {
	return log.G(ctx).WithFields(logrus.Fields{
		"container.id": c.name,
	})
}

func (c *containerAdapter) pullImage(ctx context.Context) error {
	image, err := c.client.Pull(ctx, c.spec.Image, containerd.WithPullUnpack)
	if err != nil {
		return errors.Wrap(err, "pulling container image")
	}
	c.image = image

	return nil
}

func withMounts(ctx context.Context, ms []api.Mount) oci.SpecOpts {
	sort.Sort(mounts(ms))

	return func(_ gocontext.Context, _ oci.Client, _ *containers.Container, s *specs.Spec) error {
		for _, m := range ms {
			if !filepath.IsAbs(m.Target) {
				return errors.Errorf("mount %s is not absolute", m.Target)
			}

			switch m.Type {
			case api.MountTypeTmpfs:
				opts := []string{"noexec", "nosuid", "nodev", "rprivate"}
				if m.TmpfsOptions != nil {
					if m.TmpfsOptions.SizeBytes <= 0 {
						return errors.New("invalid tmpfs size give")
					}
					opts = append(opts, fmt.Sprintf("size=%d", m.TmpfsOptions.SizeBytes))
					opts = append(opts, fmt.Sprintf("mode=%o", m.TmpfsOptions.Mode))
				}
				if m.ReadOnly {
					opts = append(opts, "ro")
				} else {
					opts = append(opts, "rw")
				}

				s.Mounts = append(s.Mounts, specs.Mount{
					Destination: m.Target,
					Type:        "tmpfs",
					Source:      "tmpfs",
					Options:     opts,
				})

			case api.MountTypeVolume:
				return errors.Errorf("volume mounts not implemented, ignoring %v", m)

			case api.MountTypeBind:
				opts := []string{"rbind"}
				if m.ReadOnly {
					opts = append(opts, "ro")
				} else {
					opts = append(opts, "rw")
				}

				propagation := "rprivate"
				if m.BindOptions != nil {
					if p, ok := mountPropagationReverseMap[m.BindOptions.Propagation]; ok {
						propagation = p
					} else {
						log.G(ctx).Warningf("unknown bind mount propagation, using %q", propagation)
					}
				}
				opts = append(opts, propagation)

				s.Mounts = append(s.Mounts, specs.Mount{
					Destination: m.Target,
					Type:        "bind",
					Source:      m.Source,
					Options:     opts,
				})
			}
		}
		return nil
	}
}

func (c *containerAdapter) isPrepared() bool {
	return c.container != nil && c.task != nil
}

func (c *containerAdapter) prepare(ctx context.Context) error {
	if c.isPrepared() {
		return errors.New("adapter already prepared")
	}
	if c.image == nil {
		return errors.New("image has not been pulled")
	}

	specOpts := []oci.SpecOpts{
		oci.WithImageConfig(c.image),
		withMounts(ctx, c.spec.Mounts),
	}

	// spec.Process.Args is config.Entrypoint + config.Cmd at this
	// point from WithImageConfig above. If the ContainerSpec
	// specifies a Command then we can completely override. If it
	// does not then all we can do is append our Args and hope
	// they do not conflict.
	// TODO(ijc) Improve this
	if len(c.spec.Command) > 0 {
		args := append(c.spec.Command, c.spec.Args...)
		specOpts = append(specOpts, oci.WithProcessArgs(args...))
	} else {
		specOpts = append(specOpts, func(_ gocontext.Context, _ oci.Client, _ *containers.Container, s *specs.Spec) error {
			s.Process.Args = append(s.Process.Args, c.spec.Args...)
			return nil
		})
	}

	spec, err := oci.GenerateSpec(ctx, c.client, &containers.Container{}, specOpts...)
	if err != nil {
		return err
	}

	// TODO(ijc) Consider an addition to container library which
	// directly attaches stdin to /dev/null.
	if devNull == nil {
		if devNull, err = os.Open(os.DevNull); err != nil {
			return errors.Wrap(err, "opening null device")
		}
	}

	c.container, err = c.client.NewContainer(ctx, c.name,
		containerd.WithSpec(spec),
		containerd.WithNewSnapshot(c.name, c.image))
	if err != nil {
		return errors.Wrap(err, "creating container")
	}

	// TODO(ijc) support ControllerLogs interface.
	io := cio.NewIOWithTerminal(devNull, os.Stdout, os.Stderr, spec.Process.Terminal)

	c.task, err = c.container.NewTask(ctx, io)
	if err != nil {
		// Destroy the container we created above, but
		// propagate the original error.
		if err2 := c.container.Delete(ctx); err2 != nil {
			c.log(ctx).WithError(err2).Error("failed to delete container on prepare failure")
		}
		c.container = nil
		return errors.Wrap(err, "creating task")
	}

	return nil
}

func (c *containerAdapter) start(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}
	err := c.task.Start(ctx)
	return errors.Wrap(err, "starting")
}

func (c *containerAdapter) wait(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}
	status, err := c.task.Wait(ctx)
	if err != nil {
		return errors.Wrap(err, "waiting")
	}
	// Should update c.exitStatus or not?
	ec := <-status
	if ec.Error() != nil {
		return err
	}
	return makeExitError(ec.ExitCode(), "")
}

type status struct {
	ID         string
	Pid        uint32
	Status     containerd.Status
	ExitStatus error
}

func (c *containerAdapter) inspect(ctx context.Context) (status, error) {
	if !c.isPrepared() {
		return status{}, errAdapterNotPrepared
	}

	ts, err := c.task.Status(ctx)
	if err != nil {
		return status{}, err
	}
	s := status{
		ID:         c.container.ID(),
		Pid:        c.task.Pid(),
		Status:     ts,
		ExitStatus: c.exitStatus,
	}
	return s, nil
}

func (c *containerAdapter) shutdown(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	var (
		sig     syscall.Signal
		timeout = time.Duration(10 * time.Second)
		err     error
	)

	if c.spec.StopSignal != "" {
		if sig, err = signal.ParseSignal(c.spec.StopSignal); err != nil {
			sig = syscall.SIGTERM
			c.log(ctx).WithError(err).Errorf("unknown StopSignal, using %q", sig)
		}
	} else {
		sig = syscall.SIGTERM
		c.log(ctx).Infof("no StopSignal given, using %q", sig)
	}

	if c.spec.StopGracePeriod != nil {
		timeout, _ = gogotypes.DurationFromProto(c.spec.StopGracePeriod)
	}

	deleteErr := make(chan error, 1)
	deleteCtx, deleteCancel := context.WithCancel(ctx)
	defer deleteCancel()

	go func(ctx context.Context, ch chan error) {
		status, err := c.task.Delete(ctx)
		if err != nil {
			c.log(ctx).WithError(err).Debug("Task.Delete failed")
			ch <- err
		}
		c.log(ctx).Debugf("Task.Delete success, status=%d", status)
		ch <- makeExitError(status.ExitCode(), "")
	}(deleteCtx, deleteErr)

	c.log(ctx).Debugf("Killing task with %q signal", sig)
	if err := c.task.Kill(ctx, sig); err != nil {
		return errors.Wrapf(err, "killing task with %q", sig)
	}

	select {
	case c.exitStatus = <-deleteErr:
		return c.exitStatus
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(timeout):
		c.log(ctx).Infof("Task did not exit after %s", timeout)
		// Fall through
	}

	if sig == syscall.SIGKILL {
		// We've tried as hard as we can.
		return errors.New("task is unkillable")
	}

	// Bring out the big guns
	sig = syscall.SIGKILL
	c.log(ctx).Debugf("Killing task harder with %q signal", sig)
	if err := c.task.Kill(ctx, sig); err != nil {
		return errors.Wrapf(err, "killing task with %q", sig)
	}

	select {
	case c.exitStatus = <-deleteErr:
		return c.exitStatus
	case <-ctx.Done():
		return ctx.Err()
	}

}

func (c *containerAdapter) terminate(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	c.log(ctx).Debug("Terminate")
	return errors.New("terminate not implemented")
}

func (c *containerAdapter) remove(ctx context.Context) error {
	// Unlike most other entry points we don't use c.isPrepared
	// here so that we can clean up a container which was
	// partially reattached (via c.attach).
	if c.container == nil {
		return errAdapterNotPrepared
	}

	c.log(ctx).Debug("Remove")
	err := c.container.Delete(ctx)
	return errors.Wrap(err, "removing container")
}

func isContainerCreateNameConflict(err error) bool {
	// container ".*" already exists
	splits := strings.SplitN(err.Error(), "\"", 3)
	return splits[0] == "container " && splits[2] == " already exists"
}

func isUnknownContainer(err error) bool {
	return strings.Contains(err.Error(), "container does not exist")
}

// For sort.Sort
type mounts []api.Mount

// Len returns the number of mounts. Used in sorting.
func (m mounts) Len() int {
	return len(m)
}

// Less returns true if the number of parts (a/b/c would be 3 parts) in the
// mount indexed by parameter 1 is less than that of the mount indexed by
// parameter 2. Used in sorting.
func (m mounts) Less(i, j int) bool {
	return m.parts(i) < m.parts(j)
}

// Swap swaps two items in an array of mounts. Used in sorting
func (m mounts) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

// parts returns the number of parts in the destination of a mount. Used in sorting.
func (m mounts) parts(i int) int {
	return strings.Count(filepath.Clean(m[i].Target), string(os.PathSeparator))
}
<commit_msg>Update agent containerd adapter for changes in cio<commit_after>package containerd

import (
	gocontext "context"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"syscall"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/cio"
	"github.com/containerd/containerd/containers"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/oci"
	"github.com/docker/docker/pkg/signal"
	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/api/naming"
	"github.com/docker/swarmkit/log"
	gogotypes "github.com/gogo/protobuf/types"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/net/context"
)

var (
	devNull                    *os.File
	errAdapterNotPrepared      = errors.New("container adapter not prepared")
	mountPropagationReverseMap = map[api.Mount_BindOptions_MountPropagation]string{
		api.MountPropagationPrivate:  "private",
		api.MountPropagationRPrivate: "rprivate",
		api.MountPropagationShared:   "shared",
		api.MountPropagationRShared:  "rshared",
		api.MountPropagationRSlave:   "slave",
		api.MountPropagationSlave:    "rslave",
	}
)

// containerAdapter conducts remote operations for a container. All calls
// are mostly naked calls to the client API, seeded with information from
// containerConfig.
type containerAdapter struct {
	client     *containerd.Client
	spec       *api.ContainerSpec
	secrets    exec.SecretGetter
	name       string
	image      containerd.Image // Pulled image
	container  containerd.Container
	task       containerd.Task
	exitStatus error
}

func newContainerAdapter(client *containerd.Client, task *api.Task, secrets exec.SecretGetter) (*containerAdapter, error) {
	spec := task.Spec.GetContainer()
	if spec == nil {
		return nil, exec.ErrRuntimeUnsupported
	}

	c := &containerAdapter{
		client:  client,
		spec:    spec,
		secrets: secrets,
		name:    naming.Task(task),
	}

	if err := c.reattach(context.Background()); err != nil {
		return nil, err
	}

	return c, nil
}

// reattaches to an existing container. If the container is found but
// the task is missing then still succeeds, allowing subsequent use of
// c.delete()
func (c *containerAdapter) reattach(ctx context.Context) error {
	container, err := c.client.LoadContainer(ctx, c.name)
	if err != nil {
		if errdefs.IsNotFound(err) {
			c.log(ctx).Debug("reattach: container not found")
			return nil
		}

		return errors.Wrap(err, "reattach: loading container")
	}
	c.log(ctx).Debug("reattach: loaded container")
	c.container = container

	// TODO(ijc) Consider an addition to container library which
	// directly attaches stdin to /dev/null.
	if devNull == nil {
		if devNull, err = os.Open(os.DevNull); err != nil {
			return errors.Wrap(err, "reattach: opening null device")
		}
	}

	task, err := container.Task(ctx, cio.NewAttach(cio.WithStreams(devNull, os.Stdout, os.Stderr)))
	if err != nil {
		if errdefs.IsNotFound(err) {
			c.log(ctx).WithError(err).Info("reattach: no running task")
			return nil
		}
		return errors.Wrap(err, "reattach: reattaching task")
	}
	c.task = task
	c.log(ctx).Debug("reattach: successful")
	return nil
}

func (c *containerAdapter) log(ctx context.Context) *logrus.Entry {
	return log.G(ctx).WithFields(logrus.Fields{
		"container.id": c.name,
	})
}

func (c *containerAdapter) pullImage(ctx context.Context) error {
	image, err := c.client.Pull(ctx, c.spec.Image, containerd.WithPullUnpack)
	if err != nil {
		return errors.Wrap(err, "pulling container image")
	}
	c.image = image

	return nil
}

func withMounts(ctx context.Context, ms []api.Mount) oci.SpecOpts {
	sort.Sort(mounts(ms))

	return func(_ gocontext.Context, _ oci.Client, _ *containers.Container, s *specs.Spec) error {
		for _, m := range ms {
			if !filepath.IsAbs(m.Target) {
				return errors.Errorf("mount %s is not absolute", m.Target)
			}

			switch m.Type {
			case api.MountTypeTmpfs:
				opts := []string{"noexec", "nosuid", "nodev", "rprivate"}
				if m.TmpfsOptions != nil {
					if m.TmpfsOptions.SizeBytes <= 0 {
						return errors.New("invalid tmpfs size give")
					}
					opts = append(opts, fmt.Sprintf("size=%d", m.TmpfsOptions.SizeBytes))
					opts = append(opts, fmt.Sprintf("mode=%o", m.TmpfsOptions.Mode))
				}
				if m.ReadOnly {
					opts = append(opts, "ro")
				} else {
					opts = append(opts, "rw")
				}

				s.Mounts = append(s.Mounts, specs.Mount{
					Destination: m.Target,
					Type:        "tmpfs",
					Source:      "tmpfs",
					Options:     opts,
				})

			case api.MountTypeVolume:
				return errors.Errorf("volume mounts not implemented, ignoring %v", m)

			case api.MountTypeBind:
				opts := []string{"rbind"}
				if m.ReadOnly {
					opts = append(opts, "ro")
				} else {
					opts = append(opts, "rw")
				}

				propagation := "rprivate"
				if m.BindOptions != nil {
					if p, ok := mountPropagationReverseMap[m.BindOptions.Propagation]; ok {
						propagation = p
					} else {
						log.G(ctx).Warningf("unknown bind mount propagation, using %q", propagation)
					}
				}
				opts = append(opts, propagation)

				s.Mounts = append(s.Mounts, specs.Mount{
					Destination: m.Target,
					Type:        "bind",
					Source:      m.Source,
					Options:     opts,
				})
			}
		}
		return nil
	}
}

func (c *containerAdapter) isPrepared() bool {
	return c.container != nil && c.task != nil
}

func (c *containerAdapter) prepare(ctx context.Context) error {
	if c.isPrepared() {
		return errors.New("adapter already prepared")
	}
	if c.image == nil {
		return errors.New("image has not been pulled")
	}

	specOpts := []oci.SpecOpts{
		oci.WithImageConfig(c.image),
		withMounts(ctx, c.spec.Mounts),
	}

	// spec.Process.Args is config.Entrypoint + config.Cmd at this
	// point from WithImageConfig above. If the ContainerSpec
	// specifies a Command then we can completely override. If it
	// does not then all we can do is append our Args and hope
	// they do not conflict.
	// TODO(ijc) Improve this
	if len(c.spec.Command) > 0 {
		args := append(c.spec.Command, c.spec.Args...)
		specOpts = append(specOpts, oci.WithProcessArgs(args...))
	} else {
		specOpts = append(specOpts, func(_ gocontext.Context, _ oci.Client, _ *containers.Container, s *specs.Spec) error {
			s.Process.Args = append(s.Process.Args, c.spec.Args...)
			return nil
		})
	}

	spec, err := oci.GenerateSpec(ctx, c.client, &containers.Container{}, specOpts...)
	if err != nil {
		return err
	}

	// TODO(ijc) Consider an addition to container library which
	// directly attaches stdin to /dev/null.
	if devNull == nil {
		if devNull, err = os.Open(os.DevNull); err != nil {
			return errors.Wrap(err, "opening null device")
		}
	}

	c.container, err = c.client.NewContainer(ctx, c.name,
		containerd.WithSpec(spec),
		containerd.WithNewSnapshot(c.name, c.image))
	if err != nil {
		return errors.Wrap(err, "creating container")
	}

	// TODO(ijc) support ControllerLogs interface.
	opts := []cio.Opt{cio.WithStreams(devNull, os.Stdout, os.Stderr)}
	if spec.Process.Terminal {
		opts = append(opts, cio.WithTerminal)
	}
	c.task, err = c.container.NewTask(ctx, cio.NewCreator(opts...))
	if err != nil {
		// Destroy the container we created above, but
		// propagate the original error.
		if err2 := c.container.Delete(ctx); err2 != nil {
			c.log(ctx).WithError(err2).Error("failed to delete container on prepare failure")
		}
		c.container = nil
		return errors.Wrap(err, "creating task")
	}

	return nil
}

func (c *containerAdapter) start(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}
	err := c.task.Start(ctx)
	return errors.Wrap(err, "starting")
}

func (c *containerAdapter) wait(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}
	status, err := c.task.Wait(ctx)
	if err != nil {
		return errors.Wrap(err, "waiting")
	}
	// Should update c.exitStatus or not?
	ec := <-status
	if ec.Error() != nil {
		return err
	}
	return makeExitError(ec.ExitCode(), "")
}

type status struct {
	ID         string
	Pid        uint32
	Status     containerd.Status
	ExitStatus error
}

func (c *containerAdapter) inspect(ctx context.Context) (status, error) {
	if !c.isPrepared() {
		return status{}, errAdapterNotPrepared
	}

	ts, err := c.task.Status(ctx)
	if err != nil {
		return status{}, err
	}
	s := status{
		ID:         c.container.ID(),
		Pid:        c.task.Pid(),
		Status:     ts,
		ExitStatus: c.exitStatus,
	}
	return s, nil
}

func (c *containerAdapter) shutdown(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	var (
		sig     syscall.Signal
		timeout = time.Duration(10 * time.Second)
		err     error
	)

	if c.spec.StopSignal != "" {
		if sig, err = signal.ParseSignal(c.spec.StopSignal); err != nil {
			sig = syscall.SIGTERM
			c.log(ctx).WithError(err).Errorf("unknown StopSignal, using %q", sig)
		}
	} else {
		sig = syscall.SIGTERM
		c.log(ctx).Infof("no StopSignal given, using %q", sig)
	}

	if c.spec.StopGracePeriod != nil {
		timeout, _ = gogotypes.DurationFromProto(c.spec.StopGracePeriod)
	}

	deleteErr := make(chan error, 1)
	deleteCtx, deleteCancel := context.WithCancel(ctx)
	defer deleteCancel()

	go func(ctx context.Context, ch chan error) {
		status, err := c.task.Delete(ctx)
		if err != nil {
			c.log(ctx).WithError(err).Debug("Task.Delete failed")
			ch <- err
		}
		c.log(ctx).Debugf("Task.Delete success, status=%d", status)
		ch <- makeExitError(status.ExitCode(), "")
	}(deleteCtx, deleteErr)

	c.log(ctx).Debugf("Killing task with %q signal", sig)
	if err := c.task.Kill(ctx, sig); err != nil {
		return errors.Wrapf(err, "killing task with %q", sig)
	}

	select {
	case c.exitStatus = <-deleteErr:
		return c.exitStatus
	case <-ctx.Done():
		return ctx.Err()
	case <-time.After(timeout):
		c.log(ctx).Infof("Task did not exit after %s", timeout)
		// Fall through
	}

	if sig == syscall.SIGKILL {
		// We've tried as hard as we can.
		return errors.New("task is unkillable")
	}

	// Bring out the big guns
	sig = syscall.SIGKILL
	c.log(ctx).Debugf("Killing task harder with %q signal", sig)
	if err := c.task.Kill(ctx, sig); err != nil {
		return errors.Wrapf(err, "killing task with %q", sig)
	}

	select {
	case c.exitStatus = <-deleteErr:
		return c.exitStatus
	case <-ctx.Done():
		return ctx.Err()
	}

}

func (c *containerAdapter) terminate(ctx context.Context) error {
	if !c.isPrepared() {
		return errAdapterNotPrepared
	}

	c.log(ctx).Debug("Terminate")
	return errors.New("terminate not implemented")
}

func (c *containerAdapter) remove(ctx context.Context) error {
	// Unlike most other entry points we don't use c.isPrepared
	// here so that we can clean up a container which was
	// partially reattached (via c.attach).
	if c.container == nil {
		return errAdapterNotPrepared
	}

	c.log(ctx).Debug("Remove")
	err := c.container.Delete(ctx)
	return errors.Wrap(err, "removing container")
}

func isContainerCreateNameConflict(err error) bool {
	// container ".*" already exists
	splits := strings.SplitN(err.Error(), "\"", 3)
	return splits[0] == "container " && splits[2] == " already exists"
}

func isUnknownContainer(err error) bool {
	return strings.Contains(err.Error(), "container does not exist")
}

// For sort.Sort
type mounts []api.Mount

// Len returns the number of mounts. Used in sorting.
func (m mounts) Len() int {
	return len(m)
}

// Less returns true if the number of parts (a/b/c would be 3 parts) in the
// mount indexed by parameter 1 is less than that of the mount indexed by
// parameter 2. Used in sorting.
func (m mounts) Less(i, j int) bool {
	return m.parts(i) < m.parts(j)
}

// Swap swaps two items in an array of mounts. Used in sorting
func (m mounts) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

// parts returns the number of parts in the destination of a mount. Used in sorting.
func (m mounts) parts(i int) int {
	return strings.Count(filepath.Clean(m[i].Target), string(os.PathSeparator))
}
<|endoftext|>


<commit_before>package app

import (
	"bytes"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"syscall"
	"time"

	"github.com/koding/klient/Godeps/_workspace/src/github.com/hashicorp/go-version"
	"github.com/koding/klient/Godeps/_workspace/src/github.com/inconshreveable/go-update"
	"github.com/koding/klient/Godeps/_workspace/src/github.com/koding/kite"
	"github.com/koding/klient/Godeps/_workspace/src/github.com/mitchellh/osext"
	"github.com/koding/klient/protocol"
)

type Updater struct {
	Endpoint       string
	Interval       time.Duration
	CurrentVersion string
	Log            kite.Logger
}

type UpdateData struct {
	KlientURL string
}

var AuthenticatedUser = "koding"

func (u *Updater) ServeKite(r *kite.Request) (interface{}, error) {
	if r.Username != AuthenticatedUser {
		return nil, fmt.Errorf("Not authenticated to make an update: %s", r.Username)
	}

	go func() {
		r.LocalKite.Log.Info("klient.Update is called. Updating binary via latest version")

		updatingState(true)
		if err := u.checkAndUpdate(); err != nil {
			u.Log.Warning("klient.update: %s", err)
		}
		updatingState(false)
	}()

	return true, nil
}

func updatingState(state bool) {
	updatingMu.Lock()
	updating = state
	updatingMu.Unlock()
}

func (u *Updater) checkAndUpdate() error {
	l, err := u.latestVersion()
	if err != nil {
		return err
	}

	latestVer := "0.1." + l
	latest, err := version.NewVersion(latestVer)
	if err != nil {
		return err
	}

	current, err := version.NewVersion(u.CurrentVersion)
	if err != nil {
		return err
	}

	if !current.LessThan(latest) {
		// current running binary version is equal or greater than what we fetched, so return we don't need to update
		return nil
	}

	u.Log.Info("Current version: %s is old. Going to update to: %s", u.CurrentVersion, latestVer)

	basePath := "https://s3.amazonaws.com/koding-klient/" + protocol.Environment + "/latest"
	latestKlientURL := basePath + "/klient-" + latestVer + ".gz"

	return u.updateBinary(latestKlientURL)
}

func (u *Updater) updateBinary(url string) error {
	updater := update.New()
	err := updater.CanUpdate()
	if err != nil {
		return err
	}

	self, err := osext.Executable()
	if err != nil {
		return err
	}

	u.Log.Info("Going to update binary at: %s", self)
	bin, err := u.fetch(url)
	if err != nil {
		return err
	}

	u.Log.Info("Replacing new binary with the old one.")
	err, errRecover := updater.FromStream(bytes.NewBuffer(bin))
	if err != nil {
		if errRecover != nil {
			return errRecover
		}

		return err
	}

	env := os.Environ()

	// TODO: os.Args[1:] should come also from the endpoint if the new binary
	// has a different flag!
	args := []string{self}
	args = append(args, os.Args[1:]...)

	// we need to call it here now too, because syscall.Exec will prevent to
	// call the defer that we've defined in the beginning.

	u.Log.Info("Updating was successfull. Replacing current process with args: %v\n=====> RESTARTING...\n\n", args)

	execErr := syscall.Exec(self, args, env)
	if execErr != nil {
		return err
	}

	return nil
}

func (u *Updater) latestVersion() (string, error) {
	resp, err := http.Get(u.Endpoint)
	if err != nil {
		u.Log.Debug("Getting latest version from %s", u.Endpoint)
		return "", err
	}
	defer resp.Body.Close()

	latest, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(latest)), nil
}

func (u *Updater) fetch(url string) ([]byte, error) {
	u.Log.Info("Fetching binary %s", url)
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("bad http status from %s: %v", url, resp.Status)
	}

	buf := new(bytes.Buffer)
	gz, err := gzip.NewReader(resp.Body)
	if err != nil {
		return nil, err
	}
	defer gz.Close()

	if _, err = io.Copy(buf, gz); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// Run runs the updater in the background for the interval of updater interval.
func (u *Updater) Run() {
	u.Log.Info("Starting Updater with following options:\n\tinterval of: %s\n\tendpoint: %s",
		u.Interval, u.Endpoint)

	for _ = range time.Tick(u.Interval) {
		updatingState(true)

		if err := u.checkAndUpdate(); err != nil {
			u.Log.Warning("Self-update: %s", err)
		}

		updatingState(false)
	}
}

// haveFreeSpace checks whether the disk has free space to provide the update
func haveFreeSpace() error {
	stat := new(syscall.Statfs_t)

	if err := syscall.Statfs("/", stat); err != nil {
		return err
	}

	bsize := stat.Bsize / 512

	free := (uint64(stat.Bfree) * uint64(bsize)) >> 1
	freeInMB := free / 1024

	if freeInMB < 100 {
		return errors.New("No enough space to upgrade klient")
	}

	return nil
}
<commit_msg>update: check size before we update<commit_after>package app

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"syscall"
	"time"

	"github.com/koding/klient/Godeps/_workspace/src/github.com/hashicorp/go-version"
	"github.com/koding/klient/Godeps/_workspace/src/github.com/inconshreveable/go-update"
	"github.com/koding/klient/Godeps/_workspace/src/github.com/koding/kite"
	"github.com/koding/klient/Godeps/_workspace/src/github.com/mitchellh/osext"
	"github.com/koding/klient/protocol"
)

type Updater struct {
	Endpoint       string
	Interval       time.Duration
	CurrentVersion string
	Log            kite.Logger
}

type UpdateData struct {
	KlientURL string
}

var AuthenticatedUser = "koding"

func (u *Updater) ServeKite(r *kite.Request) (interface{}, error) {
	if r.Username != AuthenticatedUser {
		return nil, fmt.Errorf("Not authenticated to make an update: %s", r.Username)
	}

	go func() {
		r.LocalKite.Log.Info("klient.Update is called. Updating binary via latest version")

		updatingState(true)
		if err := u.checkAndUpdate(); err != nil {
			u.Log.Warning("klient.update: %s", err)
		}
		updatingState(false)
	}()

	return true, nil
}

func updatingState(state bool) {
	updatingMu.Lock()
	updating = state
	updatingMu.Unlock()
}

func (u *Updater) checkAndUpdate() error {
	if err := hasFreeSpace(100); err != nil {
		return err
	}

	l, err := u.latestVersion()
	if err != nil {
		return err
	}

	latestVer := "0.1." + l
	latest, err := version.NewVersion(latestVer)
	if err != nil {
		return err
	}

	current, err := version.NewVersion(u.CurrentVersion)
	if err != nil {
		return err
	}

	if !current.LessThan(latest) {
		// current running binary version is equal or greater than what we fetched, so return we don't need to update
		return nil
	}

	u.Log.Info("Current version: %s is old. Going to update to: %s", u.CurrentVersion, latestVer)

	basePath := "https://s3.amazonaws.com/koding-klient/" + protocol.Environment + "/latest"
	latestKlientURL := basePath + "/klient-" + latestVer + ".gz"

	return u.updateBinary(latestKlientURL)
}

func (u *Updater) updateBinary(url string) error {
	updater := update.New()
	err := updater.CanUpdate()
	if err != nil {
		return err
	}

	self, err := osext.Executable()
	if err != nil {
		return err
	}

	u.Log.Info("Going to update binary at: %s", self)
	bin, err := u.fetch(url)
	if err != nil {
		return err
	}

	u.Log.Info("Replacing new binary with the old one.")
	err, errRecover := updater.FromStream(bytes.NewBuffer(bin))
	if err != nil {
		if errRecover != nil {
			return errRecover
		}

		return err
	}

	env := os.Environ()

	// TODO: os.Args[1:] should come also from the endpoint if the new binary
	// has a different flag!
	args := []string{self}
	args = append(args, os.Args[1:]...)

	// we need to call it here now too, because syscall.Exec will prevent to
	// call the defer that we've defined in the beginning.

	u.Log.Info("Updating was successfull. Replacing current process with args: %v\n=====> RESTARTING...\n\n", args)

	execErr := syscall.Exec(self, args, env)
	if execErr != nil {
		return err
	}

	return nil
}

func (u *Updater) latestVersion() (string, error) {
	resp, err := http.Get(u.Endpoint)
	if err != nil {
		u.Log.Debug("Getting latest version from %s", u.Endpoint)
		return "", err
	}
	defer resp.Body.Close()

	latest, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(latest)), nil
}

func (u *Updater) fetch(url string) ([]byte, error) {
	u.Log.Info("Fetching binary %s", url)
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("bad http status from %s: %v", url, resp.Status)
	}

	buf := new(bytes.Buffer)
	gz, err := gzip.NewReader(resp.Body)
	if err != nil {
		return nil, err
	}
	defer gz.Close()

	if _, err = io.Copy(buf, gz); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// Run runs the updater in the background for the interval of updater interval.
func (u *Updater) Run() {
	u.Log.Info("Starting Updater with following options:\n\tinterval of: %s\n\tendpoint: %s",
		u.Interval, u.Endpoint)

	for _ = range time.Tick(u.Interval) {
		updatingState(true)

		if err := u.checkAndUpdate(); err != nil {
			u.Log.Warning("Self-update: %s", err)
		}

		updatingState(false)
	}
}

// hasFreeSpace checks whether the disk has free space to provide the update.
// If free space is lower than then the given mustHas it returns an error.
func hasFreeSpace(mustHas uint64) error {
	stat := new(syscall.Statfs_t)

	if err := syscall.Statfs("/", stat); err != nil {
		return err
	}

	bsize := stat.Bsize / 512

	free := (uint64(stat.Bfree) * uint64(bsize)) >> 1
	freeInMB := free / 1024

	if freeInMB < mustHas {
		return fmt.Errorf("No enough space to upgrade klient. Need '%d'. Have: '%d'",
			mustHas, freeInMB)
	}

	return nil
}
<|endoftext|>