<commit_before>package lzma

import (
	"io"

	"github.com/uli-go/xz/xlog"
)

// Writer supports the LZMA compression of a file.
//
// Using an arithmetic coder it cannot support flushing. A writer must be
// closed.
type Writer struct {
	w         io.Writer
	ow        *opWriter
	props     Properties
	unpackLen int64
	dict      *writerDict
	// hash table for four-byte sequences
	t4 *hashTable
}

// Default defines the properties used by NewWriter.
var Default = Properties{
	LC:      3,
	LP:      0,
	PB:      2,
	DictLen: 1 << 12}

// NewWriter creates a new writer. The properties stored in Default will be
// used.
//
// Don't forget to call Close() for the writer after all data has been written.
func NewWriter(w io.Writer) (*Writer, error) {
	return NewWriterP(w, Default)
}

// NewWriterP creates a new writer with the given Properties.
//
// Don't forget to call Close() for the writer after all data has been written.
func NewWriterP(w io.Writer, p Properties) (*Writer, error) {
	if w == nil {
		return nil, newError("can't support a nil writer")
	}
	var err error
	if err = verifyProperties(&p); err != nil {
		return nil, err
	}
	length := p.Len
	if length == 0 {
		length = noUnpackLen
		p.EOS = true
	}
	lw := &Writer{
		w:         w,
		props:     p,
		unpackLen: length,
	}
	lw.dict, err = newWriterDict(int(p.DictLen), defaultBufferLen)
	if err != nil {
		return nil, err
	}
	lw.ow, err = newOpWriter(w, &p, lw.dict)
	if err != nil {
		return nil, err
	}
	lw.t4, err = newHashTable(int(p.DictLen), 4)
	if err != nil {
		return nil, err
	}
	if err = lw.writeHeader(); err != nil {
		return nil, err
	}
	return lw, nil
}

// Properties returns the properties of the LZMA writer.
func (lw *Writer) Properties() Properties {
	return lw.props
}

// putUint64LE puts the uint64 value into the byte slice as little endian
// value. The byte slice b must have at least place for 8 bytes.
func putUint64LE(b []byte, x uint64) {
	b[0] = byte(x)
	b[1] = byte(x >> 8)
	b[2] = byte(x >> 16)
	b[3] = byte(x >> 24)
	b[4] = byte(x >> 32)
	b[5] = byte(x >> 40)
	b[6] = byte(x >> 48)
	b[7] = byte(x >> 56)
}

// writeHeader writes the classic header into the output writer.
func (lw *Writer) writeHeader() error {
	err := writeProperties(lw.w, &lw.props)
	if err != nil {
		return err
	}
	b := make([]byte, 8)
	putUint64LE(b, uint64(lw.unpackLen))
	_, err = lw.w.Write(b)
	return err
}

// Write moves data into the internal buffer and triggers its compression.
func (lw *Writer) Write(p []byte) (n int, err error) {
	for n < len(p) {
		k, err := lw.dict.Write(p[n:])
		n += k
		if err != nil && err != errAgain {
			return n, err
		}
		if err = lw.process(0); err != nil {
			return n, err
		}
	}
	return n, nil
}

// Close terminates the LZMA stream. It doesn't close the underlying writer
// though.
func (lw *Writer) Close() error {
	var err error
	if err = lw.process(allData); err != nil {
		return err
	}
	if lw.props.EOS {
		if err = lw.ow.writeEOS(); err != nil {
			return err
		}
	}
	if err = lw.ow.Close(); err != nil {
		return err
	}
	return lw.dict.Close()
}

// The allData flag tells the process method that all data must be processed.
const allData = 1

// indicates an empty buffer
var errEmptyBuf = newError("empty buffer")

// potentialOffsets creates a list of potential offsets for matches.
func (lw *Writer) potentialOffsets(p []byte) []int64 {
	head := lw.dict.Offset()
	start := lw.dict.start
	offs := make([]int64, 0, 32)
	// add potential offsets with highest priority at the top
	for i := 1; i < 11; i++ {
		// distance 1 to 8
		off := head - int64(i)
		if start <= off {
			offs = append(offs, off)
		}
	}
	if len(p) == 4 {
		// distances from the hash table
		offs = append(offs, lw.t4.Offsets(p)...)
	}
	for i := 3; i >= 0; i-- {
		// distances from the repetition for length less than 4
		dist := int64(lw.ow.rep[i]) + minDistance
		off := head - dist
		if start <= off {
			offs = append(offs, off)
		}
	}
	return offs
}

// errNoMatch indicates that no match could be found
var errNoMatch = newError("no match found")

// bestMatch finds the best match for the given offsets.
//
// TODO: compare all possible commands for compressed bits per encoded bits.
func (lw *Writer) bestMatch(offsets []int64) (m match, err error) {
	// creates a match for 1
	head := lw.dict.Offset()
	off := int64(-1)
	length := 0
	for i := len(offsets) - 1; i >= 0; i-- {
		n := lw.dict.EqualBytes(head, offsets[i], maxLength)
		if n > length {
			off, length = offsets[i], n
		}
	}
	if off < 0 {
		err = errNoMatch
		return
	}
	if length == 1 {
		dist := int64(lw.ow.rep[0]) + minDistance
		offRep0 := head - dist
		if off != offRep0 {
			err = errNoMatch
			return
		}
	}
	return match{distance: head - off, length: length}, nil
}

// findOp finds an operation for the head of the dictionary.
func (lw *Writer) findOp() (op operation, err error) {
	p := make([]byte, 4)
	n, err := lw.dict.PeekHead(p)
	if err != nil && err != errAgain && err != io.EOF {
		return nil, err
	}
	if n <= 0 {
		if n < 0 {
			panic("strange n")
		}
		return nil, errEmptyBuf
	}
	offs := lw.potentialOffsets(p[:n])
	m, err := lw.bestMatch(offs)
	if err == errNoMatch {
		return lit{b: p[0]}, nil
	}
	if err != nil {
		return nil, err
	}
	return m, nil
}

// discardOp advances the head of the dictionary and writes the the bytes into
// the hash table.
func (lw *Writer) discardOp(op operation) error {
	n, err := lw.dict.Copy(lw.t4, op.Len())
	if err != nil {
		return err
	}
	if n < op.Len() {
		return errAgain
	}
	return nil
}

// process encodes the data written into the dictionary buffer. The allData
// flag requires all data remaining in the buffer to be encoded.
func (lw *Writer) process(flags int) error {
	var lowMark int
	if flags&allData == 0 {
		lowMark = maxLength - 1
	}
	for lw.dict.Readable() > lowMark {
		op, err := lw.findOp()
		if err != nil {
			xlog.Printf(debug, "findOp error %s\n", err)
			return err
		}
		if err = lw.ow.WriteOp(op); err != nil {
			return err
		}
		xlog.Printf(debug, "op %s", op)
		if err = lw.discardOp(op); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>lzma: actually handling unpackLen in writer<commit_after>package lzma

import (
	"io"

	"github.com/uli-go/xz/xlog"
)

// Writer supports the LZMA compression of a file.
//
// Using an arithmetic coder it cannot support flushing. A writer must be
// closed.
type Writer struct {
	w         io.Writer
	ow        *opWriter
	props     Properties
	unpackLen int64
	dict      *writerDict
	// hash table for four-byte sequences
	t4 *hashTable
}

// Default defines the properties used by NewWriter.
var Default = Properties{
	LC:      3,
	LP:      0,
	PB:      2,
	DictLen: 1 << 12}

// NewWriter creates a new writer. The properties stored in Default will be
// used.
//
// Don't forget to call Close() for the writer after all data has been written.
func NewWriter(w io.Writer) (*Writer, error) {
	return NewWriterP(w, Default)
}

// NewWriterP creates a new writer with the given Properties.
//
// Don't forget to call Close() for the writer after all data has been written.
func NewWriterP(w io.Writer, p Properties) (*Writer, error) {
	if w == nil {
		return nil, newError("can't support a nil writer")
	}
	var err error
	if err = verifyProperties(&p); err != nil {
		return nil, err
	}
	length := p.Len
	if length == 0 {
		length = noUnpackLen
		p.EOS = true
	}
	lw := &Writer{
		w:         w,
		props:     p,
		unpackLen: length,
	}
	lw.dict, err = newWriterDict(int(p.DictLen), defaultBufferLen)
	if err != nil {
		return nil, err
	}
	lw.ow, err = newOpWriter(w, &p, lw.dict)
	if err != nil {
		return nil, err
	}
	lw.t4, err = newHashTable(int(p.DictLen), 4)
	if err != nil {
		return nil, err
	}
	if err = lw.writeHeader(); err != nil {
		return nil, err
	}
	return lw, nil
}

// Properties returns the properties of the LZMA writer.
func (lw *Writer) Properties() Properties {
	return lw.props
}

// putUint64LE puts the uint64 value into the byte slice as little endian
// value. The byte slice b must have at least place for 8 bytes.
func putUint64LE(b []byte, x uint64) {
	b[0] = byte(x)
	b[1] = byte(x >> 8)
	b[2] = byte(x >> 16)
	b[3] = byte(x >> 24)
	b[4] = byte(x >> 32)
	b[5] = byte(x >> 40)
	b[6] = byte(x >> 48)
	b[7] = byte(x >> 56)
}

// writeHeader writes the classic header into the output writer.
func (lw *Writer) writeHeader() error {
	err := writeProperties(lw.w, &lw.props)
	if err != nil {
		return err
	}
	b := make([]byte, 8)
	putUint64LE(b, uint64(lw.unpackLen))
	_, err = lw.w.Write(b)
	return err
}

// Write moves data into the internal buffer and triggers its compression.
func (lw *Writer) Write(p []byte) (n int, err error) {
	end := lw.dict.end + int64(len(p))
	if end < 0 {
		panic("end counter overflow")
	}
	var rerr error
	if lw.unpackLen != noUnpackLen && end > lw.unpackLen {
		m := lw.unpackLen - end
		p = p[:m]
		rerr = newError("write exceeds unpackLen")
	}
	for n < len(p) {
		k, err := lw.dict.Write(p[n:])
		n += k
		if err != nil && err != errAgain {
			return n, err
		}
		if err = lw.process(0); err != nil {
			return n, err
		}
	}
	return n, rerr
}

// Close terminates the LZMA stream. It doesn't close the underlying writer
// though.
func (lw *Writer) Close() error {
	var err error
	if err = lw.process(allData); err != nil {
		return err
	}
	if lw.props.EOS {
		if err = lw.ow.writeEOS(); err != nil {
			return err
		}
	}
	if err = lw.ow.Close(); err != nil {
		return err
	}
	return lw.dict.Close()
}

// The allData flag tells the process method that all data must be processed.
const allData = 1

// indicates an empty buffer
var errEmptyBuf = newError("empty buffer")

// potentialOffsets creates a list of potential offsets for matches.
func (lw *Writer) potentialOffsets(p []byte) []int64 {
	head := lw.dict.Offset()
	start := lw.dict.start
	offs := make([]int64, 0, 32)
	// add potential offsets with highest priority at the top
	for i := 1; i < 11; i++ {
		// distance 1 to 8
		off := head - int64(i)
		if start <= off {
			offs = append(offs, off)
		}
	}
	if len(p) == 4 {
		// distances from the hash table
		offs = append(offs, lw.t4.Offsets(p)...)
	}
	for i := 3; i >= 0; i-- {
		// distances from the repetition for length less than 4
		dist := int64(lw.ow.rep[i]) + minDistance
		off := head - dist
		if start <= off {
			offs = append(offs, off)
		}
	}
	return offs
}

// errNoMatch indicates that no match could be found
var errNoMatch = newError("no match found")

// bestMatch finds the best match for the given offsets.
//
// TODO: compare all possible commands for compressed bits per encoded bits.
func (lw *Writer) bestMatch(offsets []int64) (m match, err error) {
	// creates a match for 1
	head := lw.dict.Offset()
	off := int64(-1)
	length := 0
	for i := len(offsets) - 1; i >= 0; i-- {
		n := lw.dict.EqualBytes(head, offsets[i], maxLength)
		if n > length {
			off, length = offsets[i], n
		}
	}
	if off < 0 {
		err = errNoMatch
		return
	}
	if length == 1 {
		dist := int64(lw.ow.rep[0]) + minDistance
		offRep0 := head - dist
		if off != offRep0 {
			err = errNoMatch
			return
		}
	}
	return match{distance: head - off, length: length}, nil
}

// findOp finds an operation for the head of the dictionary.
func (lw *Writer) findOp() (op operation, err error) {
	p := make([]byte, 4)
	n, err := lw.dict.PeekHead(p)
	if err != nil && err != errAgain && err != io.EOF {
		return nil, err
	}
	if n <= 0 {
		if n < 0 {
			panic("strange n")
		}
		return nil, errEmptyBuf
	}
	offs := lw.potentialOffsets(p[:n])
	m, err := lw.bestMatch(offs)
	if err == errNoMatch {
		return lit{b: p[0]}, nil
	}
	if err != nil {
		return nil, err
	}
	return m, nil
}

// discardOp advances the head of the dictionary and writes the the bytes into
// the hash table.
func (lw *Writer) discardOp(op operation) error {
	n, err := lw.dict.Copy(lw.t4, op.Len())
	if err != nil {
		return err
	}
	if n < op.Len() {
		return errAgain
	}
	return nil
}

// process encodes the data written into the dictionary buffer. The allData
// flag requires all data remaining in the buffer to be encoded.
func (lw *Writer) process(flags int) error {
	var lowMark int
	if flags&allData == 0 {
		lowMark = maxLength - 1
	}
	for lw.dict.Readable() > lowMark {
		op, err := lw.findOp()
		if err != nil {
			xlog.Printf(debug, "findOp error %s\n", err)
			return err
		}
		if err = lw.ow.WriteOp(op); err != nil {
			return err
		}
		xlog.Printf(debug, "op %s", op)
		if err = lw.discardOp(op); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/terraform/flatmap"
	"github.com/hashicorp/terraform/helper/config"
	"github.com/hashicorp/terraform/helper/diff"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/route53"
)

func resource_aws_r53_record_validation() *config.Validator {
	return &config.Validator{
		Required: []string{
			"zone_id",
			"name",
			"type",
			"ttl",
			"records.*",
		},
	}
}

func resource_aws_r53_record_create(
	s *terraform.InstanceState,
	d *terraform.InstanceDiff,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	conn := p.route53

	// Merge the diff into the state so that we have all the attributes
	// properly.
	rs := s.MergeDiff(d)

	// Get the record
	rec, err := resource_aws_r53_build_record_set(rs)
	if err != nil {
		return rs, err
	}

	// Create the new records. We abuse StateChangeConf for this to
	// retry for us since Route53 sometimes returns errors about another
	// operation happening at the same time.
	req := &route53.ChangeResourceRecordSetsRequest{
		Comment: "Managed by Terraform",
		Changes: []route53.Change{
			route53.Change{
				Action: "UPSERT",
				Record: *rec,
			},
		},
	}
	zone := rs.Attributes["zone_id"]
	log.Printf("[DEBUG] Creating resource records for zone: %s, name: %s",
		zone, rs.Attributes["name"])
	wait := resource.StateChangeConf{
		Pending:    []string{"rejected"},
		Target:     "accepted",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			resp, err := conn.ChangeResourceRecordSets(zone, req)
			if err != nil {
				if strings.Contains(err.Error(), "PriorRequestNotComplete") {
					// There is some pending operation, so just retry
					// in a bit.
					return nil, "rejected", nil
				}

				return nil, "failure", err
			}

			return resp.ChangeInfo, "accepted", nil
		},
	}
	respRaw, err := wait.WaitForState()
	if err != nil {
		return rs, err
	}
	changeInfo := respRaw.(route53.ChangeInfo)

	// Generate an ID
	rs.ID = fmt.Sprintf("%s_%s_%s", zone, rs.Attributes["name"], rs.Attributes["type"])

	// Wait until we are done
	wait = resource.StateChangeConf{
		Delay:      30 * time.Second,
		Pending:    []string{"PENDING"},
		Target:     "INSYNC",
		Timeout:    10 * time.Minute,
		MinTimeout: 5 * time.Second,
		Refresh: func() (result interface{}, state string, err error) {
			return resource_aws_r53_wait(conn, changeInfo.ID)
		},
	}
	_, err = wait.WaitForState()
	if err != nil {
		return rs, err
	}
	return rs, nil
}

func resource_aws_r53_build_record_set(s *terraform.InstanceState) (*route53.ResourceRecordSet, error) {
	// Parse the TTL
	ttl, err := strconv.ParseInt(s.Attributes["ttl"], 10, 32)
	if err != nil {
		return nil, err
	}

	// Expand the records
	recRaw := flatmap.Expand(s.Attributes, "records")
	var records []string
	for _, raw := range recRaw.([]interface{}) {
		records = append(records, raw.(string))
	}

	rec := &route53.ResourceRecordSet{
		Name:    s.Attributes["name"],
		Type:    s.Attributes["type"],
		TTL:     int(ttl),
		Records: records,
	}
	return rec, nil
}

func resource_aws_r53_record_destroy(
	s *terraform.InstanceState,
	meta interface{}) error {
	p := meta.(*ResourceProvider)
	conn := p.route53

	// Get the record
	rec, err := resource_aws_r53_build_record_set(s)
	if err != nil {
		return err
	}

	// Create the new records
	req := &route53.ChangeResourceRecordSetsRequest{
		Comment: "Deleted by Terraform",
		Changes: []route53.Change{
			route53.Change{
				Action: "DELETE",
				Record: *rec,
			},
		},
	}
	zone := s.Attributes["zone_id"]
	log.Printf("[DEBUG] Deleting resource records for zone: %s, name: %s",
		zone, s.Attributes["name"])
	wait := resource.StateChangeConf{
		Pending:    []string{"rejected"},
		Target:     "accepted",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			_, err := conn.ChangeResourceRecordSets(zone, req)
			if err != nil {
				if strings.Contains(err.Error(), "PriorRequestNotComplete") {
					// There is some pending operation, so just retry
					// in a bit.
					return nil, "rejected", nil
				}

				return nil, "failure", err
			}

			return nil, "accepted", nil
		},
	}
	if _, err := wait.WaitForState(); err != nil {
		return err
	}

	return nil
}

func resource_aws_r53_record_refresh(
	s *terraform.InstanceState,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	conn := p.route53

	zone := s.Attributes["zone_id"]
	lopts := &route53.ListOpts{
		Name: s.Attributes["name"],
		Type: s.Attributes["type"],
	}
	resp, err := conn.ListResourceRecordSets(zone, lopts)
	if err != nil {
		return s, err
	}

	// Scan for a matching record
	found := false
	for _, record := range resp.Records {
		if route53.FQDN(record.Name) != route53.FQDN(lopts.Name) {
			continue
		}
		if strings.ToUpper(record.Type) != strings.ToUpper(lopts.Type) {
			continue
		}

		found = true
		resource_aws_r53_record_update_state(s, &record)
		break
	}
	if !found {
		s.ID = ""
	}
	return s, nil
}

func resource_aws_r53_record_update_state(
	s *terraform.InstanceState,
	rec *route53.ResourceRecordSet) {

	flatRec := flatmap.Flatten(map[string]interface{}{
		"records": rec.Records,
	})
	for k, v := range flatRec {
		s.Attributes[k] = v
	}

	s.Attributes["ttl"] = strconv.FormatInt(int64(rec.TTL), 10)
}

func resource_aws_r53_record_diff(
	s *terraform.InstanceState,
	c *terraform.ResourceConfig,
	meta interface{}) (*terraform.InstanceDiff, error) {
	b := &diff.ResourceBuilder{
		Attrs: map[string]diff.AttrType{
			"zone_id": diff.AttrTypeCreate,
			"name":    diff.AttrTypeCreate,
			"type":    diff.AttrTypeCreate,
			"ttl":     diff.AttrTypeUpdate,
			"records": diff.AttrTypeUpdate,
		},
	}
	return b.Diff(s, c)
}
<commit_msg>providers/aws: not an error if route53 record is already destroy [GH-342]<commit_after>package aws

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/hashicorp/terraform/flatmap"
	"github.com/hashicorp/terraform/helper/config"
	"github.com/hashicorp/terraform/helper/diff"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/route53"
)

func resource_aws_r53_record_validation() *config.Validator {
	return &config.Validator{
		Required: []string{
			"zone_id",
			"name",
			"type",
			"ttl",
			"records.*",
		},
	}
}

func resource_aws_r53_record_create(
	s *terraform.InstanceState,
	d *terraform.InstanceDiff,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	conn := p.route53

	// Merge the diff into the state so that we have all the attributes
	// properly.
	rs := s.MergeDiff(d)

	// Get the record
	rec, err := resource_aws_r53_build_record_set(rs)
	if err != nil {
		return rs, err
	}

	// Create the new records. We abuse StateChangeConf for this to
	// retry for us since Route53 sometimes returns errors about another
	// operation happening at the same time.
	req := &route53.ChangeResourceRecordSetsRequest{
		Comment: "Managed by Terraform",
		Changes: []route53.Change{
			route53.Change{
				Action: "UPSERT",
				Record: *rec,
			},
		},
	}
	zone := rs.Attributes["zone_id"]
	log.Printf("[DEBUG] Creating resource records for zone: %s, name: %s",
		zone, rs.Attributes["name"])
	wait := resource.StateChangeConf{
		Pending:    []string{"rejected"},
		Target:     "accepted",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			resp, err := conn.ChangeResourceRecordSets(zone, req)
			if err != nil {
				if strings.Contains(err.Error(), "PriorRequestNotComplete") {
					// There is some pending operation, so just retry
					// in a bit.
					return nil, "rejected", nil
				}

				return nil, "failure", err
			}

			return resp.ChangeInfo, "accepted", nil
		},
	}
	respRaw, err := wait.WaitForState()
	if err != nil {
		return rs, err
	}
	changeInfo := respRaw.(route53.ChangeInfo)

	// Generate an ID
	rs.ID = fmt.Sprintf("%s_%s_%s", zone, rs.Attributes["name"], rs.Attributes["type"])

	// Wait until we are done
	wait = resource.StateChangeConf{
		Delay:      30 * time.Second,
		Pending:    []string{"PENDING"},
		Target:     "INSYNC",
		Timeout:    10 * time.Minute,
		MinTimeout: 5 * time.Second,
		Refresh: func() (result interface{}, state string, err error) {
			return resource_aws_r53_wait(conn, changeInfo.ID)
		},
	}
	_, err = wait.WaitForState()
	if err != nil {
		return rs, err
	}
	return rs, nil
}

func resource_aws_r53_build_record_set(s *terraform.InstanceState) (*route53.ResourceRecordSet, error) {
	// Parse the TTL
	ttl, err := strconv.ParseInt(s.Attributes["ttl"], 10, 32)
	if err != nil {
		return nil, err
	}

	// Expand the records
	recRaw := flatmap.Expand(s.Attributes, "records")
	var records []string
	for _, raw := range recRaw.([]interface{}) {
		records = append(records, raw.(string))
	}

	rec := &route53.ResourceRecordSet{
		Name:    s.Attributes["name"],
		Type:    s.Attributes["type"],
		TTL:     int(ttl),
		Records: records,
	}
	return rec, nil
}

func resource_aws_r53_record_destroy(
	s *terraform.InstanceState,
	meta interface{}) error {
	p := meta.(*ResourceProvider)
	conn := p.route53

	// Get the record
	rec, err := resource_aws_r53_build_record_set(s)
	if err != nil {
		return err
	}

	// Create the new records
	req := &route53.ChangeResourceRecordSetsRequest{
		Comment: "Deleted by Terraform",
		Changes: []route53.Change{
			route53.Change{
				Action: "DELETE",
				Record: *rec,
			},
		},
	}
	zone := s.Attributes["zone_id"]
	log.Printf("[DEBUG] Deleting resource records for zone: %s, name: %s",
		zone, s.Attributes["name"])
	wait := resource.StateChangeConf{
		Pending:    []string{"rejected"},
		Target:     "accepted",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			_, err := conn.ChangeResourceRecordSets(zone, req)
			if err != nil {
				if strings.Contains(err.Error(), "PriorRequestNotComplete") {
					// There is some pending operation, so just retry
					// in a bit.
					return 42, "rejected", nil
				}

				if strings.Contains(err.Error(), "InvalidChangeBatch") {
					// This means that the record is already gone.
					return 42, "accepted", nil
				}

				return 42, "failure", err
			}

			return 42, "accepted", nil
		},
	}
	if _, err := wait.WaitForState(); err != nil {
		return err
	}

	return nil
}

func resource_aws_r53_record_refresh(
	s *terraform.InstanceState,
	meta interface{}) (*terraform.InstanceState, error) {
	p := meta.(*ResourceProvider)
	conn := p.route53

	zone := s.Attributes["zone_id"]
	lopts := &route53.ListOpts{
		Name: s.Attributes["name"],
		Type: s.Attributes["type"],
	}
	resp, err := conn.ListResourceRecordSets(zone, lopts)
	if err != nil {
		return s, err
	}

	// Scan for a matching record
	found := false
	for _, record := range resp.Records {
		if route53.FQDN(record.Name) != route53.FQDN(lopts.Name) {
			continue
		}
		if strings.ToUpper(record.Type) != strings.ToUpper(lopts.Type) {
			continue
		}

		found = true
		resource_aws_r53_record_update_state(s, &record)
		break
	}
	if !found {
		s.ID = ""
	}
	return s, nil
}

func resource_aws_r53_record_update_state(
	s *terraform.InstanceState,
	rec *route53.ResourceRecordSet) {

	flatRec := flatmap.Flatten(map[string]interface{}{
		"records": rec.Records,
	})
	for k, v := range flatRec {
		s.Attributes[k] = v
	}

	s.Attributes["ttl"] = strconv.FormatInt(int64(rec.TTL), 10)
}

func resource_aws_r53_record_diff(
	s *terraform.InstanceState,
	c *terraform.ResourceConfig,
	meta interface{}) (*terraform.InstanceDiff, error) {
	b := &diff.ResourceBuilder{
		Attrs: map[string]diff.AttrType{
			"zone_id": diff.AttrTypeCreate,
			"name":    diff.AttrTypeCreate,
			"type":    diff.AttrTypeCreate,
			"ttl":     diff.AttrTypeUpdate,
			"records": diff.AttrTypeUpdate,
		},
	}
	return b.Diff(s, c)
}
<|endoftext|>


<commit_before>package paperless

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"unicode"

	"github.com/kopoli/go-util"
)

// Environment is the run environment for each command. It is supplied as part
// of Status for a Link when it is Run or Validated.
type Environment struct {

	// Constants are variables that are defined before a command chain is run
	Constants map[string]string

	// Tempfiles are constants that house a name of a temporary file. The
	// files are created before the chain is run and they are removed at
	// the end.
	TempFiles []string

	// The directory where the commands are run. This is a created
	// temporary directory
	RootDir string

	// AllowedCommands contain the commands that are allowed. If this is
	// nil, all commands are allowed.
	AllowedCommands map[string]bool

	initialized bool
}

func (e *Environment) initEnv() (err error) {
	if e.initialized {
		return
	}

	if e.Constants == nil {
		return util.E.New("field Constants not initialized")
	}

	e.RootDir, err = ioutil.TempDir("", "chain")
	if err != nil {
		return util.E.Annotate(err, "rootdir creation failed")
	}

	var fp *os.File
	for _, name := range e.TempFiles {
		fp, err = ioutil.TempFile(e.RootDir, "tmp")
		if err != nil {
			err = util.E.Annotate(err, "tempfile creation failed")
			e.initialized = true
			e2 := e.deinitEnv()
			if e2 != nil {
				err = util.E.Annotate(err, "temproot removal failed:", e2)
				e.initialized = false
			}
			return
		}
		e.Constants[name] = fp.Name()
		fp.Close()
	}

	e.initialized = true
	return
}

func (e *Environment) deinitEnv() (err error) {
	if !e.initialized {
		return
	}

	if !strings.HasPrefix(e.RootDir, os.TempDir()) || e.RootDir == os.TempDir() {
		err = util.E.New("Temporary directory path is corrupted: %s", e.RootDir)
		return
	}

	err = os.RemoveAll(e.RootDir)
	if err != nil {
		err = util.E.Annotate(err, "tempdir removal failed:")
	}

	// Clear the temporary file and directory names
	e.RootDir = ""
	for _, n := range e.TempFiles {
		e.Constants[n] = ""
	}

	e.initialized = false
	return
}

func (e *Environment) validate() (err error) {
	if len(e.RootDir) == 0 {
		return util.E.New("the RootDir must be defined")
	}
	info, err := os.Stat(e.RootDir)
	if err != nil || info.Mode()&os.ModeDir == 0 {
		return util.E.Annotate(err, "file ", e.RootDir, " is not a proper directory")
	}

	return
}

// Status is the runtime status of the command chain
type Status struct {
	// The log output will be written to this
	Log io.Writer

	Environment
}

type Link interface {
	Validate(*Environment) error
	Run(*Status) error
}

type CmdChain struct {

	//TODO remove this (this should come from outside)
	Environment

	Links []Link
}

func (c *CmdChain) Validate(e *Environment) (err error) {
	for _, l := range c.Links {
		err = l.Validate(e)
		if err != nil {
			return
		}
	}
	return
}

func (c *CmdChain) Run(s *Status) (err error) {
	err = c.Validate(&s.Environment)
	if err != nil {
		return
	}

	for i := range c.Links {
		err = c.Links[i].Run(s)
		if err != nil {
			return
		}
	}

	return
}

func RunCmdChain(c *CmdChain, s *Status) (err error) {
	err = s.Environment.initEnv()
	if err != nil {
		return
	}

	err = c.Run(s)
	e2 := s.Environment.deinitEnv()
	if e2 != nil {
		err = util.E.Annotate(err, "cmdchain deinit failed: ", e2)
	}

	return
}

////////////////////////////////////////////////////////////

type Cmd struct {
	Cmd []string
}

// NewCmd creates a new Cmd from given command string
func NewCmd(cmdstr string) (c *Cmd, err error) {
	command := splitWsQuote(cmdstr)

	if len(command) == 0 {
		return nil, util.E.New("A command could not be parsed from:", cmdstr)
	}

	c = &Cmd{command}

	_, err = exec.LookPath(c.Cmd[0])
	if err != nil {
		return nil, util.E.Annotate(err, "Command", c.Cmd[0], "could not be found")

	}

	return
}

// Validate makes sure the command is proper and can be run
func (c *Cmd) Validate(e *Environment) (err error) {
	if len(c.Cmd) == 0 {
		return util.E.New("command string must be non-empty")
	}

	if e.AllowedCommands != nil {
		if _, ok := e.AllowedCommands[c.Cmd[0]]; ok != true {
			return util.E.New("command is not allowed")
		}
	}

	_, err = exec.LookPath(c.Cmd[0])
	if err != nil {
		return
	}

	err = e.validate()

	for idx, a := range c.Cmd {
		consts := parseConsts(a)
		if len(consts) > 0 {
			for _, co := range consts {
				if _, ok := e.Constants[co]; !ok {
					return util.E.New("constant \"%s\" not defined", co)
				}
			}
		}

		// Output redirection to a file
		if a == ">" && (idx == len(c.Cmd)-1 || c.Cmd[idx+1] == "") {
			return util.E.New("The output redirection requires a string")
		}
	}

	return
}

func (c *Cmd) Run(s *Status) (err error) {
	err = c.Validate(&s.Environment)
	if err != nil {
		return
	}

	var args []string
	for i := range c.Cmd {
		args = append(args, expandConsts(c.Cmd[i], s.Constants))
	}

	if s.Log != nil {
		fmt.Fprintln(s.Log, "# Running command:", strings.Join(args, " "))
	}

	var output io.Writer = s.Log

	redirout, pos := getRedirectFile(">", args)
	if redirout != "" {
		var fp *os.File
		redirout = PathAbs(s.RootDir, redirout)
		fp, err = os.OpenFile(redirout, os.O_WRONLY | os.O_CREATE, 0666)
		if err != nil {
			err = util.E.Annotate(err, "Could not open file",redirout,"for redirection")
			return
		}
		defer fp.Close()
		output = fp

		// Remove the redirection and the file argument from the command
		args = append(args[:pos], args[pos+2:]...)
	}

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Dir = s.RootDir
	cmd.Stdout = output
	cmd.Stderr = s.Log

	return cmd.Run()
}

////////////////////////////////////////////////////////////

var (
	constRe         = regexp.MustCompile(`\$(\w+)`)
	tmpfileConstRe  = regexp.MustCompile(`\$(tmp\w+)`)
	commentRe       = regexp.MustCompile(`#.*$`)
	preWhitespaceRe = regexp.MustCompile(`^\s+`)
)

// parseConsts parses the constants from a string. Returns a list of constant names
func parseConsts(s string) (ret []string) {
	ret = []string{}

	matches := constRe.FindAllStringSubmatch(s, -1)
	if matches == nil {
		return
	}
	for _, m := range matches {
		ret = append(ret, m[1])
	}

	return
}

func expandConsts(s string, constants map[string]string) string {
	return constRe.ReplaceAllStringFunc(s, func(match string) string {
		cs := parseConsts(match)
		if len(cs) != 1 {
			panic("Invalid Regexp parsing")
		}

		ret, ok := constants[cs[0]]
		if !ok {
			ret = ""
		}

		return ret
	})
}

// Gets the string after the given redir string. If not found, returns empty
// string.
func getRedirectFile(redir string, args []string) (file string, pos int) {
	for i := range args {
		if args[i] == redir {
			if i+1 < len(args) {
				file = args[i+1]
				pos = i
			}
			return
		}
	}
	return
}

// NewCmdChainScript creates a CmdChain from a script where each command is on a separate line. The following syntax elements are supported:
//
// - Empty lines are filtered out.
//
// - Comments start with # and end with EOL.
//
// - Constants are strings that begin with $ and they can be set before running the cmdchain.
//
// - Temporary files are strings that start with $tmp and they are automatically created before running the cmdchain and removed afterwards.
func NewCmdChainScript(script string) (c *CmdChain, err error) {
	c = &CmdChain{}
	c.Constants = make(map[string]string)

	for _, line := range strings.Split(script, "\n") {
		line = commentRe.ReplaceAllString(line, "")
		line = preWhitespaceRe.ReplaceAllString(line, "")

		if len(line) == 0 {
			continue
		}

		constants := parseConsts(line)
		for _, co := range constants {
			c.Constants[co] = ""
			if tmpfileConstRe.MatchString("$" + co) {
				c.TempFiles = append(c.TempFiles, co)
			}
		}

		var cmd *Cmd
		cmd, err = NewCmd(line)
		if err != nil {
			return nil, util.E.Annotate(err, "improper command")
		}

		c.Links = append(c.Links, cmd)
	}

	e := c.Environment
	e.RootDir = "/"

	err = c.Validate(&e)
	if err != nil {
		return nil, util.E.Annotate(err, "invalid command chain")
	}

	return
}

// splitWsQuote splits a string by whitespace, but takes doublequotes into
// account
func splitWsQuote(s string) []string {

	quote := rune(0)

	return strings.FieldsFunc(s, func(r rune) bool {
		switch {
		case r == quote:
			quote = rune(0)
			return true
		case quote != rune(0):
			return false
		case unicode.In(r, unicode.Quotation_Mark):
			quote = r
			return true
		default:
			return unicode.IsSpace(r)
		}
	})
}
<commit_msg>cmdchain: Fix format string<commit_after>package paperless

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"unicode"

	"github.com/kopoli/go-util"
)

// Environment is the run environment for each command. It is supplied as part
// of Status for a Link when it is Run or Validated.
type Environment struct {

	// Constants are variables that are defined before a command chain is run
	Constants map[string]string

	// Tempfiles are constants that house a name of a temporary file. The
	// files are created before the chain is run and they are removed at
	// the end.
	TempFiles []string

	// The directory where the commands are run. This is a created
	// temporary directory
	RootDir string

	// AllowedCommands contain the commands that are allowed. If this is
	// nil, all commands are allowed.
	AllowedCommands map[string]bool

	initialized bool
}

func (e *Environment) initEnv() (err error) {
	if e.initialized {
		return
	}

	if e.Constants == nil {
		return util.E.New("field Constants not initialized")
	}

	e.RootDir, err = ioutil.TempDir("", "chain")
	if err != nil {
		return util.E.Annotate(err, "rootdir creation failed")
	}

	var fp *os.File
	for _, name := range e.TempFiles {
		fp, err = ioutil.TempFile(e.RootDir, "tmp")
		if err != nil {
			err = util.E.Annotate(err, "tempfile creation failed")
			e.initialized = true
			e2 := e.deinitEnv()
			if e2 != nil {
				err = util.E.Annotate(err, "temproot removal failed:", e2)
				e.initialized = false
			}
			return
		}
		e.Constants[name] = fp.Name()
		fp.Close()
	}

	e.initialized = true
	return
}

func (e *Environment) deinitEnv() (err error) {
	if !e.initialized {
		return
	}

	if !strings.HasPrefix(e.RootDir, os.TempDir()) || e.RootDir == os.TempDir() {
		err = util.E.New("Temporary directory path is corrupted: %s", e.RootDir)
		return
	}

	err = os.RemoveAll(e.RootDir)
	if err != nil {
		err = util.E.Annotate(err, "tempdir removal failed:")
	}

	// Clear the temporary file and directory names
	e.RootDir = ""
	for _, n := range e.TempFiles {
		e.Constants[n] = ""
	}

	e.initialized = false
	return
}

func (e *Environment) validate() (err error) {
	if len(e.RootDir) == 0 {
		return util.E.New("the RootDir must be defined")
	}
	info, err := os.Stat(e.RootDir)
	if err != nil || info.Mode()&os.ModeDir == 0 {
		return util.E.Annotate(err, "file ", e.RootDir, " is not a proper directory")
	}

	return
}

// Status is the runtime status of the command chain
type Status struct {
	// The log output will be written to this
	Log io.Writer

	Environment
}

type Link interface {
	Validate(*Environment) error
	Run(*Status) error
}

type CmdChain struct {

	//TODO remove this (this should come from outside)
	Environment

	Links []Link
}

func (c *CmdChain) Validate(e *Environment) (err error) {
	for _, l := range c.Links {
		err = l.Validate(e)
		if err != nil {
			return
		}
	}
	return
}

func (c *CmdChain) Run(s *Status) (err error) {
	err = c.Validate(&s.Environment)
	if err != nil {
		return
	}

	for i := range c.Links {
		err = c.Links[i].Run(s)
		if err != nil {
			return
		}
	}

	return
}

func RunCmdChain(c *CmdChain, s *Status) (err error) {
	err = s.Environment.initEnv()
	if err != nil {
		return
	}

	err = c.Run(s)
	e2 := s.Environment.deinitEnv()
	if e2 != nil {
		err = util.E.Annotate(err, "cmdchain deinit failed: ", e2)
	}

	return
}

////////////////////////////////////////////////////////////

type Cmd struct {
	Cmd []string
}

// NewCmd creates a new Cmd from given command string
func NewCmd(cmdstr string) (c *Cmd, err error) {
	command := splitWsQuote(cmdstr)

	if len(command) == 0 {
		return nil, util.E.New("A command could not be parsed from: %s", cmdstr)
	}

	c = &Cmd{command}

	_, err = exec.LookPath(c.Cmd[0])
	if err != nil {
		return nil, util.E.Annotate(err, "Command", c.Cmd[0], "could not be found")

	}

	return
}

// Validate makes sure the command is proper and can be run
func (c *Cmd) Validate(e *Environment) (err error) {
	if len(c.Cmd) == 0 {
		return util.E.New("command string must be non-empty")
	}

	if e.AllowedCommands != nil {
		if _, ok := e.AllowedCommands[c.Cmd[0]]; ok != true {
			return util.E.New("command is not allowed")
		}
	}

	_, err = exec.LookPath(c.Cmd[0])
	if err != nil {
		return
	}

	err = e.validate()

	for idx, a := range c.Cmd {
		consts := parseConsts(a)
		if len(consts) > 0 {
			for _, co := range consts {
				if _, ok := e.Constants[co]; !ok {
					return util.E.New("constant \"%s\" not defined", co)
				}
			}
		}

		// Output redirection to a file
		if a == ">" && (idx == len(c.Cmd)-1 || c.Cmd[idx+1] == "") {
			return util.E.New("The output redirection requires a string")
		}
	}

	return
}

func (c *Cmd) Run(s *Status) (err error) {
	err = c.Validate(&s.Environment)
	if err != nil {
		return
	}

	var args []string
	for i := range c.Cmd {
		args = append(args, expandConsts(c.Cmd[i], s.Constants))
	}

	if s.Log != nil {
		fmt.Fprintln(s.Log, "# Running command:", strings.Join(args, " "))
	}

	var output io.Writer = s.Log

	redirout, pos := getRedirectFile(">", args)
	if redirout != "" {
		var fp *os.File
		redirout = PathAbs(s.RootDir, redirout)
		fp, err = os.OpenFile(redirout, os.O_WRONLY | os.O_CREATE, 0666)
		if err != nil {
			err = util.E.Annotate(err, "Could not open file",redirout,"for redirection")
			return
		}
		defer fp.Close()
		output = fp

		// Remove the redirection and the file argument from the command
		args = append(args[:pos], args[pos+2:]...)
	}

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Dir = s.RootDir
	cmd.Stdout = output
	cmd.Stderr = s.Log

	return cmd.Run()
}

////////////////////////////////////////////////////////////

var (
	constRe         = regexp.MustCompile(`\$(\w+)`)
	tmpfileConstRe  = regexp.MustCompile(`\$(tmp\w+)`)
	commentRe       = regexp.MustCompile(`#.*$`)
	preWhitespaceRe = regexp.MustCompile(`^\s+`)
)

// parseConsts parses the constants from a string. Returns a list of constant names
func parseConsts(s string) (ret []string) {
	ret = []string{}

	matches := constRe.FindAllStringSubmatch(s, -1)
	if matches == nil {
		return
	}
	for _, m := range matches {
		ret = append(ret, m[1])
	}

	return
}

func expandConsts(s string, constants map[string]string) string {
	return constRe.ReplaceAllStringFunc(s, func(match string) string {
		cs := parseConsts(match)
		if len(cs) != 1 {
			panic("Invalid Regexp parsing")
		}

		ret, ok := constants[cs[0]]
		if !ok {
			ret = ""
		}

		return ret
	})
}

// Gets the string after the given redir string. If not found, returns empty
// string.
func getRedirectFile(redir string, args []string) (file string, pos int) {
	for i := range args {
		if args[i] == redir {
			if i+1 < len(args) {
				file = args[i+1]
				pos = i
			}
			return
		}
	}
	return
}

// NewCmdChainScript creates a CmdChain from a script where each command is on a separate line. The following syntax elements are supported:
//
// - Empty lines are filtered out.
//
// - Comments start with # and end with EOL.
//
// - Constants are strings that begin with $ and they can be set before running the cmdchain.
//
// - Temporary files are strings that start with $tmp and they are automatically created before running the cmdchain and removed afterwards.
func NewCmdChainScript(script string) (c *CmdChain, err error) {
	c = &CmdChain{}
	c.Constants = make(map[string]string)

	for _, line := range strings.Split(script, "\n") {
		line = commentRe.ReplaceAllString(line, "")
		line = preWhitespaceRe.ReplaceAllString(line, "")

		if len(line) == 0 {
			continue
		}

		constants := parseConsts(line)
		for _, co := range constants {
			c.Constants[co] = ""
			if tmpfileConstRe.MatchString("$" + co) {
				c.TempFiles = append(c.TempFiles, co)
			}
		}

		var cmd *Cmd
		cmd, err = NewCmd(line)
		if err != nil {
			return nil, util.E.Annotate(err, "improper command")
		}

		c.Links = append(c.Links, cmd)
	}

	e := c.Environment
	e.RootDir = "/"

	err = c.Validate(&e)
	if err != nil {
		return nil, util.E.Annotate(err, "invalid command chain")
	}

	return
}

// splitWsQuote splits a string by whitespace, but takes doublequotes into
// account
func splitWsQuote(s string) []string {

	quote := rune(0)

	return strings.FieldsFunc(s, func(r rune) bool {
		switch {
		case r == quote:
			quote = rune(0)
			return true
		case quote != rune(0):
			return false
		case unicode.In(r, unicode.Quotation_Mark):
			quote = r
			return true
		default:
			return unicode.IsSpace(r)
		}
	})
}
<|endoftext|>


<commit_before>// Package lib contains: comnands.go - snapshot.go - uuid.go - written.go
//
// Commands provides all ZFS functions to manage the datasets and backups
//
package lib

import (
	"strings"

	"github.com/nfrance-conseil/zeplic/config"
	"github.com/IgnacioCarbajoVallejo/go-zfs"
)

var (
	w, _ = config.LogBook()
)

// TakeOrder takes a snapshot based on the order received from director
func TakeOrder(j int, DestDataset string) {
	// Check if dataset is configured
	index := -1
	for i := 0; i < j; i++ {
		pieces := config.Extract(i)
		dataset := pieces[2].(string)

		if dataset == DestDataset {
			index = i
			break
		} else {
			continue
		}
	}

	if index > -1 {
		// Extract data of dataset
		pieces := config.Extract(index)
		enable := pieces[0].(bool)
		clone := pieces[1].(string)
		dataset := pieces[2].(string)
		snapshot := pieces[3].(string)
		retain := pieces[4].(int)
		backup := pieces[5].(bool)
//		getClone := pieces[6].(bool)

		if dataset == DestDataset && enable == true {
			ds := Dataset(dataset)
			Snapshot(dataset, snapshot, ds)
			DeleteBackup(dataset)
			Policy(dataset, retain, clone)
			Backup(backup, dataset, ds)
		} else if dataset == DestDataset && enable == false {
			w.Notice("[NOTICE] the dataset '"+dataset+"' is disabled.")
		}
	} else {
		w.Notice("[NOTICE] the dataset '"+DestDataset+"' is not configured.")
	}
}

// DestroyOrder destroys a snapshot based on the order received from director
func DestroyOrder(j int, SnapshotName string) (bool, string) {
	// Define return variables
	var destroy bool
	var clone string

	// Get the snapshot
	snap, err := zfs.GetDataset(SnapshotName)
	if err != nil {
		w.Err("[ERROR] it was not possible to get the snapshot '"+SnapshotName+"'.")
	}

	// Get dataset of snapshot
	ds := DatasetName(SnapshotName)
	for i := 0; i < j; i++ {
		pieces := config.Extract(i)
		dataset := pieces[2].(string)

		if dataset == ds {
			clone = pieces[1].(string)
			break
		} else {
			continue
		}
	}

	err = snap.Destroy(zfs.DestroyDefault)
	if err != nil {
		destroy = false
	} else {
		destroy = true
	}
	return destroy, clone
}

// Runner is a loop that executes 'ZFS' functions for each dataset enabled
func Runner(j int) int {
	for i := 0; i < j; i++ {
		// Extract all data stored in JSON file
		pieces := config.Extract(i)

		// This value returns if the dataset is enable
		takedataset := pieces[0].(bool)

		// Get variables
		clone := pieces[1].(string)
		dataset := pieces[2].(string)
		snapshot := pieces[3].(string)
		retain := pieces[4].(int)
		backup := pieces[5].(bool)
		getClone := pieces[6].(bool)

		// Execute functions
		if takedataset == true {
			DeleteClone(clone)
			ds := Dataset(dataset)
			s, SnapshotName := Snapshot(dataset, snapshot, ds)
			DeleteBackup(dataset)
			Policy(dataset, retain, clone)
			Backup(backup, dataset, ds)
			Clone(getClone, clone, dataset, SnapshotName, s)
		} else if takedataset == false && dataset != "" {
			w.Notice("[NOTICE] the dataset '"+dataset+"' is disabled.")
		}
	}
	return 0
}

// DeleteClone deletes an existing clone
func DeleteClone(clone string) {
	// Get clones dataset
	cl, err := zfs.GetDataset(clone)
	if err != nil {
		w.Info("[INFO] the clone '"+clone+"' does not exist.")
	} else {
		// Destroy clones dataset
		err := cl.Destroy(zfs.DestroyRecursiveClones)
		if err != nil {
			w.Err("[ERROR] it was not possible to destroy the clone '"+clone+"'.")
		} else {
			w.Info("[INFO] the clone '"+clone+"' has been destroyed.")
		}
	}
}

// Dataset creates a dataset or get an existing one
func Dataset(dataset string) (*zfs.Dataset) {
	ds, err := zfs.GetDataset(dataset)

	// Destroy dataset (optional)
/*	err := ds.Destroy(zfs.DestroyRecursive)
	if err != nil {
		w.Err("[ERROR] it was not possible to destroy the dataset '"+dataset+"'.")
	} else {
		w.Info("[INFO] the dataset '"+dataset+"' has been destroyed.")
	}
	ds, err = zfs.GetDataset(dataset)*/

	if err != nil {
		w.Info("[INFO] the dataset '"+dataset+"' does not exist.")

		// Create dataset if it does not exist
		_, err := zfs.CreateFilesystem(dataset, nil)
		if err != nil {
			w.Err("[ERROR] it was not possible to create the dataset '"+dataset+"'.")
		} else {
			w.Info("[INFO] the dataset '"+dataset+"' has been created.")
		}
		ds, _ = zfs.GetDataset(dataset)
	}
	return ds
}

// Snapshot creates a new snapshot
func Snapshot(dataset, name string, ds *zfs.Dataset) (*zfs.Dataset, string) {
	SnapshotName := SnapName(name)
	s, err := ds.Snapshot(SnapshotName, false)

	// Get the snapshot created
	list, err := zfs.Snapshots(dataset)
	count := len(list)
	take := list[count-1].Name
	if strings.Contains(take, "BACKUP") {
		take = list[count-1].Name
	}
	// Check if it was created
	if err != nil {
		w.Err("[ERROR] it was not possible to create the snapshot '"+dataset+"@"+SnapshotName+"'.")
	} else {
		w.Info("[INFO] the snapshot '"+dataset+"@"+SnapshotName+"' has been created.")
		// Assign an uuid to the snapshot
		go UUID(take)
	}
	return s, SnapshotName
}

// DeleteBackup deletes the backup snapshot
func DeleteBackup(dataset string) {
	// Delete the backup snapshot
	list, err := zfs.Snapshots(dataset)
	if err != nil {
		w.Err("[ERROR] it was not possible to access of snapshots list.")
	}
	count := len(list)
	for k := 0; k < count; k++ {
		take := list[k].Name
		if strings.Contains(take, "BACKUP") {
			snap, err := zfs.GetDataset(take)
			if err != nil {
				w.Err("[ERROR] it was not possible to get the snapshot '"+take+"'.")
			}
			err = snap.Destroy(zfs.DestroyDefault)
			if err != nil {
				w.Err("[ERROR] it was not possible to destroy the snapshot '"+take+"'.")
			} else {
				w.Info("[INFO] the snapshot '"+take+"' has been destroyed.")
			}
		}
	}
}

// Policy apply the retention policy
func Policy(dataset string, retain int, clone string) {
	list, err := zfs.Snapshots(dataset)
	if err != nil {
		w.Err("[ERROR] it was not possible to access of snapshots list.")
	}
	count := len(list)
	for k := 0; count > retain; k++ {
		take := list[k].Name
		snap, err := zfs.GetDataset(take)
		if err != nil {
			w.Err("[ERROR] it was not possible to get the snapshot '"+take+"'.")
		}
		err = snap.Destroy(zfs.DestroyDefault)
		if err != nil {
			w.Err("[ERROR] the snapshot '"+take+"' has dependent clones: '"+clone+"'.")
		} else {
			w.Info("[INFO] the snapshot '"+take+"' has been destroyed.")
		}
		list, err = zfs.Snapshots(dataset)
		if err != nil {
			w.Err("[ERROR] it was not possible to access of snapshots list.")
		}
		count = len(list)
	}
}

// Backup creates a backup snapshot
func Backup(backup bool, dataset string, ds *zfs.Dataset) {
	if backup == true {
		_, err := ds.Snapshot(SnapBackup(dataset), false)

		// Get the snapshot created
		list, err := zfs.Snapshots(dataset)
		count := len(list)
		take := list[count-1].Name

		// Check if it was created
		if err != nil {
			w.Err("[ERROR] it was not possible to create the backup snapshot '"+dataset+"@"+SnapBackup(dataset)+"'.")
		} else {
			w.Info("[INFO] the backup snapshot '"+dataset+"@"+SnapBackup(dataset)+"' has been created.")
			// Assign an uuid to the snapshot
			go UUID(take)
		}
	}
}

// Clone creates a clone of last snapshot
func Clone(takeclone bool, clone string, dataset string, SnapshotName string, s *zfs.Dataset) {
	if takeclone == true {
		_, err := s.Clone(clone, nil)
		if err != nil {
			w.Err("[ERROR] it was not possible to clone the snapshot '"+dataset+"@"+SnapshotName+"'.")
		} else {
			w.Info("[INFO] the snapshot '"+dataset+"@"+SnapshotName+"' has been clone.")
		}
	}
}

// Rollback of last snapshot
func Rollback(rollback bool, dataset string, SnapshotName string, s *zfs.Dataset) {
	if rollback == true {
		err := s.Rollback(true)
		if err != nil {
			w.Err("[ERROR] it was not possible to rolling back the snapshot '"+SnapshotName+"'.")
		} else {
			w.Info("[INFO] the snapshot '"+dataset+"@"+SnapshotName+"' has been restored.")
		}
	}
}
<commit_msg>Update commands.go<commit_after>// Package lib contains: comnands.go - snapshot.go - uuid.go - written.go
//
// Commands provides all ZFS functions to manage the datasets and backups
//
package lib

import (
	"strings"

	"github.com/nfrance-conseil/zeplic/config"
	"github.com/IgnacioCarbajoVallejo/go-zfs"
)

var (
	w, _ = config.LogBook()
)

// TakeOrder takes a snapshot based on the order received from director
func TakeOrder(j int, DestDataset string) {
	// Check if dataset is configured
	index := -1
	for i := 0; i < j; i++ {
		pieces := config.Extract(i)
		dataset := pieces[2].(string)

		if dataset == DestDataset {
			index = i
			break
		} else {
			continue
		}
	}

	if index > -1 {
		// Extract data of dataset
		pieces := config.Extract(index)
		enable := pieces[0].(bool)
		clone := pieces[1].(string)
		dataset := pieces[2].(string)
		snapshot := pieces[3].(string)
		retain := pieces[4].(int)
		backup := pieces[5].(bool)
//		getClone := pieces[6].(bool)

		if dataset == DestDataset && enable == true {
			ds := Dataset(dataset)
			Snapshot(dataset, snapshot, ds)
			DeleteBackup(dataset)
			Policy(dataset, retain, clone)
			Backup(backup, dataset, ds)
		} else if dataset == DestDataset && enable == false {
			w.Notice("[NOTICE] the dataset '"+dataset+"' is disabled.")
		}
	} else {
		w.Notice("[NOTICE] the dataset '"+DestDataset+"' is not configured.")
	}
}

// DestroyOrder destroys a snapshot based on the order received from director
func DestroyOrder(j int, SnapshotName string) (bool, string) {
	// Define return variables
	var destroy bool
	var clone string

	// Get the snapshot
	snap, err := zfs.GetDataset(SnapshotName)
	if err != nil {
		w.Err("[ERROR] it was not possible to get the snapshot '"+SnapshotName+"'.")
	}

	// Get dataset of snapshot
	ds := DatasetName(SnapshotName)
	for i := 0; i < j; i++ {
		pieces := config.Extract(i)
		dataset := pieces[2].(string)

		if dataset == ds {
			clone = pieces[1].(string)
			break
		} else {
			continue
		}
	}

	err = snap.Destroy(zfs.DestroyDefault)
	if err != nil {
		destroy = false
	} else {
		destroy = true
	}
	return destroy, clone
}

// Runner is a loop that executes 'ZFS' functions for each dataset enabled
func Runner(j int) int {
	for i := 0; i < j; i++ {
		// Extract all data stored in JSON file
		pieces := config.Extract(i)

		// This value returns if the dataset is enable
		takedataset := pieces[0].(bool)

		// Get variables
		clone := pieces[1].(string)
		dataset := pieces[2].(string)
		snapshot := pieces[3].(string)
		retain := pieces[4].(int)
		backup := pieces[5].(bool)
		getClone := pieces[6].(bool)

		// Execute functions
		if takedataset == true {
			DeleteClone(clone)
			ds := Dataset(dataset)
			s, SnapshotName := Snapshot(dataset, snapshot, ds)
			DeleteBackup(dataset)
			Policy(dataset, retain, clone)
			Backup(backup, dataset, ds)
			Clone(getClone, clone, dataset, SnapshotName, s)
		} else if takedataset == false && dataset != "" {
			w.Notice("[NOTICE] the dataset '"+dataset+"' is disabled.")
		}
	}
	return 0
}

// DeleteClone deletes an existing clone
func DeleteClone(clone string) {
	// Get clones dataset
	cl, err := zfs.GetDataset(clone)
	if err != nil {
		w.Info("[INFO] the clone '"+clone+"' does not exist.")
	} else {
		// Destroy clones dataset
		err := cl.Destroy(zfs.DestroyRecursiveClones)
		if err != nil {
			w.Err("[ERROR] it was not possible to destroy the clone '"+clone+"'.")
		} else {
			w.Info("[INFO] the clone '"+clone+"' has been destroyed.")
		}
	}
}

// Dataset creates a dataset or get an existing one
func Dataset(dataset string) (*zfs.Dataset) {
	ds, err := zfs.GetDataset(dataset)

	// Destroy dataset (optional)
/*	err := ds.Destroy(zfs.DestroyRecursive)
	if err != nil {
		w.Err("[ERROR] it was not possible to destroy the dataset '"+dataset+"'.")
	} else {
		w.Info("[INFO] the dataset '"+dataset+"' has been destroyed.")
	}
	ds, err = zfs.GetDataset(dataset)*/

	if err != nil {
		w.Info("[INFO] the dataset '"+dataset+"' does not exist.")

		// Create dataset if it does not exist
		_, err := zfs.CreateFilesystem(dataset, nil)
		if err != nil {
			w.Err("[ERROR] it was not possible to create the dataset '"+dataset+"'.")
		} else {
			w.Info("[INFO] the dataset '"+dataset+"' has been created.")
		}
		ds, _ = zfs.GetDataset(dataset)
	}
	return ds
}

// Snapshot creates a new snapshot
func Snapshot(dataset, name string, ds *zfs.Dataset) (*zfs.Dataset, string) {
	SnapshotName := SnapName(name)
	s, err := ds.Snapshot(SnapshotName, false)

	// Get the snapshot created
	list, err := zfs.Snapshots(dataset)
	count := len(list)
	take := list[count-1].Name
	if strings.Contains(take, "BACKUP") {
		take = list[count-1].Name
	}
	// Check if it was created
	if err != nil {
		w.Err("[ERROR] it was not possible to create the snapshot '"+dataset+"@"+SnapshotName+"'.")
	} else {
		w.Info("[INFO] the snapshot '"+dataset+"@"+SnapshotName+"' has been created.")
		// Assign an uuid to the snapshot
		err := UUID(take)
		if err != nil {
			w.Err("[ERROR] it was not possible to assign an uuid to the snapshot '"+take+"'.")
		}
	}
	return s, SnapshotName
}

// DeleteBackup deletes the backup snapshot
func DeleteBackup(dataset string) {
	// Delete the backup snapshot
	list, err := zfs.Snapshots(dataset)
	if err != nil {
		w.Err("[ERROR] it was not possible to access of snapshots list.")
	}
	count := len(list)
	for k := 0; k < count; k++ {
		take := list[k].Name
		if strings.Contains(take, "BACKUP") {
			snap, err := zfs.GetDataset(take)
			if err != nil {
				w.Err("[ERROR] it was not possible to get the snapshot '"+take+"'.")
			}
			err = snap.Destroy(zfs.DestroyDefault)
			if err != nil {
				w.Err("[ERROR] it was not possible to destroy the snapshot '"+take+"'.")
			} else {
				w.Info("[INFO] the snapshot '"+take+"' has been destroyed.")
			}
		}
	}
}

// Policy apply the retention policy
func Policy(dataset string, retain int, clone string) {
	list, err := zfs.Snapshots(dataset)
	if err != nil {
		w.Err("[ERROR] it was not possible to access of snapshots list.")
	}
	count := len(list)
	for k := 0; count > retain; k++ {
		take := list[k].Name
		snap, err := zfs.GetDataset(take)
		if err != nil {
			w.Err("[ERROR] it was not possible to get the snapshot '"+take+"'.")
		}
		err = snap.Destroy(zfs.DestroyDefault)
		if err != nil {
			w.Err("[ERROR] the snapshot '"+take+"' has dependent clones: '"+clone+"'.")
		} else {
			w.Info("[INFO] the snapshot '"+take+"' has been destroyed.")
		}
		list, err = zfs.Snapshots(dataset)
		if err != nil {
			w.Err("[ERROR] it was not possible to access of snapshots list.")
		}
		count = len(list)
	}
}

// Backup creates a backup snapshot
func Backup(backup bool, dataset string, ds *zfs.Dataset) {
	if backup == true {
		_, err := ds.Snapshot(SnapBackup(dataset), false)

		// Get the snapshot created
		list, err := zfs.Snapshots(dataset)
		count := len(list)
		take := list[count-1].Name

		// Check if it was created
		if err != nil {
			w.Err("[ERROR] it was not possible to create the backup snapshot '"+dataset+"@"+SnapBackup(dataset)+"'.")
		} else {
			w.Info("[INFO] the backup snapshot '"+dataset+"@"+SnapBackup(dataset)+"' has been created.")
			// Assign an uuid to the snapshot
			err := UUID(take)
			if err != nil {
				w.Err("[ERROR] it was not possible to assign an uuid to the snapshot '"+take+"'.")
			}
		}
	}
}

// Clone creates a clone of last snapshot
func Clone(takeclone bool, clone string, dataset string, SnapshotName string, s *zfs.Dataset) {
	if takeclone == true {
		_, err := s.Clone(clone, nil)
		if err != nil {
			w.Err("[ERROR] it was not possible to clone the snapshot '"+dataset+"@"+SnapshotName+"'.")
		} else {
			w.Info("[INFO] the snapshot '"+dataset+"@"+SnapshotName+"' has been clone.")
		}
	}
}

// Rollback of last snapshot
func Rollback(rollback bool, dataset string, SnapshotName string, s *zfs.Dataset) {
	if rollback == true {
		err := s.Rollback(true)
		if err != nil {
			w.Err("[ERROR] it was not possible to rolling back the snapshot '"+SnapshotName+"'.")
		} else {
			w.Info("[INFO] the snapshot '"+dataset+"@"+SnapshotName+"' has been restored.")
		}
	}
}
<|endoftext|>


<commit_before>package lib

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"
	"github.com/op/go-logging"
	"strconv"
	"github.com/go-sql-driver/mysql"
)

// This is the one and only Database-object.
var db *sql.DB

// Opens a new connection-pool to the database using the given databaseConfiguration.
func OpenDatabase(config databaseConfiguration) {
	logging.MustGetLogger("logger").Info("Opening Database...")
	var err error = nil

	// username:password@protocol(address)/dbname
	db, err = sql.Open("mysql", config.User+":"+config.Password+"@tcp("+config.Host+":"+strconv.Itoa(config.Port)+")/"+config.Database)
	if err != nil {
		logging.MustGetLogger("logger").Fatal("Unable to open Database-Connection: ", err)
	}
	db.SetMaxOpenConns(config.ConnectionLimit)

	err = db.Ping()
	if err != nil {
		logging.MustGetLogger("logger").Fatal("Unable to reach Database: ", err)
	}

	prepareDatabase()
}

// Creates the needed tables in the database.
func prepareDatabase() {
	logging.MustGetLogger("logger").Debug("Preparing Database...")

	_, err := db.Exec("CREATE TABLE IF NOT EXISTS `website` (`id` int(11) NOT NULL AUTO_INCREMENT, `name` varchar(50) NOT NULL, `enabled` int(1) NOT NULL DEFAULT '1', `visible` int(1) NOT NULL DEFAULT '1', `protocol` varchar(8) NOT NULL DEFAULT 'https', `url` varchar(100) NOT NULL, `checkMethod` VARCHAR(10) NOT NULL DEFAULT 'HEAD', `status` varchar(50) NOT NULL DEFAULT 'unknown', `time` datetime NOT NULL DEFAULT '0000-00-00 00:00:00', `lastFailStatus` varchar(50) NOT NULL DEFAULT 'unknown', `lastFailTime` datetime NOT NULL DEFAULT '0000-00-00 00:00:00', `ups` int(11) NOT NULL DEFAULT '0', `downs` int(11) NOT NULL DEFAULT '0', `totalChecks` int(11) NOT NULL DEFAULT '0', `avgAvail` float NOT NULL DEFAULT '100', PRIMARY KEY (`id`), UNIQUE KEY `url` (`url`)) ENGINE=InnoDB AUTO_INCREMENT=1 DEFAULT CHARSET=utf8;")
	if err != nil {
		logging.MustGetLogger("logger").Fatal("Unable to create table 'website': ", err)
	}

	// v2.0.0 Beta
	_, err = db.Exec("ALTER TABLE `website` ADD `checkMethod` VARCHAR(10) NOT NULL DEFAULT 'HEAD' AFTER `url`;")
	if mysqlError, ok := err.(*mysql.MySQLError); !ok {
		if mysqlError.Number != 1060 {
			logging.MustGetLogger("logger").Fatal("Unable to add 'checkMethod'-column: ", err)
		}
	}

	_, err = db.Exec("CREATE TABLE IF NOT EXISTS `settings` (`id` int(11) NOT NULL AUTO_INCREMENT, `name` varchar(20) NOT NULL, `value` varchar(1024) NOT NULL, PRIMARY KEY (`id`), UNIQUE KEY `name` (`name`)) ENGINE=InnoDB AUTO_INCREMENT=1 DEFAULT CHARSET=utf8;")
	if err != nil {
		logging.MustGetLogger("logger").Fatal("Unable to create table 'settings': ", err)
	}
}

// Returns the current Database-object.
func GetDatabase() *sql.DB {
	if db == nil {
		logging.MustGetLogger("logger").Fatal("No active Database found.")
	} else {
		return db
	}
	return nil
}
<commit_msg>gofmt -s -w .<commit_after>package lib

import (
	"database/sql"
	"github.com/go-sql-driver/mysql"
	_ "github.com/go-sql-driver/mysql"
	"github.com/op/go-logging"
	"strconv"
)

// This is the one and only Database-object.
var db *sql.DB

// Opens a new connection-pool to the database using the given databaseConfiguration.
func OpenDatabase(config databaseConfiguration) {
	logging.MustGetLogger("logger").Info("Opening Database...")
	var err error = nil

	// username:password@protocol(address)/dbname
	db, err = sql.Open("mysql", config.User+":"+config.Password+"@tcp("+config.Host+":"+strconv.Itoa(config.Port)+")/"+config.Database)
	if err != nil {
		logging.MustGetLogger("logger").Fatal("Unable to open Database-Connection: ", err)
	}
	db.SetMaxOpenConns(config.ConnectionLimit)

	err = db.Ping()
	if err != nil {
		logging.MustGetLogger("logger").Fatal("Unable to reach Database: ", err)
	}

	prepareDatabase()
}

// Creates the needed tables in the database.
func prepareDatabase() {
	logging.MustGetLogger("logger").Debug("Preparing Database...")

	_, err := db.Exec("CREATE TABLE IF NOT EXISTS `website` (`id` int(11) NOT NULL AUTO_INCREMENT, `name` varchar(50) NOT NULL, `enabled` int(1) NOT NULL DEFAULT '1', `visible` int(1) NOT NULL DEFAULT '1', `protocol` varchar(8) NOT NULL DEFAULT 'https', `url` varchar(100) NOT NULL, `checkMethod` VARCHAR(10) NOT NULL DEFAULT 'HEAD', `status` varchar(50) NOT NULL DEFAULT 'unknown', `time` datetime NOT NULL DEFAULT '0000-00-00 00:00:00', `lastFailStatus` varchar(50) NOT NULL DEFAULT 'unknown', `lastFailTime` datetime NOT NULL DEFAULT '0000-00-00 00:00:00', `ups` int(11) NOT NULL DEFAULT '0', `downs` int(11) NOT NULL DEFAULT '0', `totalChecks` int(11) NOT NULL DEFAULT '0', `avgAvail` float NOT NULL DEFAULT '100', PRIMARY KEY (`id`), UNIQUE KEY `url` (`url`)) ENGINE=InnoDB AUTO_INCREMENT=1 DEFAULT CHARSET=utf8;")
	if err != nil {
		logging.MustGetLogger("logger").Fatal("Unable to create table 'website': ", err)
	}

	// v2.0.0 Beta
	_, err = db.Exec("ALTER TABLE `website` ADD `checkMethod` VARCHAR(10) NOT NULL DEFAULT 'HEAD' AFTER `url`;")
	if mysqlError, ok := err.(*mysql.MySQLError); !ok {
		if mysqlError.Number != 1060 {
			logging.MustGetLogger("logger").Fatal("Unable to add 'checkMethod'-column: ", err)
		}
	}

	_, err = db.Exec("CREATE TABLE IF NOT EXISTS `settings` (`id` int(11) NOT NULL AUTO_INCREMENT, `name` varchar(20) NOT NULL, `value` varchar(1024) NOT NULL, PRIMARY KEY (`id`), UNIQUE KEY `name` (`name`)) ENGINE=InnoDB AUTO_INCREMENT=1 DEFAULT CHARSET=utf8;")
	if err != nil {
		logging.MustGetLogger("logger").Fatal("Unable to create table 'settings': ", err)
	}
}

// Returns the current Database-object.
func GetDatabase() *sql.DB {
	if db == nil {
		logging.MustGetLogger("logger").Fatal("No active Database found.")
	} else {
		return db
	}
	return nil
}
<|endoftext|>


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

package run

import (
	"bytes"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/beego/bee/config"
	beeLogger "github.com/beego/bee/logger"
	"github.com/beego/bee/logger/colors"
	"github.com/beego/bee/utils"
	"github.com/fsnotify/fsnotify"
)

var (
	cmd                 *exec.Cmd
	state               sync.Mutex
	eventTime           = make(map[string]int64)
	scheduleTime        time.Time
	watchExts           = []string{".go"}
	watchExtsStatic     = []string{".html", ".tpl", ".js", ".css"}
	ignoredFilesRegExps = []string{
		`.#(\w+).go`,
		`.(\w+).go.swp`,
		`(\w+).go~`,
		`(\w+).tmp`,
	}
)

// NewWatcher starts an fsnotify Watcher on the specified paths
func NewWatcher(paths []string, files []string, isgenerate bool) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		beeLogger.Log.Fatalf("Failed to create watcher: %s", err)
	}

	go func() {
		for {
			select {
			case e := <-watcher.Events:
				isBuild := true

				if ifStaticFile(e.Name) && config.Conf.EnableReload {
					sendReload(e.String())
					continue
				}
				// Skip ignored files
				if shouldIgnoreFile(e.Name) {
					continue
				}
				if !shouldWatchFileWithExtension(e.Name) {
					continue
				}

				mt := utils.GetFileModTime(e.Name)
				if t := eventTime[e.Name]; mt == t {
					beeLogger.Log.Hintf(colors.Bold("Skipping: ")+"%s", e.String())
					isBuild = false
				}

				eventTime[e.Name] = mt

				if isBuild {
					beeLogger.Log.Hintf("Event fired: %s", e)
					go func() {
						// Wait 1s before autobuild until there is no file change.
						scheduleTime = time.Now().Add(1 * time.Second)
						time.Sleep(scheduleTime.Sub(time.Now()))
						AutoBuild(files, isgenerate)
					}()
				}
			case err := <-watcher.Errors:
				beeLogger.Log.Warnf("Watcher error: %s", err.Error()) // No need to exit here
			}
		}
	}()

	beeLogger.Log.Info("Initializing watcher...")
	for _, path := range paths {
		beeLogger.Log.Hintf(colors.Bold("Watching: ")+"%s", path)
		err = watcher.Add(path)
		if err != nil {
			beeLogger.Log.Fatalf("Failed to watch directory: %s", err)
		}
	}
}

// AutoBuild builds the specified set of files
func AutoBuild(files []string, isgenerate bool) {
	state.Lock()
	defer state.Unlock()

	os.Chdir(currpath)

	cmdName := "go"

	var (
		err    error
		stderr bytes.Buffer
	)
	// For applications use full import path like "github.com/.../.."
	// are able to use "go install" to reduce build time.
	if config.Conf.GoInstall {
		icmd := exec.Command(cmdName, "install", "-v")
		icmd.Stdout = os.Stdout
		icmd.Stderr = os.Stderr
		icmd.Env = append(os.Environ(), "GOGC=off")
		icmd.Run()
	}

	if isgenerate {
		beeLogger.Log.Info("Generating the docs...")
		icmd := exec.Command("bee", "generate", "docs")
		icmd.Env = append(os.Environ(), "GOGC=off")
		err = icmd.Run()
		if err != nil {
			utils.Notify("", "Failed to generate the docs.")
			beeLogger.Log.Errorf("Failed to generate the docs.")
			return
		}
		beeLogger.Log.Success("Docs generated!")
	}

	if err == nil {
		appName := appname
		if runtime.GOOS == "windows" {
			appName += ".exe"
		}

		args := []string{"build"}
		args = append(args, "-o", appName)
		if buildTags != "" {
			args = append(args, "-tags", buildTags)
		}
		args = append(args, files...)

		bcmd := exec.Command(cmdName, args...)
		bcmd.Env = append(os.Environ(), "GOGC=off")
		bcmd.Stderr = &stderr
		err = bcmd.Run()
		if err != nil {
			utils.Notify(stderr.String(), "Build Failed")
			beeLogger.Log.Errorf("Failed to build the application: %s", stderr.String())
			return
		}
	}

	beeLogger.Log.Success("Built Successfully!")
	Restart(appname)
}

// Kill kills the running command process
func Kill() {
	defer func() {
		if e := recover(); e != nil {
			beeLogger.Log.Infof("Kill recover: %s", e)
		}
	}()
	if cmd != nil && cmd.Process != nil {
		err := cmd.Process.Kill()
		if err != nil {
			beeLogger.Log.Errorf("Error while killing cmd process: %s", err)
		}
	}
}

// Restart kills the running command process and starts it again
func Restart(appname string) {
	beeLogger.Log.Debugf("Kill running process", utils.FILE(), utils.LINE())
	Kill()
	go Start(appname)
}

// Start starts the command process
func Start(appname string) {
	beeLogger.Log.Infof("Restarting '%s'...", appname)
	if !strings.Contains(appname, "./") {
		appname = "./" + appname
	}

	cmd = exec.Command(appname)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Args = append([]string{appname}, config.Conf.CmdArgs...)
	cmd.Env = append(os.Environ(), config.Conf.Envs...)

	go cmd.Run()
	beeLogger.Log.Successf("'%s' is running...", appname)
	started <- true
}

func ifStaticFile(filename string) bool {
	for _, s := range watchExtsStatic {
		if strings.HasSuffix(filename, s) {
			return true
		}
	}
	return false
}

// shouldIgnoreFile ignores filenames generated by Emacs, Vim or SublimeText.
// It returns true if the file should be ignored, false otherwise.
func shouldIgnoreFile(filename string) bool {
	for _, regex := range ignoredFilesRegExps {
		r, err := regexp.Compile(regex)
		if err != nil {
			beeLogger.Log.Fatalf("Could not compile regular expression: %s", err)
		}
		if r.MatchString(filename) {
			return true
		}
		continue
	}
	return false
}

// shouldWatchFileWithExtension returns true if the name of the file
// hash a suffix that should be watched.
func shouldWatchFileWithExtension(name string) bool {
	for _, s := range watchExts {
		if strings.HasSuffix(name, s) {
			return true
		}
	}
	return false
}
<commit_msg>after build the browser is refreshed<commit_after>// Copyright 2013 bee authors
//
// Licensed under the Apache License, Version 2.0 (the "License"): you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

package run

import (
	"bytes"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/beego/bee/config"
	beeLogger "github.com/beego/bee/logger"
	"github.com/beego/bee/logger/colors"
	"github.com/beego/bee/utils"
	"github.com/fsnotify/fsnotify"
)

var (
	cmd                 *exec.Cmd
	state               sync.Mutex
	eventTime           = make(map[string]int64)
	scheduleTime        time.Time
	watchExts           = []string{".go"}
	watchExtsStatic     = []string{".html", ".tpl", ".js", ".css"}
	ignoredFilesRegExps = []string{
		`.#(\w+).go`,
		`.(\w+).go.swp`,
		`(\w+).go~`,
		`(\w+).tmp`,
	}
)

// NewWatcher starts an fsnotify Watcher on the specified paths
func NewWatcher(paths []string, files []string, isgenerate bool) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		beeLogger.Log.Fatalf("Failed to create watcher: %s", err)
	}

	go func() {
		for {
			select {
			case e := <-watcher.Events:
				isBuild := true

				if ifStaticFile(e.Name) && config.Conf.EnableReload {
					sendReload(e.String())
					continue
				}
				// Skip ignored files
				if shouldIgnoreFile(e.Name) {
					continue
				}
				if !shouldWatchFileWithExtension(e.Name) {
					continue
				}

				mt := utils.GetFileModTime(e.Name)
				if t := eventTime[e.Name]; mt == t {
					beeLogger.Log.Hintf(colors.Bold("Skipping: ")+"%s", e.String())
					isBuild = false
				}

				eventTime[e.Name] = mt

				if isBuild {
					beeLogger.Log.Hintf("Event fired: %s", e)
					go func() {
						// Wait 1s before autobuild until there is no file change.
						scheduleTime = time.Now().Add(1 * time.Second)
						time.Sleep(scheduleTime.Sub(time.Now()))
						AutoBuild(files, isgenerate)

						if config.Conf.EnableReload {
							// Wait 100ms more before refreshing the browser
							time.Sleep(100 * time.Millisecond)
							sendReload(e.String())
						}
					}()
				}
			case err := <-watcher.Errors:
				beeLogger.Log.Warnf("Watcher error: %s", err.Error()) // No need to exit here
			}
		}
	}()

	beeLogger.Log.Info("Initializing watcher...")
	for _, path := range paths {
		beeLogger.Log.Hintf(colors.Bold("Watching: ")+"%s", path)
		err = watcher.Add(path)
		if err != nil {
			beeLogger.Log.Fatalf("Failed to watch directory: %s", err)
		}
	}
}

// AutoBuild builds the specified set of files
func AutoBuild(files []string, isgenerate bool) {
	state.Lock()
	defer state.Unlock()

	os.Chdir(currpath)

	cmdName := "go"

	var (
		err    error
		stderr bytes.Buffer
	)
	// For applications use full import path like "github.com/.../.."
	// are able to use "go install" to reduce build time.
	if config.Conf.GoInstall {
		icmd := exec.Command(cmdName, "install", "-v")
		icmd.Stdout = os.Stdout
		icmd.Stderr = os.Stderr
		icmd.Env = append(os.Environ(), "GOGC=off")
		icmd.Run()
	}

	if isgenerate {
		beeLogger.Log.Info("Generating the docs...")
		icmd := exec.Command("bee", "generate", "docs")
		icmd.Env = append(os.Environ(), "GOGC=off")
		err = icmd.Run()
		if err != nil {
			utils.Notify("", "Failed to generate the docs.")
			beeLogger.Log.Errorf("Failed to generate the docs.")
			return
		}
		beeLogger.Log.Success("Docs generated!")
	}

	if err == nil {
		appName := appname
		if runtime.GOOS == "windows" {
			appName += ".exe"
		}

		args := []string{"build"}
		args = append(args, "-o", appName)
		if buildTags != "" {
			args = append(args, "-tags", buildTags)
		}
		args = append(args, files...)

		bcmd := exec.Command(cmdName, args...)
		bcmd.Env = append(os.Environ(), "GOGC=off")
		bcmd.Stderr = &stderr
		err = bcmd.Run()
		if err != nil {
			utils.Notify(stderr.String(), "Build Failed")
			beeLogger.Log.Errorf("Failed to build the application: %s", stderr.String())
			return
		}
	}

	beeLogger.Log.Success("Built Successfully!")
	Restart(appname)
}

// Kill kills the running command process
func Kill() {
	defer func() {
		if e := recover(); e != nil {
			beeLogger.Log.Infof("Kill recover: %s", e)
		}
	}()
	if cmd != nil && cmd.Process != nil {
		err := cmd.Process.Kill()
		if err != nil {
			beeLogger.Log.Errorf("Error while killing cmd process: %s", err)
		}
	}
}

// Restart kills the running command process and starts it again
func Restart(appname string) {
	beeLogger.Log.Debugf("Kill running process", utils.FILE(), utils.LINE())
	Kill()
	go Start(appname)
}

// Start starts the command process
func Start(appname string) {
	beeLogger.Log.Infof("Restarting '%s'...", appname)
	if !strings.Contains(appname, "./") {
		appname = "./" + appname
	}

	cmd = exec.Command(appname)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Args = append([]string{appname}, config.Conf.CmdArgs...)
	cmd.Env = append(os.Environ(), config.Conf.Envs...)

	go cmd.Run()
	beeLogger.Log.Successf("'%s' is running...", appname)
	started <- true
}

func ifStaticFile(filename string) bool {
	for _, s := range watchExtsStatic {
		if strings.HasSuffix(filename, s) {
			return true
		}
	}
	return false
}

// shouldIgnoreFile ignores filenames generated by Emacs, Vim or SublimeText.
// It returns true if the file should be ignored, false otherwise.
func shouldIgnoreFile(filename string) bool {
	for _, regex := range ignoredFilesRegExps {
		r, err := regexp.Compile(regex)
		if err != nil {
			beeLogger.Log.Fatalf("Could not compile regular expression: %s", err)
		}
		if r.MatchString(filename) {
			return true
		}
		continue
	}
	return false
}

// shouldWatchFileWithExtension returns true if the name of the file
// hash a suffix that should be watched.
func shouldWatchFileWithExtension(name string) bool {
	for _, s := range watchExts {
		if strings.HasSuffix(name, s) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package cmd

import (
	"github.com/drud/ddev/pkg/ddevapp"
	"github.com/drud/ddev/pkg/dockerutil"
	"github.com/drud/ddev/pkg/util"
	"github.com/spf13/cobra"
)

var outFileName string
var gzipOption bool
var exportTargetDB string

// ExportDBCmd is the `ddev export-db` command.
var ExportDBCmd = &cobra.Command{
	Use:   "export-db [project]",
	Short: "Dump a database to a file or to stdout",
	Long:  `Dump a database to a file or to stdout`,
	Example: `ddev export-db --file=/tmp/db.sql.gz'
ddev export-db -f /tmp/db.sql.gz
ddev export-db --gzip=false --file /tmp/db.sql
ddev export-db > /tmp/db.sql.gz
ddev export-db --gzip=false > /tmp/db.sql
ddev export-db myproject --gzip=false --file=/tmp/myproject.sql
ddev export-db someproject --gzip=false --file=/tmp/someproject.sql `,
	Args: cobra.RangeArgs(0, 1),
	PreRun: func(cmd *cobra.Command, args []string) {
		dockerutil.EnsureDdevNetwork()
	},
	Run: func(cmd *cobra.Command, args []string) {
		projects, err := getRequestedProjects(args, false)
		if err != nil {
			util.Failed("Unable to get project(s): %v", err)
		}

		app := projects[0]

		if app.SiteStatus() != ddevapp.SiteRunning {
			util.Failed("ddev can't export-db until the project is started, please use ddev start.")
		}

		err = app.ExportDB(outFileName, gzipOption, exportTargetDB)
		if err != nil {
			util.Failed("Failed to export database for %s: %v", app.GetName(), err)
		}
	},
}

func init() {
	ExportDBCmd.Flags().StringVarP(&outFileName, "file", "f", "", "Provide the path to output the dump")
	ExportDBCmd.Flags().BoolVarP(&gzipOption, "gzip", "z", true, "If provided asset is an archive, provide the path to extract within the archive.")
	ExportDBCmd.Flags().StringVarP(&exportTargetDB, "target-db", "d", "db", "If provided, target-db is alternate database to export")
	RootCmd.AddCommand(ExportDBCmd)
}
<commit_msg>Remove stray apostrophe in export-db help, fixes #3710<commit_after>package cmd

import (
	"github.com/drud/ddev/pkg/ddevapp"
	"github.com/drud/ddev/pkg/dockerutil"
	"github.com/drud/ddev/pkg/util"
	"github.com/spf13/cobra"
)

var outFileName string
var gzipOption bool
var exportTargetDB string

// ExportDBCmd is the `ddev export-db` command.
var ExportDBCmd = &cobra.Command{
	Use:   "export-db [project]",
	Short: "Dump a database to a file or to stdout",
	Long:  `Dump a database to a file or to stdout`,
	Example: `ddev export-db --file=/tmp/db.sql.gz
ddev export-db -f /tmp/db.sql.gz
ddev export-db --gzip=false --file /tmp/db.sql
ddev export-db > /tmp/db.sql.gz
ddev export-db --gzip=false > /tmp/db.sql
ddev export-db myproject --gzip=false --file=/tmp/myproject.sql
ddev export-db someproject --gzip=false --file=/tmp/someproject.sql `,
	Args: cobra.RangeArgs(0, 1),
	PreRun: func(cmd *cobra.Command, args []string) {
		dockerutil.EnsureDdevNetwork()
	},
	Run: func(cmd *cobra.Command, args []string) {
		projects, err := getRequestedProjects(args, false)
		if err != nil {
			util.Failed("Unable to get project(s): %v", err)
		}

		app := projects[0]

		if app.SiteStatus() != ddevapp.SiteRunning {
			util.Failed("ddev can't export-db until the project is started, please use ddev start.")
		}

		err = app.ExportDB(outFileName, gzipOption, exportTargetDB)
		if err != nil {
			util.Failed("Failed to export database for %s: %v", app.GetName(), err)
		}
	},
}

func init() {
	ExportDBCmd.Flags().StringVarP(&outFileName, "file", "f", "", "Provide the path to output the dump")
	ExportDBCmd.Flags().BoolVarP(&gzipOption, "gzip", "z", true, "If provided asset is an archive, provide the path to extract within the archive.")
	ExportDBCmd.Flags().StringVarP(&exportTargetDB, "target-db", "d", "db", "If provided, target-db is alternate database to export")
	RootCmd.AddCommand(ExportDBCmd)
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/cmd/kateway/structs"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/pmylund/sortutil"
	"github.com/ryanuber/columnize"
)

type hostLoadInfo struct {
	host            string
	cluster         string
	topicPartitions []structs.TopicPartition
	qps             int64
}

type hostOffsetInfo struct {
	host      string
	offsetMap map[string]map[structs.TopicPartition]int64 // cluster:tp:offset
}

func (ho hostOffsetInfo) Clusters() []string {
	var r []string
	for cluster, _ := range ho.offsetMap {
		r = append(r, cluster)
	}
	return r
}

func (ho hostOffsetInfo) Total() (t int64) {
	for _, tps := range ho.offsetMap {
		for _, qps := range tps {
			t += qps
		}
	}

	return
}

func (ho hostOffsetInfo) ClusterTotal(cluster string) (t int64) {
	for _, qps := range ho.offsetMap[cluster] {
		t += qps
	}

	return
}

type Balance struct {
	Ui  cli.Ui
	Cmd string

	zone, cluster string
	interval      time.Duration
	summaryMode   bool
	host          string

	offsets     map[string]int64 // host => offset sum TODO
	lastOffsets map[string]int64

	lastHostOffsets map[string]hostOffsetInfo

	hostOffsetCh chan map[string]hostOffsetInfo // key is host
	signalsCh    map[string]chan struct{}
}

func (this *Balance) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("balance", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&this.cluster, "c", "", "")
	cmdFlags.DurationVar(&this.interval, "i", time.Second*5, "")
	cmdFlags.StringVar(&this.host, "host", "", "")
	cmdFlags.BoolVar(&this.summaryMode, "sum", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	this.signalsCh = make(map[string]chan struct{})
	this.hostOffsetCh = make(chan map[string]hostOffsetInfo)

	this.lastHostOffsets = make(map[string]hostOffsetInfo)
	this.offsets = make(map[string]int64)
	this.lastOffsets = make(map[string]int64)

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), this.cluster) {
			return
		}

		this.signalsCh[zkcluster.Name()] = make(chan struct{})

		go this.clusterTopProducers(zkcluster)
	})

	this.drawBalance()

	return
}

func (this *Balance) startAll() {
	for _, ch := range this.signalsCh {
		ch <- struct{}{}
	}
}

func (this *Balance) collectAll(seq int) {
	for _, _ = range this.signalsCh {
		offsets := <-this.hostOffsetCh
		if seq == 0 {
			// record into lastHostOffsets
			for host, offsetInfo := range offsets {
				if _, present := this.lastHostOffsets[host]; !present {
					this.lastHostOffsets[host] = hostOffsetInfo{host: host, offsetMap: make(map[string]map[structs.TopicPartition]int64)}
				}

				for cluster, tps := range offsetInfo.offsetMap {
					if _, present := this.lastHostOffsets[host].offsetMap[cluster]; !present {
						this.lastHostOffsets[host].offsetMap[cluster] = make(map[structs.TopicPartition]int64)
					}

					for tp, off := range tps {
						// 1st loop, offset
						this.lastHostOffsets[host].offsetMap[cluster][tp] = off
					}
				}
			}
		} else {
			for host, offsetInfo := range offsets {
				for cluster, tps := range offsetInfo.offsetMap {
					for tp, off := range tps {
						// 2nd loop, qps
						// FIXME hard coding
						this.lastHostOffsets[host].offsetMap[cluster][tp] = (off - this.lastHostOffsets[host].offsetMap[cluster][tp]) / int64(this.interval.Seconds())
					}
				}
			}
		}

	}

}

func (this *Balance) drawBalance() {
	for i := 0; i < 2; i++ {
		this.startAll()
		time.Sleep(this.interval)
		this.collectAll(i)
	}

	type hostTps struct {
		host string
		tps  int64
	}
	var sortedHosts []hostTps
	for host, info := range this.lastHostOffsets {
		sortedHosts = append(sortedHosts, hostTps{host, info.Total()})
	}
	sortutil.AscByField(sortedHosts, "tps")

	var hosts []string
	for _, h := range sortedHosts {
		hosts = append(hosts, h.host)
	}

	if this.summaryMode {
		this.drawSummary(hosts)
		return
	}

	this.drawDetail(hosts)
}

func (this *Balance) drawDetail(sortedHosts []string) {
	type hostSummary struct {
		cluster string
		tp      structs.TopicPartition
		qps     int64
	}

	for _, host := range sortedHosts {
		var summary []hostSummary
		offsetInfo := this.lastHostOffsets[host]

		for cluster, tps := range offsetInfo.offsetMap {
			for tp, qps := range tps {
				summary = append(summary, hostSummary{cluster, tp, qps})
			}
		}

		sortutil.DescByField(summary, "qps")

		this.Ui.Output(color.Green("%16s %8s %+v", host, gofmt.Comma(offsetInfo.Total()), offsetInfo.Clusters()))
		for _, sum := range summary {
			if sum.qps < 100 {
				continue
			}
			this.Ui.Output(fmt.Sprintf("    %30s %8s %s", sum.cluster, gofmt.Comma(sum.qps), sum.tp))
		}
	}

}

func (this *Balance) drawSummary(sortedHosts []string) {
	lines := []string{"#|Broker|Total|Cluster|OPS"}
	var totalTps int64
	var stripes = []string{"◉", "◎"}
	for i, host := range sortedHosts {
		offsetInfo := this.lastHostOffsets[host]
		for cluster, _ := range offsetInfo.offsetMap {
			hostTps := offsetInfo.Total()
			clusterTps := offsetInfo.ClusterTotal(cluster)
			totalTps += hostTps

			lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s",
				stripes[i%2], host,
				gofmt.Comma(hostTps), cluster, gofmt.Comma(clusterTps)))
		}
	}
	this.Ui.Output(columnize.SimpleFormat(lines))
	this.Ui.Output(fmt.Sprintf("-Total- Hosts:%d Tps:%s", len(sortedHosts), gofmt.Comma(totalTps)))
}

func (this *Balance) clusterTopProducers(zkcluster *zk.ZkCluster) {
	kfk, err := sarama.NewClient(zkcluster.BrokerList(), sarama.NewConfig())
	if err != nil {
		return
	}
	defer kfk.Close()

	for {
		hostOffsets := make(map[string]hostOffsetInfo)

		topics, err := kfk.Topics()
		swallow(err)

		<-this.signalsCh[zkcluster.Name()]

		for _, topic := range topics {
			partions, err := kfk.WritablePartitions(topic)
			swallow(err)
			for _, partitionID := range partions {
				leader, err := kfk.Leader(topic, partitionID)
				swallow(err)

				latestOffset, err := kfk.GetOffset(topic, partitionID, sarama.OffsetNewest)
				swallow(err)

				host, _, err := net.SplitHostPort(leader.Addr())
				swallow(err)

				if !patternMatched(host, this.host) {
					continue
				}

				if _, present := hostOffsets[host]; !present {
					hostOffsets[host] = hostOffsetInfo{host: host, offsetMap: make(map[string]map[structs.TopicPartition]int64)}
				}
				if _, present := hostOffsets[host].offsetMap[zkcluster.Name()]; !present {
					hostOffsets[host].offsetMap[zkcluster.Name()] = make(map[structs.TopicPartition]int64)
				}

				tp := structs.TopicPartition{Topic: topic, PartitionID: partitionID}
				hostOffsets[host].offsetMap[zkcluster.Name()][tp] = latestOffset
			}
		}

		this.hostOffsetCh <- hostOffsets

		kfk.RefreshMetadata(topics...)
	}
}

func (*Balance) Synopsis() string {
	return "Balance topics distribution according to load instead of count"
}

func (this *Balance) Help() string {
	help := fmt.Sprintf(`
Usage: %s balance [options]

    %s

Options:

    -z zone
      Default %s

    -c cluster pattern

    -host broker ip

    -sum
      Display in summary mode.

`, this.Cmd, this.Synopsis(), ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<commit_msg>balance is very helpful now<commit_after>package command

import (
	"flag"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/cmd/kateway/structs"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/pmylund/sortutil"
	"github.com/ryanuber/columnize"
)

type hostLoadInfo struct {
	host            string
	cluster         string
	topicPartitions []structs.TopicPartition
	qps             int64
}

type hostOffsetInfo struct {
	host      string
	offsetMap map[string]map[structs.TopicPartition]int64 // cluster:tp:offset
}

func (ho hostOffsetInfo) Clusters() []string {
	var r []string
	for cluster, _ := range ho.offsetMap {
		r = append(r, cluster)
	}
	return r
}

func (ho hostOffsetInfo) Total() (t int64) {
	for _, tps := range ho.offsetMap {
		for _, qps := range tps {
			t += qps
		}
	}

	return
}

func (ho hostOffsetInfo) ClusterPartitions(cluster string) int {
	return len(ho.offsetMap[cluster])
}

func (ho hostOffsetInfo) ClusterTotal(cluster string) (t int64) {
	for _, qps := range ho.offsetMap[cluster] {
		t += qps
	}

	return
}

type Balance struct {
	Ui  cli.Ui
	Cmd string

	zone, cluster string
	interval      time.Duration
	detailMode    bool
	host          string

	offsets     map[string]int64 // host => offset sum TODO
	lastOffsets map[string]int64

	lastHostOffsets map[string]hostOffsetInfo

	hostOffsetCh chan map[string]hostOffsetInfo // key is host
	signalsCh    map[string]chan struct{}
}

func (this *Balance) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("balance", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&this.cluster, "c", "", "")
	cmdFlags.DurationVar(&this.interval, "i", time.Second*5, "")
	cmdFlags.StringVar(&this.host, "host", "", "")
	cmdFlags.BoolVar(&this.detailMode, "d", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	this.signalsCh = make(map[string]chan struct{})
	this.hostOffsetCh = make(chan map[string]hostOffsetInfo)

	this.lastHostOffsets = make(map[string]hostOffsetInfo)
	this.offsets = make(map[string]int64)
	this.lastOffsets = make(map[string]int64)

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), this.cluster) {
			return
		}

		this.signalsCh[zkcluster.Name()] = make(chan struct{})

		go this.clusterTopProducers(zkcluster)
	})

	this.drawBalance()

	return
}

func (this *Balance) startAll() {
	for _, ch := range this.signalsCh {
		ch <- struct{}{}
	}
}

func (this *Balance) collectAll(seq int) {
	for _, _ = range this.signalsCh {
		offsets := <-this.hostOffsetCh
		if seq == 0 {
			// record into lastHostOffsets
			for host, offsetInfo := range offsets {
				if _, present := this.lastHostOffsets[host]; !present {
					this.lastHostOffsets[host] = hostOffsetInfo{host: host, offsetMap: make(map[string]map[structs.TopicPartition]int64)}
				}

				for cluster, tps := range offsetInfo.offsetMap {
					if _, present := this.lastHostOffsets[host].offsetMap[cluster]; !present {
						this.lastHostOffsets[host].offsetMap[cluster] = make(map[structs.TopicPartition]int64)
					}

					for tp, off := range tps {
						// 1st loop, offset
						this.lastHostOffsets[host].offsetMap[cluster][tp] = off
					}
				}
			}
		} else {
			for host, offsetInfo := range offsets {
				for cluster, tps := range offsetInfo.offsetMap {
					for tp, off := range tps {
						// 2nd loop, qps
						// FIXME hard coding
						this.lastHostOffsets[host].offsetMap[cluster][tp] = (off - this.lastHostOffsets[host].offsetMap[cluster][tp]) / int64(this.interval.Seconds())
					}
				}
			}
		}
	}

}

func (this *Balance) drawBalance() {
	for i := 0; i < 2; i++ {
		this.startAll()
		time.Sleep(this.interval)
		this.collectAll(i)
	}

	type hostTps struct {
		host string
		tps  int64
	}
	var sortedHosts []hostTps
	for host, info := range this.lastHostOffsets {
		sortedHosts = append(sortedHosts, hostTps{host, info.Total()})
	}
	sortutil.AscByField(sortedHosts, "tps")

	var hosts []string
	for _, h := range sortedHosts {
		hosts = append(hosts, h.host)
	}

	if !this.detailMode {
		this.drawSummary(hosts)
		return
	}

	this.drawDetail(hosts)
}

func (this *Balance) drawDetail(sortedHosts []string) {
	type hostSummary struct {
		cluster string
		tp      structs.TopicPartition
		qps     int64
	}

	for _, host := range sortedHosts {
		var summary []hostSummary
		offsetInfo := this.lastHostOffsets[host]

		for cluster, tps := range offsetInfo.offsetMap {
			for tp, qps := range tps {
				summary = append(summary, hostSummary{cluster, tp, qps})
			}
		}

		sortutil.DescByField(summary, "qps")

		this.Ui.Output(color.Green("%16s %8s %+v", host, gofmt.Comma(offsetInfo.Total()), offsetInfo.Clusters()))
		for _, sum := range summary {
			if sum.qps < 100 {
				continue
			}
			this.Ui.Output(fmt.Sprintf("    %30s %8s %s", sum.cluster, gofmt.Comma(sum.qps), sum.tp))
		}
	}

}

type clusterQps struct {
	cluster    string
	qps        int64
	partitions int
}

func (c clusterQps) String() string {
	partitions := fmt.Sprintf("%d", c.partitions)
	if c.qps < 1000 {
		return fmt.Sprintf("%s/%s/%d", c.cluster, partitions, c.qps)
	} else if c.qps < 5000 {
		return fmt.Sprintf("%s/%s/%s", c.cluster, partitions, color.Magenta("%d", c.qps))
	} else if c.qps < 10000 {
		return fmt.Sprintf("%s/%s/%s", c.cluster, partitions, color.Yellow("%d", c.qps))
	}

	return fmt.Sprintf("%s/%s/%s", c.cluster, partitions, color.Red("%d", c.qps))
}

func (this *Balance) drawSummary(sortedHosts []string) {
	lines := []string{"Broker|P|TPS|Cluster/OPS"}
	var totalTps int64
	var totalPartitions int
	for _, host := range sortedHosts {
		hostPartitions := 0
		offsetInfo := this.lastHostOffsets[host]
		var clusters []clusterQps
		for cluster, _ := range offsetInfo.offsetMap {
			clusterTps := offsetInfo.ClusterTotal(cluster)
			clusterPartitions := offsetInfo.ClusterPartitions(cluster)
			hostPartitions += clusterPartitions
			totalTps += clusterTps
			totalPartitions += clusterPartitions

			clusters = append(clusters, clusterQps{cluster, clusterTps, clusterPartitions})
		}

		sortutil.AscByField(clusters, "cluster")

		lines = append(lines, fmt.Sprintf("%s|%d|%s|%+v",
			host, hostPartitions, gofmt.Comma(offsetInfo.Total()), clusters))
	}
	this.Ui.Output(columnize.SimpleFormat(lines))
	this.Ui.Output(fmt.Sprintf("-Total- Hosts:%d Partitions:%d Tps:%s",
		len(sortedHosts), totalPartitions, gofmt.Comma(totalTps)))
}

func (this *Balance) clusterTopProducers(zkcluster *zk.ZkCluster) {
	kfk, err := sarama.NewClient(zkcluster.BrokerList(), sarama.NewConfig())
	if err != nil {
		return
	}
	defer kfk.Close()

	for {
		hostOffsets := make(map[string]hostOffsetInfo)

		topics, err := kfk.Topics()
		swallow(err)

		<-this.signalsCh[zkcluster.Name()]

		for _, topic := range topics {
			partions, err := kfk.WritablePartitions(topic)
			swallow(err)
			for _, partitionID := range partions {
				leader, err := kfk.Leader(topic, partitionID)
				swallow(err)

				latestOffset, err := kfk.GetOffset(topic, partitionID, sarama.OffsetNewest)
				swallow(err)

				host, _, err := net.SplitHostPort(leader.Addr())
				swallow(err)

				if !patternMatched(host, this.host) {
					continue
				}

				if _, present := hostOffsets[host]; !present {
					hostOffsets[host] = hostOffsetInfo{host: host, offsetMap: make(map[string]map[structs.TopicPartition]int64)}
				}
				if _, present := hostOffsets[host].offsetMap[zkcluster.Name()]; !present {
					hostOffsets[host].offsetMap[zkcluster.Name()] = make(map[structs.TopicPartition]int64)
				}

				tp := structs.TopicPartition{Topic: topic, PartitionID: partitionID}
				hostOffsets[host].offsetMap[zkcluster.Name()][tp] = latestOffset
			}
		}

		this.hostOffsetCh <- hostOffsets

		kfk.RefreshMetadata(topics...)
	}
}

func (*Balance) Synopsis() string {
	return "Balance topics distribution according to load instead of count"
}

func (this *Balance) Help() string {
	help := fmt.Sprintf(`
Usage: %s balance [options]

    %s

Options:

    -z zone
      Default %s

    -c cluster pattern

    -host broker ip

    -d
      Display in detailed mode.

`, this.Cmd, this.Synopsis(), ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<|endoftext|>