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

package cmd

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"path/filepath"
	"runtime"
	"sync"

	"github.com/cespare/xxhash/v2"
	"github.com/shenwei356/bio/seq"
	"github.com/shenwei356/bio/seqio/fastx"
	"github.com/shenwei356/xopen"
	"github.com/spf13/cobra"
	"github.com/twotwotwo/sorts/sortutil"
)

// sumCmd represents the sum command
var sumCmd = &cobra.Command{
	Use:   "sum",
	Short: "compute message digest for all sequences in FASTA/Q files",
	Long: `compute message digest for all sequences in FASTA/Q files

Attentions:
  1. Sequence headers and qualities are skipped, only sequences matter.
  2. The order of sequences records does not matter.
  3. Circular complete genomes are supported with the flag -c/--circular.
     - The same genomes with different start positions or in reverse
       complement strand will not affect the result.
     - The message digest would change with different values of k-mer size.
  4. Multiple files are processed in parallel (-j/--threads).

Method:
  1. Converting the sequences to low cases, optionally removing gaps (-g).
  2. Computing the hash (xxhash) for all sequences or k-mers of a circular
     complete genome (-c/--circular).
  3. Sorting all hash values, for ignoring the order of sequences.
  4. Computing MD5 digest from the hash values, sequences length, and
     the number of sequences.

`,
	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)
		alphabet := config.Alphabet
		idRegexp := config.IDRegexp
		// lineWidth := config.LineWidth
		outFile := config.OutFile
		seq.AlphabetGuessSeqLengthThreshold = config.AlphabetGuessSeqLength
		seq.ValidateSeq = false
		runtime.GOMAXPROCS(config.Threads)

		basename := getFlagBool(cmd, "basename")
		circular := getFlagBool(cmd, "circular")
		k := getFlagPositiveInt(cmd, "kmer-size")
		removeGaps := getFlagBool(cmd, "remove-gaps")
		gapLetters := getFlagString(cmd, "gap-letters")

		files := getFileListFromArgsAndFile(cmd, args, true, "infile-list", true)

		outfh, err := xopen.Wopen(outFile)
		checkError(err)
		defer outfh.Close()

		tokens := make(chan int, config.Threads)
		done := make(chan int)
		var wg sync.WaitGroup

		type Aresult struct {
			ok     bool
			id     uint64
			result *SumResult
		}
		ch := make(chan *Aresult, config.Threads)

		go func() {
			m := make(map[uint64]*Aresult, config.Threads)
			var id, _id uint64
			var ok bool
			var _r *Aresult

			id = 1
			for r := range ch {
				_id = r.id

				if _id == id { // right there
					if r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", r.result.Digest, r.result.File)
						// fmt.Fprintf(outfh, "%s\t%d\t%d\t%s\n", r.result.Digest, r.result.SeqLen, r.result.SeqNum, r.result.File)
						outfh.Flush()
					}
					id++
					continue
				}

				m[_id] = r // save for later check

				if _r, ok = m[id]; ok { // check buffered
					if r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", _r.result.Digest, _r.result.File)
						// fmt.Fprintf(outfh, "%s\t%d\t%d\t%s\n", _r.result.Digest, _r.result.SeqLen, _r.result.SeqNum, _r.result.File)
						outfh.Flush()
					}
					delete(m, id)
					id++
				}
			}

			if len(m) > 0 {
				ids := make([]uint64, len(m))
				i := 0
				for _id = range m {
					ids[i] = _id
					i++
				}
				sortutil.Uint64s(ids)
				for _, _id = range ids {
					_r = m[_id]

					if _r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", _r.result.Digest, _r.result.File)
						// fmt.Fprintf(outfh, "%s\t%d\t%d\t%s\n", _r.result.Digest, _r.result.SeqLen, _r.result.SeqNum, _r.result.File)
						outfh.Flush()
					}
				}
			}
			done <- 1
		}()

		var id uint64
		for _, file := range files {
			tokens <- 1
			wg.Add(1)
			id++

			go func(file string, id uint64) {
				defer func() {
					<-tokens
					wg.Done()
				}()

				var n int // number of sequences
				var l int // lengths of all seqs
				var h uint64
				hashes := make([]uint64, 0, 1024)

				var record *fastx.Record
				var fastxReader *fastx.Reader
				var _seq *seq.Seq

				fastxReader, err = fastx.NewReader(alphabet, file, idRegexp)
				// checkError(err)
				if err != nil {
					ch <- &Aresult{
						id:     id,
						ok:     false,
						result: nil,
					}
					log.Warningf(fmt.Sprintf("%s: %s", file, err))
					return
				}

				if circular {
					// var ok bool

					// var iter *sketches.Iterator
					var rc *seq.Seq
					var s, src []byte
					var h2 uint64
					var i, j, e, l, end, originalLen int
					for {
						record, err = fastxReader.Read()
						if err != nil {
							if err == io.EOF {
								break
							}
							// checkError(err)
							// break
							ch <- &Aresult{
								id:     id,
								ok:     false,
								result: nil,
							}
							log.Warningf(fmt.Sprintf("%s: %s", file, err))
							return
						}

						_seq = record.Seq

						if k > len(_seq.Seq) {
							checkError(fmt.Errorf("k is too big for sequence of %d bp: %s", len(record.Seq.Seq), file))
						}

						if n >= 1 {
							checkError(fmt.Errorf("only one sequence is allowed for circular genome"))
						}

						if removeGaps {
							_seq.RemoveGapsInplace(gapLetters)
						}

						_seq.Seq = bytes.ToLower(_seq.Seq)

						// // ntHash shoud be simpler and faster, but some thing wrong I didn't fingure out.
						// iter, err = sketches.NewHashIterator(_seq, k, true, true)
						// for {
						// 	h, ok = iter.NextHash()
						// 	if !ok {
						// 		break
						// 	}

						// 	hashes = append(hashes, h)
						// }

						rc = _seq.RevCom()

						l = len(_seq.Seq)
						originalLen = l
						end = l - 1

						i = 0
						for {
							if i > end {
								break
							}
							e = i + k

							if e > originalLen {
								e = e - originalLen
								s = _seq.Seq[i:]
								s = append(s, _seq.Seq[0:e]...)

								j = l - i
								src = rc.Seq[l-e:]
								src = append(src, rc.Seq[0:j]...)
							} else {
								s = _seq.Seq[i : i+k]

								j = l - i
								src = rc.Seq[j-k : j]
							}
							// fmt.Println(i, string(s), string(src))

							h = xxhash.Sum64(s)
							h2 = xxhash.Sum64(src)
							if h < h2 {
								hashes = append(hashes, h)
							} else {
								hashes = append(hashes, h2)
							}

							i++
						}

						n++
						l += len(record.Seq.Seq)
					}
				} else {
					for {
						record, err = fastxReader.Read()
						if err != nil {
							if err == io.EOF {
								break
							}
							// checkError(err)
							// break
							ch <- &Aresult{
								id:     id,
								ok:     false,
								result: nil,
							}
							log.Warningf(fmt.Sprintf("%s: %s", file, err))
							return
						}

						_seq = record.Seq

						if removeGaps {
							_seq.RemoveGapsInplace(gapLetters)
						}

						h = xxhash.Sum64(bytes.ToLower(_seq.Seq))
						hashes = append(hashes, h)

						n++
						l += len(_seq.Seq)
					}
				}

				// sequences
				sortutil.Uint64s(hashes)

				var le = binary.LittleEndian
				buf := make([]byte, 8)
				di := xxhash.New()

				for _, h = range hashes {
					// fmt.Println(h)
					le.PutUint64(buf, h)
					di.Write(buf)
				}

				// sequence length
				le.PutUint64(buf, uint64(l))
				di.Write(buf)

				// sequence number
				le.PutUint64(buf, uint64(n))
				di.Write(buf)

				// sum up
				le.PutUint64(buf, di.Sum64())
				digest := md5.Sum(buf)

				// return result
				if basename {
					file = filepath.Base(file)
				}
				ch <- &Aresult{
					id: id,
					ok: true,
					result: &SumResult{
						File:   file,
						SeqNum: n,
						SeqLen: l,
						Digest: hex.EncodeToString(digest[:]),
					},
				}
			}(file, id)
		}

		wg.Wait()
		close(ch)
		<-done
	},
}

func init() {
	RootCmd.AddCommand(sumCmd)

	sumCmd.Flags().BoolP("circular", "c", false, "the file contains a single cicular genome sequence")
	sumCmd.Flags().IntP("kmer-size", "k", 1000, "k-mer size for processing circular genomes")
	sumCmd.Flags().BoolP("basename", "b", false, "only output basename of files")
	sumCmd.Flags().BoolP("remove-gaps", "g", false, "remove gaps")
	sumCmd.Flags().StringP("gap-letters", "G", "- 	.", "gap letters")
}

type SumResult struct {
	File   string
	SeqNum int
	SeqLen int
	Digest string
}
<commit_msg>seqkit sum: new flag -a/--all<commit_after>// Copyright © 2016-2019 Wei Shen <shenwei356@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"path/filepath"
	"runtime"
	"sync"

	"github.com/cespare/xxhash/v2"
	"github.com/shenwei356/bio/seq"
	"github.com/shenwei356/bio/seqio/fastx"
	"github.com/shenwei356/xopen"
	"github.com/spf13/cobra"
	"github.com/twotwotwo/sorts/sortutil"
)

// sumCmd represents the sum command
var sumCmd = &cobra.Command{
	Use:   "sum",
	Short: "compute message digest for all sequences in FASTA/Q files",
	Long: `compute message digest for all sequences in FASTA/Q files

Attentions:
  1. Sequence headers and qualities are skipped, only sequences matter.
  2. The order of sequences records does not matter.
  3. Circular complete genomes are supported with the flag -c/--circular.
     - The same genomes with different start positions or in reverse
       complement strand will not affect the result.
     - The message digest would change with different values of k-mer size.
  4. Multiple files are processed in parallel (-j/--threads).

Method:
  1. Converting the sequences to low cases, optionally removing gaps (-g).
  2. Computing the hash (xxhash) for all sequences or k-mers of a circular
     complete genome (-c/--circular).
  3. Sorting all hash values, for ignoring the order of sequences.
  4. Computing MD5 digest from the hash values, sequences length, and
     the number of sequences.

`,
	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)
		alphabet := config.Alphabet
		idRegexp := config.IDRegexp
		// lineWidth := config.LineWidth
		outFile := config.OutFile
		seq.AlphabetGuessSeqLengthThreshold = config.AlphabetGuessSeqLength
		seq.ValidateSeq = false
		runtime.GOMAXPROCS(config.Threads)

		basename := getFlagBool(cmd, "basename")
		circular := getFlagBool(cmd, "circular")
		k := getFlagPositiveInt(cmd, "kmer-size")
		removeGaps := getFlagBool(cmd, "remove-gaps")
		gapLetters := getFlagString(cmd, "gap-letters")
		all := getFlagBool(cmd, "all")

		files := getFileListFromArgsAndFile(cmd, args, true, "infile-list", true)

		outfh, err := xopen.Wopen(outFile)
		checkError(err)
		defer outfh.Close()

		tokens := make(chan int, config.Threads)
		done := make(chan int)
		var wg sync.WaitGroup

		type Aresult struct {
			ok     bool
			id     uint64
			result *SumResult
		}
		ch := make(chan *Aresult, config.Threads)

		go func() {
			m := make(map[uint64]*Aresult, config.Threads)
			var id, _id uint64
			var ok bool
			var _r *Aresult

			id = 1
			for r := range ch {
				_id = r.id

				if _id == id { // right there
					if r.ok {
						if all {
							fmt.Fprintf(outfh, "%s\t%s\t%d\t%d\n", r.result.Digest, r.result.File, r.result.SeqNum, r.result.SeqLen)
						} else {
							fmt.Fprintf(outfh, "%s\t%s\n", r.result.Digest, r.result.File)
						}

						outfh.Flush()
					}
					id++
					continue
				}

				m[_id] = r // save for later check

				if _r, ok = m[id]; ok { // check buffered
					if r.ok {
						if all {
							fmt.Fprintf(outfh, "%s\t%s\t%d\t%d\n", _r.result.Digest, _r.result.File, _r.result.SeqNum, _r.result.SeqLen)
						} else {
							fmt.Fprintf(outfh, "%s\t%s\n", _r.result.Digest, _r.result.File)
						}
						outfh.Flush()
					}
					delete(m, id)
					id++
				}
			}

			if len(m) > 0 {
				ids := make([]uint64, len(m))
				i := 0
				for _id = range m {
					ids[i] = _id
					i++
				}
				sortutil.Uint64s(ids)
				for _, _id = range ids {
					_r = m[_id]

					if _r.ok {
						if all {
							fmt.Fprintf(outfh, "%s\t%s\t%d\t%d\n", _r.result.Digest, _r.result.File, _r.result.SeqNum, _r.result.SeqLen)
						} else {
							fmt.Fprintf(outfh, "%s\t%s\n", _r.result.Digest, _r.result.File)
						}
						outfh.Flush()
					}
				}
			}
			done <- 1
		}()

		var id uint64
		for _, file := range files {
			tokens <- 1
			wg.Add(1)
			id++

			go func(file string, id uint64) {
				defer func() {
					<-tokens
					wg.Done()
				}()

				var n int // number of sequences
				var l int // lengths of all seqs
				var h uint64
				hashes := make([]uint64, 0, 1024)

				var record *fastx.Record
				var fastxReader *fastx.Reader
				var _seq *seq.Seq

				fastxReader, err = fastx.NewReader(alphabet, file, idRegexp)
				// checkError(err)
				if err != nil {
					ch <- &Aresult{
						id:     id,
						ok:     false,
						result: nil,
					}
					log.Warningf(fmt.Sprintf("%s: %s", file, err))
					return
				}

				if circular {
					// var ok bool

					// var iter *sketches.Iterator
					var rc *seq.Seq
					var s, src []byte
					var h2 uint64
					var i, j, e, l, end, originalLen int
					for {
						record, err = fastxReader.Read()
						if err != nil {
							if err == io.EOF {
								break
							}
							// checkError(err)
							// break
							ch <- &Aresult{
								id:     id,
								ok:     false,
								result: nil,
							}
							log.Warningf(fmt.Sprintf("%s: %s", file, err))
							return
						}

						_seq = record.Seq

						if k > len(_seq.Seq) {
							checkError(fmt.Errorf("k is too big for sequence of %d bp: %s", len(record.Seq.Seq), file))
						}

						if n >= 1 {
							checkError(fmt.Errorf("only one sequence is allowed for circular genome"))
						}

						if removeGaps {
							_seq.RemoveGapsInplace(gapLetters)
						}

						_seq.Seq = bytes.ToLower(_seq.Seq)

						// // ntHash shoud be simpler and faster, but some thing wrong I didn't fingure out.
						// iter, err = sketches.NewHashIterator(_seq, k, true, true)
						// for {
						// 	h, ok = iter.NextHash()
						// 	if !ok {
						// 		break
						// 	}

						// 	hashes = append(hashes, h)
						// }

						rc = _seq.RevCom()

						l = len(_seq.Seq)
						originalLen = l
						end = l - 1

						i = 0
						for {
							if i > end {
								break
							}
							e = i + k

							if e > originalLen {
								e = e - originalLen
								s = _seq.Seq[i:]
								s = append(s, _seq.Seq[0:e]...)

								j = l - i
								src = rc.Seq[l-e:]
								src = append(src, rc.Seq[0:j]...)
							} else {
								s = _seq.Seq[i : i+k]

								j = l - i
								src = rc.Seq[j-k : j]
							}
							// fmt.Println(i, string(s), string(src))

							h = xxhash.Sum64(s)
							h2 = xxhash.Sum64(src)
							if h < h2 {
								hashes = append(hashes, h)
							} else {
								hashes = append(hashes, h2)
							}

							i++
						}

						n++
						l += len(record.Seq.Seq)
					}
				} else {
					for {
						record, err = fastxReader.Read()
						if err != nil {
							if err == io.EOF {
								break
							}
							// checkError(err)
							// break
							ch <- &Aresult{
								id:     id,
								ok:     false,
								result: nil,
							}
							log.Warningf(fmt.Sprintf("%s: %s", file, err))
							return
						}

						_seq = record.Seq

						if removeGaps {
							_seq.RemoveGapsInplace(gapLetters)
						}

						h = xxhash.Sum64(bytes.ToLower(_seq.Seq))
						hashes = append(hashes, h)

						n++
						l += len(_seq.Seq)
					}
				}

				// sequences
				sortutil.Uint64s(hashes)

				var le = binary.LittleEndian
				buf := make([]byte, 8)
				di := xxhash.New()

				for _, h = range hashes {
					// fmt.Println(h)
					le.PutUint64(buf, h)
					di.Write(buf)
				}

				// sequence length
				le.PutUint64(buf, uint64(l))
				di.Write(buf)

				// sequence number
				le.PutUint64(buf, uint64(n))
				di.Write(buf)

				// sum up
				le.PutUint64(buf, di.Sum64())
				digest := md5.Sum(buf)

				// return result
				if basename {
					file = filepath.Base(file)
				}
				ch <- &Aresult{
					id: id,
					ok: true,
					result: &SumResult{
						File:   file,
						SeqNum: n,
						SeqLen: l,
						Digest: hex.EncodeToString(digest[:]),
					},
				}
			}(file, id)
		}

		wg.Wait()
		close(ch)
		<-done
	},
}

func init() {
	RootCmd.AddCommand(sumCmd)

	sumCmd.Flags().BoolP("circular", "c", false, "the file contains a single cicular genome sequence")
	sumCmd.Flags().IntP("kmer-size", "k", 1000, "k-mer size for processing circular genomes")
	sumCmd.Flags().BoolP("basename", "b", false, "only output basename of files")
	sumCmd.Flags().BoolP("remove-gaps", "g", false, "remove gaps")
	sumCmd.Flags().StringP("gap-letters", "G", "- 	.", "gap letters")
	sumCmd.Flags().BoolP("all", "a", false, "all information, including the sequences length and the number of sequences")
}

type SumResult struct {
	File   string
	SeqNum int
	SeqLen int
	Digest string
}
<|endoftext|>


<commit_before>package serial

import (
	"fmt"
	"io"
	"os"
	"sync"
	"syscall"
	"unsafe"
)

type serialPort struct {
	f  *os.File
	fd syscall.Handle
	rl sync.Mutex
	wl sync.Mutex
	ro *syscall.Overlapped
	wo *syscall.Overlapped
}

type structDCB struct {
	DCBlength, BaudRate                            uint32
	flags                                          [4]byte
	wReserved, XonLim, XoffLim                     uint16
	ByteSize, Parity, StopBits                     byte
	XonChar, XoffChar, ErrorChar, EofChar, EvtChar byte
	wReserved1                                     uint16
}

type structTimeouts struct {
	ReadIntervalTimeout         uint32
	ReadTotalTimeoutMultiplier  uint32
	ReadTotalTimeoutConstant    uint32
	WriteTotalTimeoutMultiplier uint32
	WriteTotalTimeoutConstant   uint32
}

func openPort(name string, baud int) (rwc io.ReadWriteCloser, err os.Error) {
	if len(name) > 0 && name[0] != '\\' {
		name = "\\\\.\\" + name
	}

	const FILE_FLAGS_OVERLAPPED = 0x40000000
	h, e := syscall.CreateFile(syscall.StringToUTF16Ptr(name),
		syscall.GENERIC_READ|syscall.GENERIC_WRITE,
		0,
		nil,
		syscall.OPEN_EXISTING,
		syscall.FILE_ATTRIBUTE_NORMAL|FILE_FLAGS_OVERLAPPED,
		0)
	if e != 0 {
		err = &os.PathError{"open", name, os.Errno(e)}
		return
	}
	f := os.NewFile(h, name)
	defer func() {
		if err != nil {
			f.Close()
		}
	}()

	if err = setCommState(h, baud); err != nil {
		return
	}
	if err = setupComm(h, 64, 64); err != nil {
		return
	}
	if err = setCommTimeouts(h); err != nil {
		return
	}
	if err = setCommMask(h); err != nil {
		return
	}

	ro, err := newOverlapped()
	if err != nil {
		return
	}
	wo, err := newOverlapped()
	if err != nil {
		return
	}
	port := new(serialPort)
	port.f = f
	port.fd = h
	port.ro = ro
	port.wo = wo

	return port, nil
}

func (p *serialPort) Close() os.Error {
	return p.f.Close()
}

func (p *serialPort) Write(buf []byte) (int, os.Error) {
	p.wl.Lock()
	defer p.wl.Unlock()

	if err := resetEvent(p.wo.HEvent); err != nil {
		return 0, err
	}
	var n uint32
	e := syscall.WriteFile(p.fd, buf, &n, p.wo)
	if e != 0 && e != syscall.ERROR_IO_PENDING {
		return int(n), errno(uintptr(e))
	}
	return getOverlappedResult(p.fd, p.wo)
}

func (p *serialPort) Read(buf []byte) (int, os.Error) {
	if p == nil || p.f == nil {
		return 0, fmt.Errorf("Invalid port on read %v %v", p, p.f)
	}

	p.rl.Lock()
	defer p.rl.Unlock()

	if err := resetEvent(p.ro.HEvent); err != nil {
		return 0, err
	}
	var done uint32
	e := syscall.ReadFile(p.fd, buf, &done, p.ro)
	if e != 0 && e != syscall.ERROR_IO_PENDING {
		return int(done), errno(uintptr(e))
	}
	return getOverlappedResult(p.fd, p.ro)
}

var (
	nSetCommState,
	nSetCommTimeouts,
	nSetCommMask,
	nSetupComm,
	nGetOverlappedResult,
	nCreateEvent,
	nResetEvent uintptr
)

func init() {
	k32, err := syscall.LoadLibrary("kernel32.dll")
	if err != 0 {
		panic("LoadLibrary " + syscall.Errstr(err))
	}
	defer syscall.FreeLibrary(k32)

	nSetCommState = getProcAddr(k32, "SetCommState")
	nSetCommTimeouts = getProcAddr(k32, "SetCommTimeouts")
	nSetCommMask = getProcAddr(k32, "SetCommMask")
	nSetupComm = getProcAddr(k32, "SetupComm")
	nGetOverlappedResult = getProcAddr(k32, "GetOverlappedResult")
	nCreateEvent = getProcAddr(k32, "CreateEventW")
	nResetEvent = getProcAddr(k32, "ResetEvent")
}

func getProcAddr(lib syscall.Handle, name string) uintptr {
	addr, err := syscall.GetProcAddress(lib, name)
	if err != 0 {
		panic(name + " " + syscall.Errstr(err))
	}
	return addr
}

func errno(e uintptr) os.Error {
	if e != 0 {
		return os.Errno(e)
	}
	return os.Errno(syscall.EINVAL)
}

func setCommState(h syscall.Handle, baud int) os.Error {
	var params structDCB
	params.DCBlength = uint32(unsafe.Sizeof(params))

	params.flags[0] = 0x01  // fBinary
	params.flags[0] |= 0x10 // Assert DSR

	params.BaudRate = uint32(baud)
	params.ByteSize = 8

	r, _, e := syscall.Syscall(nSetCommState, 2, uintptr(h), uintptr(unsafe.Pointer(&params)), 0)
	if r == 0 {
		return errno(e)
	}
	return nil
}

func setCommTimeouts(h syscall.Handle) os.Error {
	var timeouts structTimeouts
	timeouts.ReadIntervalTimeout = 1<<32 - 1
	timeouts.ReadTotalTimeoutConstant = 0
	r, _, e := syscall.Syscall(nSetCommTimeouts, 2, uintptr(h), uintptr(unsafe.Pointer(&timeouts)), 0)
	if r == 0 {
		return errno(e)
	}
	return nil
}

func setupComm(h syscall.Handle, in, out int) os.Error {
	r, _, e := syscall.Syscall(nSetupComm, 3, uintptr(h), uintptr(in), uintptr(out))
	if r == 0 {
		return errno(e)
	}
	return nil
}

func setCommMask(h syscall.Handle) os.Error {
	const EV_RXCHAR = 0x0001
	r, _, e := syscall.Syscall(nSetCommMask, 2, uintptr(h), EV_RXCHAR, 0)
	if r == 0 {
		return errno(e)
	}
	return nil
}

func resetEvent(h syscall.Handle) os.Error {
	r, _, e := syscall.Syscall(nResetEvent, 1, uintptr(h), 0, 0)
	if r == 0 {
		return errno(e)
	}
	return nil
}

func newOverlapped() (*syscall.Overlapped, os.Error) {
	var overlapped syscall.Overlapped
	r, _, e := syscall.Syscall6(nCreateEvent, 4, 0, 1, 0, 0, 0, 0)
	if r == 0 {
		return nil, errno(e)
	}
	overlapped.HEvent = syscall.Handle(r)
	return &overlapped, nil
}

func getOverlappedResult(h syscall.Handle, overlapped *syscall.Overlapped) (int, os.Error) {
	var n int
	r, _, e := syscall.Syscall6(nGetOverlappedResult, 4,
		uintptr(h),
		uintptr(unsafe.Pointer(overlapped)),
		uintptr(unsafe.Pointer(&n)), 1, 0, 0)
	if r == 0 {
		return n, errno(e)
	}

	return n, nil
}
<commit_msg>Make windows reads block until 1 byte is read<commit_after>package serial

import (
	"fmt"
	"io"
	"os"
	"sync"
	"syscall"
	"unsafe"
)

type serialPort struct {
	f  *os.File
	fd syscall.Handle
	rl sync.Mutex
	wl sync.Mutex
	ro *syscall.Overlapped
	wo *syscall.Overlapped
}

type structDCB struct {
	DCBlength, BaudRate                            uint32
	flags                                          [4]byte
	wReserved, XonLim, XoffLim                     uint16
	ByteSize, Parity, StopBits                     byte
	XonChar, XoffChar, ErrorChar, EofChar, EvtChar byte
	wReserved1                                     uint16
}

type structTimeouts struct {
	ReadIntervalTimeout         uint32
	ReadTotalTimeoutMultiplier  uint32
	ReadTotalTimeoutConstant    uint32
	WriteTotalTimeoutMultiplier uint32
	WriteTotalTimeoutConstant   uint32
}

func openPort(name string, baud int) (rwc io.ReadWriteCloser, err os.Error) {
	if len(name) > 0 && name[0] != '\\' {
		name = "\\\\.\\" + name
	}

	const FILE_FLAGS_OVERLAPPED = 0x40000000
	h, e := syscall.CreateFile(syscall.StringToUTF16Ptr(name),
		syscall.GENERIC_READ|syscall.GENERIC_WRITE,
		0,
		nil,
		syscall.OPEN_EXISTING,
		syscall.FILE_ATTRIBUTE_NORMAL|FILE_FLAGS_OVERLAPPED,
		0)
	if e != 0 {
		err = &os.PathError{"open", name, os.Errno(e)}
		return
	}
	f := os.NewFile(h, name)
	defer func() {
		if err != nil {
			f.Close()
		}
	}()

	if err = setCommState(h, baud); err != nil {
		return
	}
	if err = setupComm(h, 64, 64); err != nil {
		return
	}
	if err = setCommTimeouts(h); err != nil {
		return
	}
	if err = setCommMask(h); err != nil {
		return
	}

	ro, err := newOverlapped()
	if err != nil {
		return
	}
	wo, err := newOverlapped()
	if err != nil {
		return
	}
	port := new(serialPort)
	port.f = f
	port.fd = h
	port.ro = ro
	port.wo = wo

	return port, nil
}

func (p *serialPort) Close() os.Error {
	return p.f.Close()
}

func (p *serialPort) Write(buf []byte) (int, os.Error) {
	p.wl.Lock()
	defer p.wl.Unlock()

	if err := resetEvent(p.wo.HEvent); err != nil {
		return 0, err
	}
	var n uint32
	e := syscall.WriteFile(p.fd, buf, &n, p.wo)
	if e != 0 && e != syscall.ERROR_IO_PENDING {
		return int(n), errno(uintptr(e))
	}
	return getOverlappedResult(p.fd, p.wo)
}

func (p *serialPort) Read(buf []byte) (int, os.Error) {
	if p == nil || p.f == nil {
		return 0, fmt.Errorf("Invalid port on read %v %v", p, p.f)
	}

	p.rl.Lock()
	defer p.rl.Unlock()

	if err := resetEvent(p.ro.HEvent); err != nil {
		return 0, err
	}
	var done uint32
	e := syscall.ReadFile(p.fd, buf, &done, p.ro)
	if e != 0 && e != syscall.ERROR_IO_PENDING {
		return int(done), errno(uintptr(e))
	}
	return getOverlappedResult(p.fd, p.ro)
}

var (
	nSetCommState,
	nSetCommTimeouts,
	nSetCommMask,
	nSetupComm,
	nGetOverlappedResult,
	nCreateEvent,
	nResetEvent uintptr
)

func init() {
	k32, err := syscall.LoadLibrary("kernel32.dll")
	if err != 0 {
		panic("LoadLibrary " + syscall.Errstr(err))
	}
	defer syscall.FreeLibrary(k32)

	nSetCommState = getProcAddr(k32, "SetCommState")
	nSetCommTimeouts = getProcAddr(k32, "SetCommTimeouts")
	nSetCommMask = getProcAddr(k32, "SetCommMask")
	nSetupComm = getProcAddr(k32, "SetupComm")
	nGetOverlappedResult = getProcAddr(k32, "GetOverlappedResult")
	nCreateEvent = getProcAddr(k32, "CreateEventW")
	nResetEvent = getProcAddr(k32, "ResetEvent")
}

func getProcAddr(lib syscall.Handle, name string) uintptr {
	addr, err := syscall.GetProcAddress(lib, name)
	if err != 0 {
		panic(name + " " + syscall.Errstr(err))
	}
	return addr
}

func errno(e uintptr) os.Error {
	if e != 0 {
		return os.Errno(e)
	}
	return os.Errno(syscall.EINVAL)
}

func setCommState(h syscall.Handle, baud int) os.Error {
	var params structDCB
	params.DCBlength = uint32(unsafe.Sizeof(params))

	params.flags[0] = 0x01  // fBinary
	params.flags[0] |= 0x10 // Assert DSR

	params.BaudRate = uint32(baud)
	params.ByteSize = 8

	r, _, e := syscall.Syscall(nSetCommState, 2, uintptr(h), uintptr(unsafe.Pointer(&params)), 0)
	if r == 0 {
		return errno(e)
	}
	return nil
}

func setCommTimeouts(h syscall.Handle) os.Error {
	var timeouts structTimeouts
	const MAXDWORD = 1<<32 - 1
	timeouts.ReadIntervalTimeout = MAXDWORD
	timeouts.ReadTotalTimeoutMultiplier = MAXDWORD
	timeouts.ReadTotalTimeoutConstant = MAXDWORD - 1

	/* From http://msdn.microsoft.com/en-us/library/aa363190(v=VS.85).aspx

	 For blocking I/O see below:

	 Remarks:

	 If an application sets ReadIntervalTimeout and
	 ReadTotalTimeoutMultiplier to MAXDWORD and sets
	 ReadTotalTimeoutConstant to a value greater than zero and
	 less than MAXDWORD, one of the following occurs when the
	 ReadFile function is called:

	 If there are any bytes in the input buffer, ReadFile returns
	       immediately with the bytes in the buffer.

	 If there are no bytes in the input buffer, ReadFile waits
               until a byte arrives and then returns immediately.

	 If no bytes arrive within the time specified by
	       ReadTotalTimeoutConstant, ReadFile times out.
	*/

	r, _, e := syscall.Syscall(nSetCommTimeouts, 2, uintptr(h), uintptr(unsafe.Pointer(&timeouts)), 0)
	if r == 0 {
		return errno(e)
	}
	return nil
}

func setupComm(h syscall.Handle, in, out int) os.Error {
	r, _, e := syscall.Syscall(nSetupComm, 3, uintptr(h), uintptr(in), uintptr(out))
	if r == 0 {
		return errno(e)
	}
	return nil
}

func setCommMask(h syscall.Handle) os.Error {
	const EV_RXCHAR = 0x0001
	r, _, e := syscall.Syscall(nSetCommMask, 2, uintptr(h), EV_RXCHAR, 0)
	if r == 0 {
		return errno(e)
	}
	return nil
}

func resetEvent(h syscall.Handle) os.Error {
	r, _, e := syscall.Syscall(nResetEvent, 1, uintptr(h), 0, 0)
	if r == 0 {
		return errno(e)
	}
	return nil
}

func newOverlapped() (*syscall.Overlapped, os.Error) {
	var overlapped syscall.Overlapped
	r, _, e := syscall.Syscall6(nCreateEvent, 4, 0, 1, 0, 0, 0, 0)
	if r == 0 {
		return nil, errno(e)
	}
	overlapped.HEvent = syscall.Handle(r)
	return &overlapped, nil
}

func getOverlappedResult(h syscall.Handle, overlapped *syscall.Overlapped) (int, os.Error) {
	var n int
	r, _, e := syscall.Syscall6(nGetOverlappedResult, 4,
		uintptr(h),
		uintptr(unsafe.Pointer(overlapped)),
		uintptr(unsafe.Pointer(&n)), 1, 0, 0)
	if r == 0 {
		return n, errno(e)
	}

	return n, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes 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 scheduling

import (
	"fmt"
	"math"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/uuid"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/test/e2e/framework"
	e2elog "k8s.io/kubernetes/test/e2e/framework/log"
	e2epod "k8s.io/kubernetes/test/e2e/framework/pod"
	testutils "k8s.io/kubernetes/test/utils"
	imageutils "k8s.io/kubernetes/test/utils/image"
)

var _ = SIGDescribe("Multi-AZ Clusters", func() {
	f := framework.NewDefaultFramework("multi-az")
	var zoneCount int
	var err error
	image := framework.ServeHostnameImage
	ginkgo.BeforeEach(func() {
		framework.SkipUnlessProviderIs("gce", "gke", "aws")
		if zoneCount <= 0 {
			zoneCount, err = getZoneCount(f.ClientSet)
			framework.ExpectNoError(err)
		}
		ginkgo.By(fmt.Sprintf("Checking for multi-zone cluster.  Zone count = %d", zoneCount))
		msg := fmt.Sprintf("Zone count is %d, only run for multi-zone clusters, skipping test", zoneCount)
		framework.SkipUnlessAtLeast(zoneCount, 2, msg)
		// TODO: SkipUnlessDefaultScheduler() // Non-default schedulers might not spread
	})
	ginkgo.It("should spread the pods of a service across zones", func() {
		SpreadServiceOrFail(f, (2*zoneCount)+1, image)
	})

	ginkgo.It("should spread the pods of a replication controller across zones", func() {
		SpreadRCOrFail(f, int32((2*zoneCount)+1), image, []string{"serve-hostname"})
	})
})

// SpreadServiceOrFail check that the pods comprising a service
// get spread evenly across available zones
func SpreadServiceOrFail(f *framework.Framework, replicaCount int, image string) {
	// First create the service
	serviceName := "test-service"
	serviceSpec := &v1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      serviceName,
			Namespace: f.Namespace.Name,
		},
		Spec: v1.ServiceSpec{
			Selector: map[string]string{
				"service": serviceName,
			},
			Ports: []v1.ServicePort{{
				Port:       80,
				TargetPort: intstr.FromInt(80),
			}},
		},
	}
	_, err := f.ClientSet.CoreV1().Services(f.Namespace.Name).Create(serviceSpec)
	framework.ExpectNoError(err)

	// Now create some pods behind the service
	podSpec := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:   serviceName,
			Labels: map[string]string{"service": serviceName},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:  "test",
					Image: imageutils.GetPauseImageName(),
				},
			},
		},
	}

	// Caution: StartPods requires at least one pod to replicate.
	// Based on the callers, replicas is always positive number: zoneCount >= 0 implies (2*zoneCount)+1 > 0.
	// Thus, no need to test for it. Once the precondition changes to zero number of replicas,
	// test for replicaCount > 0. Otherwise, StartPods panics.
	framework.ExpectNoError(testutils.StartPods(f.ClientSet, replicaCount, f.Namespace.Name, serviceName, *podSpec, false, e2elog.Logf))

	// Wait for all of them to be scheduled
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"service": serviceName}))
	pods, err := e2epod.WaitForPodsWithLabelScheduled(f.ClientSet, f.Namespace.Name, selector)
	framework.ExpectNoError(err)

	// Now make sure they're spread across zones
	zoneNames, err := framework.GetClusterZones(f.ClientSet)
	framework.ExpectNoError(err)
	ret, _ := checkZoneSpreading(f.ClientSet, pods, zoneNames.List())
	framework.ExpectEqual(ret, true)
}

// Find the name of the zone in which a Node is running
func getZoneNameForNode(node v1.Node) (string, error) {
	for key, value := range node.Labels {
		if key == v1.LabelZoneFailureDomain {
			return value, nil
		}
	}
	return "", fmt.Errorf("Zone name for node %s not found. No label with key %s",
		node.Name, v1.LabelZoneFailureDomain)
}

// Return the number of zones in which we have nodes in this cluster.
func getZoneCount(c clientset.Interface) (int, error) {
	zoneNames, err := framework.GetClusterZones(c)
	if err != nil {
		return -1, err
	}
	return len(zoneNames), nil
}

// Find the name of the zone in which the pod is scheduled
func getZoneNameForPod(c clientset.Interface, pod v1.Pod) (string, error) {
	ginkgo.By(fmt.Sprintf("Getting zone name for pod %s, on node %s", pod.Name, pod.Spec.NodeName))
	node, err := c.CoreV1().Nodes().Get(pod.Spec.NodeName, metav1.GetOptions{})
	framework.ExpectNoError(err)
	return getZoneNameForNode(*node)
}

// Determine whether a set of pods are approximately evenly spread
// across a given set of zones
func checkZoneSpreading(c clientset.Interface, pods *v1.PodList, zoneNames []string) (bool, error) {
	podsPerZone := make(map[string]int)
	for _, zoneName := range zoneNames {
		podsPerZone[zoneName] = 0
	}
	for _, pod := range pods.Items {
		if pod.DeletionTimestamp != nil {
			continue
		}
		zoneName, err := getZoneNameForPod(c, pod)
		framework.ExpectNoError(err)
		podsPerZone[zoneName] = podsPerZone[zoneName] + 1
	}
	minPodsPerZone := math.MaxInt32
	maxPodsPerZone := 0
	for _, podCount := range podsPerZone {
		if podCount < minPodsPerZone {
			minPodsPerZone = podCount
		}
		if podCount > maxPodsPerZone {
			maxPodsPerZone = podCount
		}
	}
	gomega.Expect(minPodsPerZone).To(gomega.BeNumerically("~", maxPodsPerZone, 1),
		"Pods were not evenly spread across zones.  %d in one zone and %d in another zone",
		minPodsPerZone, maxPodsPerZone)
	return true, nil
}

// SpreadRCOrFail Check that the pods comprising a replication
// controller get spread evenly across available zones
func SpreadRCOrFail(f *framework.Framework, replicaCount int32, image string, args []string) {
	name := "ubelite-spread-rc-" + string(uuid.NewUUID())
	ginkgo.By(fmt.Sprintf("Creating replication controller %s", name))
	controller, err := f.ClientSet.CoreV1().ReplicationControllers(f.Namespace.Name).Create(&v1.ReplicationController{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: f.Namespace.Name,
			Name:      name,
		},
		Spec: v1.ReplicationControllerSpec{
			Replicas: &replicaCount,
			Selector: map[string]string{
				"name": name,
			},
			Template: &v1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: map[string]string{"name": name},
				},
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						{
							Name:  name,
							Image: image,
							Args:  args,
							Ports: []v1.ContainerPort{{ContainerPort: 9376}},
						},
					},
				},
			},
		},
	})
	framework.ExpectNoError(err)
	// Cleanup the replication controller when we are done.
	defer func() {
		// Resize the replication controller to zero to get rid of pods.
		if err := framework.DeleteRCAndWaitForGC(f.ClientSet, f.Namespace.Name, controller.Name); err != nil {
			e2elog.Logf("Failed to cleanup replication controller %v: %v.", controller.Name, err)
		}
	}()
	// List the pods, making sure we observe all the replicas.
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": name}))
	pods, err := e2epod.PodsCreated(f.ClientSet, f.Namespace.Name, name, replicaCount)
	framework.ExpectNoError(err)

	// Wait for all of them to be scheduled
	ginkgo.By(fmt.Sprintf("Waiting for %d replicas of %s to be scheduled.  Selector: %v", replicaCount, name, selector))
	pods, err = e2epod.WaitForPodsWithLabelScheduled(f.ClientSet, f.Namespace.Name, selector)
	framework.ExpectNoError(err)

	// Now make sure they're spread across zones
	zoneNames, err := framework.GetClusterZones(f.ClientSet)
	framework.ExpectNoError(err)
	ret, _ := checkZoneSpreading(f.ClientSet, pods, zoneNames.List())
	framework.ExpectEqual(ret, true)
}
<commit_msg>Remove unnecessary return value check<commit_after>/*
Copyright 2015 The Kubernetes 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 scheduling

import (
	"fmt"
	"math"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/uuid"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/kubernetes/test/e2e/framework"
	e2elog "k8s.io/kubernetes/test/e2e/framework/log"
	e2epod "k8s.io/kubernetes/test/e2e/framework/pod"
	testutils "k8s.io/kubernetes/test/utils"
	imageutils "k8s.io/kubernetes/test/utils/image"
)

var _ = SIGDescribe("Multi-AZ Clusters", func() {
	f := framework.NewDefaultFramework("multi-az")
	var zoneCount int
	var err error
	image := framework.ServeHostnameImage
	ginkgo.BeforeEach(func() {
		framework.SkipUnlessProviderIs("gce", "gke", "aws")
		if zoneCount <= 0 {
			zoneCount, err = getZoneCount(f.ClientSet)
			framework.ExpectNoError(err)
		}
		ginkgo.By(fmt.Sprintf("Checking for multi-zone cluster.  Zone count = %d", zoneCount))
		msg := fmt.Sprintf("Zone count is %d, only run for multi-zone clusters, skipping test", zoneCount)
		framework.SkipUnlessAtLeast(zoneCount, 2, msg)
		// TODO: SkipUnlessDefaultScheduler() // Non-default schedulers might not spread
	})
	ginkgo.It("should spread the pods of a service across zones", func() {
		SpreadServiceOrFail(f, (2*zoneCount)+1, image)
	})

	ginkgo.It("should spread the pods of a replication controller across zones", func() {
		SpreadRCOrFail(f, int32((2*zoneCount)+1), image, []string{"serve-hostname"})
	})
})

// SpreadServiceOrFail check that the pods comprising a service
// get spread evenly across available zones
func SpreadServiceOrFail(f *framework.Framework, replicaCount int, image string) {
	// First create the service
	serviceName := "test-service"
	serviceSpec := &v1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      serviceName,
			Namespace: f.Namespace.Name,
		},
		Spec: v1.ServiceSpec{
			Selector: map[string]string{
				"service": serviceName,
			},
			Ports: []v1.ServicePort{{
				Port:       80,
				TargetPort: intstr.FromInt(80),
			}},
		},
	}
	_, err := f.ClientSet.CoreV1().Services(f.Namespace.Name).Create(serviceSpec)
	framework.ExpectNoError(err)

	// Now create some pods behind the service
	podSpec := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:   serviceName,
			Labels: map[string]string{"service": serviceName},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:  "test",
					Image: imageutils.GetPauseImageName(),
				},
			},
		},
	}

	// Caution: StartPods requires at least one pod to replicate.
	// Based on the callers, replicas is always positive number: zoneCount >= 0 implies (2*zoneCount)+1 > 0.
	// Thus, no need to test for it. Once the precondition changes to zero number of replicas,
	// test for replicaCount > 0. Otherwise, StartPods panics.
	framework.ExpectNoError(testutils.StartPods(f.ClientSet, replicaCount, f.Namespace.Name, serviceName, *podSpec, false, e2elog.Logf))

	// Wait for all of them to be scheduled
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"service": serviceName}))
	pods, err := e2epod.WaitForPodsWithLabelScheduled(f.ClientSet, f.Namespace.Name, selector)
	framework.ExpectNoError(err)

	// Now make sure they're spread across zones
	zoneNames, err := framework.GetClusterZones(f.ClientSet)
	framework.ExpectNoError(err)
	checkZoneSpreading(f.ClientSet, pods, zoneNames.List())
}

// Find the name of the zone in which a Node is running
func getZoneNameForNode(node v1.Node) (string, error) {
	for key, value := range node.Labels {
		if key == v1.LabelZoneFailureDomain {
			return value, nil
		}
	}
	return "", fmt.Errorf("Zone name for node %s not found. No label with key %s",
		node.Name, v1.LabelZoneFailureDomain)
}

// Return the number of zones in which we have nodes in this cluster.
func getZoneCount(c clientset.Interface) (int, error) {
	zoneNames, err := framework.GetClusterZones(c)
	if err != nil {
		return -1, err
	}
	return len(zoneNames), nil
}

// Find the name of the zone in which the pod is scheduled
func getZoneNameForPod(c clientset.Interface, pod v1.Pod) (string, error) {
	ginkgo.By(fmt.Sprintf("Getting zone name for pod %s, on node %s", pod.Name, pod.Spec.NodeName))
	node, err := c.CoreV1().Nodes().Get(pod.Spec.NodeName, metav1.GetOptions{})
	framework.ExpectNoError(err)
	return getZoneNameForNode(*node)
}

// Determine whether a set of pods are approximately evenly spread
// across a given set of zones
func checkZoneSpreading(c clientset.Interface, pods *v1.PodList, zoneNames []string) {
	podsPerZone := make(map[string]int)
	for _, zoneName := range zoneNames {
		podsPerZone[zoneName] = 0
	}
	for _, pod := range pods.Items {
		if pod.DeletionTimestamp != nil {
			continue
		}
		zoneName, err := getZoneNameForPod(c, pod)
		framework.ExpectNoError(err)
		podsPerZone[zoneName] = podsPerZone[zoneName] + 1
	}
	minPodsPerZone := math.MaxInt32
	maxPodsPerZone := 0
	for _, podCount := range podsPerZone {
		if podCount < minPodsPerZone {
			minPodsPerZone = podCount
		}
		if podCount > maxPodsPerZone {
			maxPodsPerZone = podCount
		}
	}
	gomega.Expect(minPodsPerZone).To(gomega.BeNumerically("~", maxPodsPerZone, 1),
		"Pods were not evenly spread across zones.  %d in one zone and %d in another zone",
		minPodsPerZone, maxPodsPerZone)
}

// SpreadRCOrFail Check that the pods comprising a replication
// controller get spread evenly across available zones
func SpreadRCOrFail(f *framework.Framework, replicaCount int32, image string, args []string) {
	name := "ubelite-spread-rc-" + string(uuid.NewUUID())
	ginkgo.By(fmt.Sprintf("Creating replication controller %s", name))
	controller, err := f.ClientSet.CoreV1().ReplicationControllers(f.Namespace.Name).Create(&v1.ReplicationController{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: f.Namespace.Name,
			Name:      name,
		},
		Spec: v1.ReplicationControllerSpec{
			Replicas: &replicaCount,
			Selector: map[string]string{
				"name": name,
			},
			Template: &v1.PodTemplateSpec{
				ObjectMeta: metav1.ObjectMeta{
					Labels: map[string]string{"name": name},
				},
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						{
							Name:  name,
							Image: image,
							Args:  args,
							Ports: []v1.ContainerPort{{ContainerPort: 9376}},
						},
					},
				},
			},
		},
	})
	framework.ExpectNoError(err)
	// Cleanup the replication controller when we are done.
	defer func() {
		// Resize the replication controller to zero to get rid of pods.
		if err := framework.DeleteRCAndWaitForGC(f.ClientSet, f.Namespace.Name, controller.Name); err != nil {
			e2elog.Logf("Failed to cleanup replication controller %v: %v.", controller.Name, err)
		}
	}()
	// List the pods, making sure we observe all the replicas.
	selector := labels.SelectorFromSet(labels.Set(map[string]string{"name": name}))
	pods, err := e2epod.PodsCreated(f.ClientSet, f.Namespace.Name, name, replicaCount)
	framework.ExpectNoError(err)

	// Wait for all of them to be scheduled
	ginkgo.By(fmt.Sprintf("Waiting for %d replicas of %s to be scheduled.  Selector: %v", replicaCount, name, selector))
	pods, err = e2epod.WaitForPodsWithLabelScheduled(f.ClientSet, f.Namespace.Name, selector)
	framework.ExpectNoError(err)

	// Now make sure they're spread across zones
	zoneNames, err := framework.GetClusterZones(f.ClientSet)
	framework.ExpectNoError(err)
	checkZoneSpreading(f.ClientSet, pods, zoneNames.List())
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/ym/oss-aliyun-go"
	"io/ioutil"
	"log"
	"mime"
	"os"
	"path/filepath"
	"strings"
)

var (
	a   oss.Auth
	o   *oss.OSS
	b   *oss.Bucket
	err error
)

const (
	DefaultContentType = "application/octet-stream"
)

func syncFiles() {
	a = oss.Auth{config.accessId, config.accessKey}
	o = &oss.OSS{a, config.endpoint, config.endpoint}

	b = o.Bucket(config.bucket)

	remoteFiles, _ := getRemoteFiles()
	localFiles, _ := getLocalFiles()

	for path, file := range localFiles {
		fn := config.prefix + strings.TrimLeft(strings.TrimPrefix(path, config.source), "/")
		if remote, ok := remoteFiles[fn]; ok {
			if file.Mode().IsDir() == false {
				data, err := getFile(path)
				if err == nil {
					hash := hashBytes(data)
					if compareHash(hash, remote.ETag) == false {
						log.Printf("Remote file %s existing, etag %s != %s.\n", fn, remote.ETag, hash)
						putFile(path, fn, data)
					}
				}
			}
		} else {
			if file.Mode().IsDir() == true {
				putDirectory(fn)
			} else {
				data, err := getFile(path)
				if err == nil {
					putFile(path, fn, data)
				}
			}
		}
	}
}

func putDirectory(filename string) {
	data := []byte("")
	err := b.Put(filename, data, "text/plain", oss.Private)
	if err != nil {
		log.Printf("Unable to create directory %s: %s.\n", filename, err.Error())
	} else {
		log.Printf("Created directory %s.\n", filename)
	}
}

func getFile(path string) ([]byte, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		log.Printf("Unable to read local file %s: %s.\n", path, err.Error())
	}
	return data, err
}

func putFile(path string, filename string, data []byte) {
	// get content type
	filetype := mime.TypeByExtension(filepath.Ext(path))
	if filetype == "" {
		log.Printf("Unable to detect content type of file %s, using %s.\n", filename, DefaultContentType)
		filetype = DefaultContentType
	}

	err = b.Put(filename, data, filetype, oss.Private)
	if err != nil {
		log.Printf("Unable to upload file %s: %s.\n", filename, err.Error())
	} else {
		log.Printf("Uploaded file: %s ... ", filename)
	}
}

func getRemoteFiles() (files map[string]oss.ContentInfo, err error) {
	var (
		marker = ""
		max    = 512
		get    func(string) error
	)

	files = make(map[string]oss.ContentInfo)

	get = func(prefix string) error {
		remoteFiles, err := b.List(prefix, "/", marker, max)
		if err != nil {
			return err
		}

		for _, p := range remoteFiles.CommonPrefixes {
			marker = ""
			get(p)
		}

		for _, file := range remoteFiles.Contents {
			files[file.Key] = file
		}

		if remoteFiles.IsTruncated {
			marker = remoteFiles.NextMarker
			get(prefix)
		}
		return nil
	}

	err = get(config.prefix)
	if err != nil {
		return files, err
	}

	return files, nil
}

func getLocalFiles() (files map[string]os.FileInfo, err error) {
	files = make(map[string]os.FileInfo)
	err = filepath.Walk(config.source, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			log.Printf("%s while looking up %s", err.Error(), path)
			return err
		}
		files[path] = info
		return nil
	})
	return files, err
}
<commit_msg>delete not existing file<commit_after>package main

import (
	"github.com/ym/oss-aliyun-go"
	"io/ioutil"
	"log"
	"mime"
	"os"
	"path/filepath"
	"strings"
)

var (
	a   oss.Auth
	o   *oss.OSS
	b   *oss.Bucket
	err error
)

const (
	DefaultContentType = "application/octet-stream"
)

func syncFiles() {
	a = oss.Auth{config.accessId, config.accessKey}
	o = &oss.OSS{a, config.endpoint, config.endpoint}

	b = o.Bucket(config.bucket)

	remoteFiles, _ := getRemoteFiles()
	localFiles, _ := getLocalFiles()

	for path, file := range localFiles {
		fn := config.prefix + strings.TrimLeft(strings.TrimPrefix(path, config.source), "/")
		if remote, ok := remoteFiles[fn]; ok {
			if file.Mode().IsDir() == false {
				data, err := getFile(path)
				if err == nil {
					hash := hashBytes(data)
					if compareHash(hash, remote.ETag) == false {
						log.Printf("Remote file %s existing, etag %s != %s.\n", fn, remote.ETag, hash)
						putFile(path, fn, data)
					}
				}
			}
		} else {
			if file.Mode().IsDir() == true {
				putDirectory(fn)
			} else {
				data, err := getFile(path)
				if err == nil {
					putFile(path, fn, data)
				}
			}
		}
	}

	if config.deleteIfNotFound {
		for path, _ := range remoteFiles {
			fn := strings.TrimLeft(strings.TrimPrefix(path, config.prefix), "/")
			local := config.source + "/" + fn
			if _, ok := localFiles[local]; ok == false {
				log.Printf("Remove remote file %s %s.\n", fn, local)
			}
		}
	}
}

func putDirectory(filename string) {
	data := []byte("")
	err := b.Put(filename, data, "text/plain", oss.Private)
	if err != nil {
		log.Printf("Unable to create directory %s: %s.\n", filename, err.Error())
	} else {
		log.Printf("Created directory %s.\n", filename)
	}
}

func getFile(path string) ([]byte, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		log.Printf("Unable to read local file %s: %s.\n", path, err.Error())
	}
	return data, err
}

func putFile(path string, filename string, data []byte) {
	// get content type
	filetype := mime.TypeByExtension(filepath.Ext(path))
	if filetype == "" {
		log.Printf("Unable to detect content type of file %s, using %s.\n", filename, DefaultContentType)
		filetype = DefaultContentType
	}

	err = b.Put(filename, data, filetype, oss.Private)
	if err != nil {
		log.Printf("Unable to upload file %s: %s.\n", filename, err.Error())
	} else {
		log.Printf("Uploaded file: %s ... ", filename)
	}
}

func getRemoteFiles() (files map[string]oss.ContentInfo, err error) {
	var (
		marker = ""
		max    = 512
		get    func(string) error
	)

	files = make(map[string]oss.ContentInfo)

	get = func(prefix string) error {
		remoteFiles, err := b.List(prefix, "/", marker, max)
		if err != nil {
			return err
		}

		for _, p := range remoteFiles.CommonPrefixes {
			marker = ""
			get(p)
		}

		for _, file := range remoteFiles.Contents {
			files[file.Key] = file
		}

		if remoteFiles.IsTruncated {
			marker = remoteFiles.NextMarker
			get(prefix)
		}
		return nil
	}

	err = get(config.prefix)
	if err != nil {
		return files, err
	}

	return files, nil
}

func getLocalFiles() (files map[string]os.FileInfo, err error) {
	files = make(map[string]os.FileInfo)
	err = filepath.Walk(config.source, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			log.Printf("%s while looking up %s", err.Error(), path)
			return err
		}
		files[path] = info
		return nil
	})
	return files, err
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"net/http"
	"strings"
)

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	//s.handlermutex.Lock()
	//defer s.handlermutex.Unlock()

	if !s.closed {
		s.clmutex.Lock()
		clients := s.clientlist
		s.clmutex.Unlock()

		for i := range clients {
			nick, ok := clients[i]["client_nickname"]
			if ok {
				if strings.Contains(clients[i]["client_type"], "0") {
					fmt.Fprintln(w, "
", nick, "
")
				}
			} else {
				fmt.Fprintln(w, "error: ", clients[i], ", empty map")
				if s.logger != nil {
					s.logger.Print("error: " + fmt.Sprint(clients[i]))
				}
			}
		}
	} else {
		http.Error(w, "internal error", 500)
	}

	return
}
<commit_msg>adding no one online msg<commit_after>package server

import (
	"fmt"
	"net/http"
	"strings"
)

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	//s.handlermutex.Lock()
	//defer s.handlermutex.Unlock()

	if !s.closed {
		s.clmutex.Lock()
		clients := s.clientlist
		s.clmutex.Unlock()

		if len(clients) > 0 {
			fmt.Fprintln(w, "
No one is online right now.
")
		} else {
			for i := range clients {
				nick, ok := clients[i]["client_nickname"]
				if ok {
					if strings.Contains(clients[i]["client_type"], "0") {
						fmt.Fprintln(w, "
", nick, "
")
					}
				} else {
					fmt.Fprintln(w, "error: ", clients[i], ", empty map")
					if s.logger != nil {
						s.logger.Print("error: " + fmt.Sprint(clients[i]))
					}
				}
			}
		}
	} else {
		http.Error(w, "internal error", 500)
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2015 Openprovider Authors. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

/*
Package tags 0.1.6

Definitions:

	- Tags - a strings which used for tag any object
	- Non-strict query tag - a string which match to string in tag ("new" -> "new")
	- Strict query tag - all strict query tags have prefix "+" for strict match ("+new")
    	and "-" for strict mismatch ("-old")

Rules:

	- All strict query tags applied with logical operator "AND" between each other
	- All non-strict query tags applied with logical operator "OR" between all tags in query

Queries:

	{"a","b"} this mean that we ask "a" OR "b" tag
	{"a","+b","+c"} this mean that we ask "a" OR ("b" AND "c") tag
	{"a","+b", "c", "-d"} this mean that we ask "a" OR "c" OR ("b" AND NOT "d") tag

Example:

	package main

	import (
		"fmt"

		"github.com/openprovider/tags"
	)

	// Product is struct with tags
	type Product struct {
		Name        string
		Description string
		Tags        tags.Tags
	}

	func main() {
		product := Product{
			Name:        "the tea",
			Description: "the boutle of ice black tea with sugar",
			Tags:        tags.Tags{"ice", "black", "sugar"},
		}

		fmt.Println("Product:", product.Description)

		// We ask for any tea "black" or "green"
		fmt.Println("Is this tea black or green?")
		query := tags.Tags{"black", "green"}

		if product.Tags.IsTagged(query) {
			fmt.Println("Yes, the tea is black.")
		} else {
			fmt.Println("No, the tea has not black or green options.")
		}

		// We ask fot strict match "green" and "sugar"
		fmt.Println("Is this tea green with sugar?")
		query = tags.Tags{"+green", "+sugar"}

		if product.Tags.IsTagged(query) {
			fmt.Println("Yes, the tea is green with sugar.")
		} else {
			fmt.Println("No, the tea with sugar, but is not green.")
		}

		// We ask for strict mismatch, not "ice"
		fmt.Println("Is this tea hot?")
		query = tags.Tags{"-ice"}

		if product.Tags.IsTagged(query) {
			fmt.Println("Yes, the tea is hot.")
		} else {
			fmt.Println("No, it is ice tea.")
		}
	}

Go Tags
*/
package tags

import (
	"strings"
)

// Tags is a slice of strings which used for tag any object
type Tags []string

// IsTagged - this method checks if Tags are matched with tags in query
func (t Tags) IsTagged(query Tags) bool {
	if len(query) == 0 {
		return true
	}
	if len(t) == 0 {
		return false
	}
	for _, qTag := range query {
		for _, tag := range t {
			if tag == qTag {
				return true
			}
		}
	}

	var strict uint8
	for _, qTag := range query {
		if strings.HasPrefix(qTag, "+") || strings.HasPrefix(qTag, "-") {
			strict++
		}
	}
	if strict == 0 {
		return false
	}
	for _, qTag := range query {
		if strings.HasPrefix(qTag, "+") {
			for _, tag := range t {
				if tag == strings.TrimPrefix(qTag, "+") {
					strict--
					break
				}
			}
		}
		if strings.HasPrefix(qTag, "-") {
			fit := true
			for _, tag := range t {
				if tag == strings.TrimPrefix(qTag, "-") {
					fit = false
					break
				}
			}
			if fit {
				strict--
			}
		}
	}

	return strict == 0
}
<commit_msg>Bumped version number to 0.2.0<commit_after>// Copyright 2015 Openprovider Authors. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

/*
Package tags 0.2.0

Definitions:

	- Tags - a strings which used for tag any object
	- Non-strict query tag - a string which match to string in tag ("new" -> "new")
	- Strict query tag - all strict query tags have prefix "+" for strict match ("+new")
    	and "-" for strict mismatch ("-old")

Rules:

	- All strict query tags applied with logical operator "AND" between each other
	- All non-strict query tags applied with logical operator "OR" between all tags in query

Queries:

	{"a","b"} this mean that we ask "a" OR "b" tag
	{"a","+b","+c"} this mean that we ask "a" OR ("b" AND "c") tag
	{"a","+b", "c", "-d"} this mean that we ask "a" OR "c" OR ("b" AND NOT "d") tag

Example:

	package main

	import (
		"fmt"

		"github.com/openprovider/tags"
	)

	// Product is struct with tags
	type Product struct {
		Name        string
		Description string
		Tags        tags.Tags
	}

	func main() {
		product := Product{
			Name:        "the tea",
			Description: "the boutle of ice black tea with sugar",
			Tags:        tags.Tags{"ice", "black", "sugar"},
		}

		fmt.Println("Product:", product.Description)

		// We ask for any tea "black" or "green"
		fmt.Println("Is this tea black or green?")
		query := tags.Tags{"black", "green"}

		if product.Tags.IsTagged(query) {
			fmt.Println("Yes, the tea is black.")
		} else {
			fmt.Println("No, the tea has not black or green options.")
		}

		// We ask fot strict match "green" and "sugar"
		fmt.Println("Is this tea green with sugar?")
		query = tags.Tags{"+green", "+sugar"}

		if product.Tags.IsTagged(query) {
			fmt.Println("Yes, the tea is green with sugar.")
		} else {
			fmt.Println("No, the tea with sugar, but is not green.")
		}

		// We ask for strict mismatch, not "ice"
		fmt.Println("Is this tea hot?")
		query = tags.Tags{"-ice"}

		if product.Tags.IsTagged(query) {
			fmt.Println("Yes, the tea is hot.")
		} else {
			fmt.Println("No, it is ice tea.")
		}
	}

Go Tags
*/
package tags

import (
	"strings"
)

// Tags is a slice of strings which used for tag any object
type Tags []string

// IsTagged - this method checks if Tags are matched with tags in query
func (t Tags) IsTagged(query Tags) bool {
	if len(query) == 0 {
		return true
	}
	if len(t) == 0 {
		return false
	}
	for _, qTag := range query {
		for _, tag := range t {
			if tag == qTag {
				return true
			}
		}
	}

	var strict uint8
	for _, qTag := range query {
		if strings.HasPrefix(qTag, "+") || strings.HasPrefix(qTag, "-") {
			strict++
		}
	}
	if strict == 0 {
		return false
	}
	for _, qTag := range query {
		if strings.HasPrefix(qTag, "+") {
			for _, tag := range t {
				if tag == strings.TrimPrefix(qTag, "+") {
					strict--
					break
				}
			}
		}
		if strings.HasPrefix(qTag, "-") {
			fit := true
			for _, tag := range t {
				if tag == strings.TrimPrefix(qTag, "-") {
					fit = false
					break
				}
			}
			if fit {
				strict--
			}
		}
	}

	return strict == 0
}
<|endoftext|>


<commit_before>package tail

import (
	"bufio"
	"fmt"
	"github.com/howeyc/fsnotify"
	"io"
	"log"
	"os"
	"path/filepath"
	"time"
)

type Line struct {
	Text     string
	UnixTime int64
}

type Tail struct {
	Filename string
	Lines    chan *Line

	maxlinesize int
	file        *os.File
	reader      *bufio.Reader
	watcher     *fsnotify.Watcher

	stop    chan bool
	created chan bool
}

// TailFile channels the lines of a logfile along with timestamp. If
// end is true, channel only newly added lines. If retry is true, tail
// the file name (not descriptor) and retry on file open/read errors.
func TailFile(filename string, maxlinesize int, end bool, retry bool) *Tail {
	t := &Tail{
		filename,
		make(chan *Line),
		maxlinesize,
		nil,
		nil,
		fileCreateWatcher(filename),
		make(chan bool),
		make(chan bool)}

	go t.tailFileSync(end, retry)

	return t
}

func (tail *Tail) Stop() {
	tail.stop <- true
	tail.close()
}

func (tail *Tail) close() {
	close(tail.Lines)
	tail.watcher.Close()
	if tail.file != nil {
		tail.file.Close()
	}
}

func (tail *Tail) reopen(wait bool) {
	if tail.file != nil {
		tail.file.Close()
	}
	for {
		var err error
		tail.file, err = os.Open(tail.Filename)
		if err != nil {
			if os.IsNotExist(err) && wait {
				for {
					evt := <-tail.watcher.Event
					if evt.Name == tail.Filename {
						break
					}
				}
				continue
			}
			// TODO: don't panic here
			panic(fmt.Sprintf("can't open file: %s", err))
		}
		return
	}
	return // unreachable
}

func (tail *Tail) readLine() ([]byte, error) {
	line, isPrefix, err := tail.reader.ReadLine()

	if isPrefix && err == nil {
		// line is longer than what we can accept.
		// ignore the rest of this line.
		for {
			_, isPrefix, err := tail.reader.ReadLine()
			if !isPrefix || err != nil {
				return line, err
			}
		}
	}
	return line, err
}

func (tail *Tail) tailFileSync(end bool, retry bool) {
	tail.reopen(retry)

	// Note: seeking to end happens only at the beginning; never
	// during subsequent re-opens.
	if end {
		_, err := tail.file.Seek(0, 2) // seek to end of the file
		if err != nil {
			// TODO: don't panic here
			panic(fmt.Sprintf("seek error: %s", err))
		}
	}

	tail.reader = bufio.NewReaderSize(tail.file, tail.maxlinesize)

	every2Seconds := time.Tick(2 * time.Second)

	for {
		line, err := tail.readLine()

		if err != nil && err != io.EOF {
			log.Println("Error reading file; skipping this file - ", err)
			tail.close()
			return
		}

		// sleep for 0.1s on inactive files, else we cause too much I/O activity
		if err == io.EOF {
			time.Sleep(100 * time.Millisecond)
		}

		if line != nil {
			tail.Lines <- &Line{string(line), getCurrentTime()}
		}

		select {
		case <-every2Seconds: // periodically stat the file to check for possibly deletion.
			if _, err := tail.file.Stat(); os.IsNotExist(err) {
				if retry {
					log.Printf("File %s has gone away; attempting to reopen it.\n", tail.Filename)
					tail.reopen(retry)
					tail.reader = bufio.NewReaderSize(tail.file, tail.maxlinesize)
					continue
				} else {
					log.Printf("File %s has gone away; skipping this file.\n", tail.Filename)
					tail.close()
					return
				}
			}
		case evt := <-tail.watcher.Event:
			if evt.Name == tail.Filename {
				log.Printf("File %s has been moved (logrotation?); reopening..", tail.Filename)
				tail.reopen(retry)
				tail.reader = bufio.NewReaderSize(tail.file, tail.maxlinesize)
				continue
			}
		case <-tail.stop: // stop the tailer if requested
			return
		default:
		}
	}
}

// returns the watcher for file create events
func fileCreateWatcher(filename string) *fsnotify.Watcher {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		panic(fmt.Sprintf("fsnotify error: %s", err))
	}

	// watch on parent directory because the file may not exit.
	err = watcher.WatchFlags(filepath.Dir(filename), fsnotify.FSN_CREATE)
	if err != nil {
		panic(fmt.Sprintf("fsnotify error on file: %s", err))
	}

	return watcher
}

// get current time in unix timestamp
func getCurrentTime() int64 {
	return time.Now().UTC().Unix()
}
<commit_msg>Bug #95787: warn and ignore tailing app logs if their base dir (container) goes away fast<commit_after>package tail

import (
	"bufio"
	"fmt"
	"github.com/howeyc/fsnotify"
	"io"
	"log"
	"os"
	"path/filepath"
	"time"
)

type Line struct {
	Text     string
	UnixTime int64
}

type Tail struct {
	Filename string
	Lines    chan *Line

	maxlinesize int
	file        *os.File
	reader      *bufio.Reader
	watcher     *fsnotify.Watcher

	stop    chan bool
	created chan bool
}

// TailFile channels the lines of a logfile along with timestamp. If
// end is true, channel only newly added lines. If retry is true, tail
// the file name (not descriptor) and retry on file open/read errors.
func TailFile(filename string, maxlinesize int, end bool, retry bool) (*Tail, error) {
	watcher, err := fileCreateWatcher(filename)
	if err != nil {
		return nil, err
	}
	t := &Tail{
		filename,
		make(chan *Line),
		maxlinesize,
		nil,
		nil,
		watcher,
		make(chan bool),
		make(chan bool)}

	go t.tailFileSync(end, retry)

	return t, nil
}

func (tail *Tail) Stop() {
	tail.stop <- true
	tail.close()
}

func (tail *Tail) close() {
	close(tail.Lines)
	tail.watcher.Close()
	if tail.file != nil {
		tail.file.Close()
	}
}

func (tail *Tail) reopen(wait bool) {
	if tail.file != nil {
		tail.file.Close()
	}
	for {
		var err error
		tail.file, err = os.Open(tail.Filename)
		if err != nil {
			if os.IsNotExist(err) && wait {
				for {
					evt := <-tail.watcher.Event
					if evt.Name == tail.Filename {
						break
					}
				}
				continue
			}
			// TODO: don't panic here
			panic(fmt.Sprintf("can't open file: %s", err))
		}
		return
	}
	return // unreachable
}

func (tail *Tail) readLine() ([]byte, error) {
	line, isPrefix, err := tail.reader.ReadLine()

	if isPrefix && err == nil {
		// line is longer than what we can accept.
		// ignore the rest of this line.
		for {
			_, isPrefix, err := tail.reader.ReadLine()
			if !isPrefix || err != nil {
				return line, err
			}
		}
	}
	return line, err
}

func (tail *Tail) tailFileSync(end bool, retry bool) {
	tail.reopen(retry)

	// Note: seeking to end happens only at the beginning; never
	// during subsequent re-opens.
	if end {
		_, err := tail.file.Seek(0, 2) // seek to end of the file
		if err != nil {
			// TODO: don't panic here
			panic(fmt.Sprintf("seek error: %s", err))
		}
	}

	tail.reader = bufio.NewReaderSize(tail.file, tail.maxlinesize)

	every2Seconds := time.Tick(2 * time.Second)

	for {
		line, err := tail.readLine()

		if err != nil && err != io.EOF {
			log.Println("Error reading file; skipping this file - ", err)
			tail.close()
			return
		}

		// sleep for 0.1s on inactive files, else we cause too much I/O activity
		if err == io.EOF {
			time.Sleep(100 * time.Millisecond)
		}

		if line != nil {
			tail.Lines <- &Line{string(line), getCurrentTime()}
		}

		select {
		case <-every2Seconds: // periodically stat the file to check for possibly deletion.
			if _, err := tail.file.Stat(); os.IsNotExist(err) {
				if retry {
					log.Printf("File %s has gone away; attempting to reopen it.\n", tail.Filename)
					tail.reopen(retry)
					tail.reader = bufio.NewReaderSize(tail.file, tail.maxlinesize)
					continue
				} else {
					log.Printf("File %s has gone away; skipping this file.\n", tail.Filename)
					tail.close()
					return
				}
			}
		case evt := <-tail.watcher.Event:
			if evt.Name == tail.Filename {
				log.Printf("File %s has been moved (logrotation?); reopening..", tail.Filename)
				tail.reopen(retry)
				tail.reader = bufio.NewReaderSize(tail.file, tail.maxlinesize)
				continue
			}
		case <-tail.stop: // stop the tailer if requested
			return
		default:
		}
	}
}

// returns the watcher for file create events
func fileCreateWatcher(filename string) (*fsnotify.Watcher, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}

	// watch on parent directory because the file may not exit.
	err = watcher.WatchFlags(filepath.Dir(filename), fsnotify.FSN_CREATE)
	if err != nil {
		return nil, err
	}

	return watcher, nil
}

// get current time in unix timestamp
func getCurrentTime() int64 {
	return time.Now().UTC().Unix()
}
<|endoftext|>


<commit_before>package kapacitor

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"time"

	"github.com/influxdata/kapacitor/pipeline"
)

// The type of a task
type TaskType int

const (
	StreamTask TaskType = iota
	BatchTask
)

func (t TaskType) String() string {
	switch t {
	case StreamTask:
		return "stream"
	case BatchTask:
		return "batch"
	default:
		return "unknown"
	}
}

type DBRP struct {
	Database        string `json:"db"`
	RetentionPolicy string `json:"rp"`
}

func CreateDBRPMap(dbrps []DBRP) map[DBRP]bool {
	dbMap := make(map[DBRP]bool, len(dbrps))
	for _, dbrp := range dbrps {
		dbMap[dbrp] = true
	}
	return dbMap
}

func (d DBRP) String() string {
	return fmt.Sprintf("%q.%q", d.Database, d.RetentionPolicy)
}

// The complete definition of a task, its name, pipeline and type.
type Task struct {
	Name             string
	Pipeline         *pipeline.Pipeline
	Type             TaskType
	DBRPs            []DBRP
	SnapshotInterval time.Duration
}

func (t *Task) Dot() []byte {
	return t.Pipeline.Dot(t.Name)
}

// ----------------------------------
// ExecutingTask

// A task that is ready for execution.
type ExecutingTask struct {
	tm      *TaskMaster
	Task    *Task
	source  Node
	outputs map[string]Output
	// node lookup from pipeline.ID -> Node
	lookup   map[pipeline.ID]Node
	nodes    []Node
	stopping chan struct{}
	wg       sync.WaitGroup
	logger   *log.Logger

	// Mutex for throughput var
	tmu        sync.RWMutex
	throughput float64
}

// Create a new  task from a defined kapacitor.
func NewExecutingTask(tm *TaskMaster, t *Task) (*ExecutingTask, error) {
	l := tm.LogService.NewLogger(fmt.Sprintf("[task:%s] ", t.Name), log.LstdFlags)
	et := &ExecutingTask{
		tm:      tm,
		Task:    t,
		outputs: make(map[string]Output),
		lookup:  make(map[pipeline.ID]Node),
		logger:  l,
	}
	err := et.link()
	if err != nil {
		return nil, err
	}
	return et, nil
}

// walks the entire pipeline applying function f.
func (et *ExecutingTask) walk(f func(n Node) error) error {
	for _, n := range et.nodes {
		err := f(n)
		if err != nil {
			return err
		}
	}
	return nil
}

// walks the entire pipeline in reverse order applying function f.
func (et *ExecutingTask) rwalk(f func(n Node) error) error {
	for i := len(et.nodes) - 1; i >= 0; i-- {
		err := f(et.nodes[i])
		if err != nil {
			return err
		}
	}
	return nil
}

// Link all the nodes together based on the task pipeline.
func (et *ExecutingTask) link() error {

	// Walk Pipeline and create equivalent executing nodes
	err := et.Task.Pipeline.Walk(func(n pipeline.Node) error {
		l := et.tm.LogService.NewLogger(
			fmt.Sprintf("[%s:%s] ", et.Task.Name, n.Name()),
			log.LstdFlags,
		)
		en, err := et.createNode(n, l)
		if err != nil {
			return err
		}
		et.lookup[n.ID()] = en
		// Save the walk order
		et.nodes = append(et.nodes, en)
		// Duplicate the Edges
		for _, p := range n.Parents() {
			ep := et.lookup[p.ID()]
			err := ep.linkChild(en)
			if err != nil {
				return err
			}
		}
		return err
	})
	if err != nil {
		return err
	}

	// The first node is always the source node
	et.source = et.nodes[0]
	return nil
}

// Start the task.
func (et *ExecutingTask) start(ins []*Edge, snapshot *TaskSnapshot) error {

	for _, in := range ins {
		et.source.addParentEdge(in)
	}
	validSnapshot := false
	if snapshot != nil {
		err := et.walk(func(n Node) error {
			_, ok := snapshot.NodeSnapshots[n.Name()]
			if !ok {
				return fmt.Errorf("task pipeline changed not using snapshot")
			}
			return nil
		})
		validSnapshot = err == nil
	}

	err := et.walk(func(n Node) error {
		if validSnapshot {
			n.start(snapshot.NodeSnapshots[n.Name()])
		} else {
			n.start(nil)
		}
		return nil
	})
	if err != nil {
		return err
	}
	et.stopping = make(chan struct{})
	if et.Task.SnapshotInterval > 0 {
		et.wg.Add(1)
		go et.runSnapshotter()
	}
	// Start calcThroughput
	et.wg.Add(1)
	go et.calcThroughput()
	return nil
}

func (et *ExecutingTask) stop() (err error) {
	close(et.stopping)
	et.walk(func(n Node) error {
		n.stop()
		e := n.Wait()
		if e != nil {
			err = e
		}
		return nil
	})
	et.wg.Wait()
	return
}

var ErrWrongTaskType = errors.New("wrong task type")

// Instruct source batch node to start querying and sending batches of data
func (et *ExecutingTask) StartBatching() error {
	if et.Task.Type != BatchTask {
		return ErrWrongTaskType
	}

	batcher := et.source.(*SourceBatchNode)

	err := et.checkDBRPs(batcher)
	if err != nil {
		batcher.Abort()
		return err
	}

	batcher.Start()
	return nil
}

func (et *ExecutingTask) BatchCount() (int, error) {
	if et.Task.Type != BatchTask {
		return 0, ErrWrongTaskType
	}

	batcher := et.source.(*SourceBatchNode)
	return batcher.Count(), nil
}

// Get the next `num` batch queries that the batcher will run starting at time `start`.
func (et *ExecutingTask) BatchQueries(start, stop time.Time) ([][]string, error) {
	if et.Task.Type != BatchTask {
		return nil, ErrWrongTaskType
	}

	batcher := et.source.(*SourceBatchNode)

	err := et.checkDBRPs(batcher)
	if err != nil {
		return nil, err
	}

	return batcher.Queries(start, stop), nil
}

// Check that the task allows access to DBRPs
func (et *ExecutingTask) checkDBRPs(batcher *SourceBatchNode) error {
	dbMap := CreateDBRPMap(et.Task.DBRPs)
	dbrps, err := batcher.DBRPs()
	if err != nil {
		return err
	}
	for _, dbrp := range dbrps {
		if !dbMap[dbrp] {
			return fmt.Errorf("batch query is not allowed to request data from %v", dbrp)
		}
	}
	return nil
}

// Stop all stats nodes
func (et *ExecutingTask) StopStats() {
	et.walk(func(n Node) error {
		if s, ok := n.(*StatsNode); ok {
			s.stopStats()
		}
		return nil
	})
}

// Wait till the task finishes and return any error
func (et *ExecutingTask) Wait() error {
	return et.rwalk(func(n Node) error {
		return n.Wait()
	})
}

// Get a named output.
func (et *ExecutingTask) GetOutput(name string) (Output, error) {
	if o, ok := et.outputs[name]; ok {
		return o, nil
	} else {
		return nil, fmt.Errorf("unknown output %s", name)
	}
}

// Register a named output.
func (et *ExecutingTask) registerOutput(name string, o Output) {
	et.outputs[name] = o
}

type ExecutionStats struct {
	TaskStats map[string]interface{}
	NodeStats map[string]map[string]interface{}
}

func (et *ExecutingTask) ExecutionStats() (ExecutionStats, error) {
	executionStats := ExecutionStats{
		TaskStats: make(map[string]interface{}),
		NodeStats: make(map[string]map[string]interface{}),
	}

	// Fill the task stats
	executionStats.TaskStats["throughput"] = et.getThroughput()

	// Fill the nodes stats
	err := et.walk(func(node Node) error {
		nodeStats := node.stats()

		// Add collected and emitted
		nodeStats["collected"] = node.collectedCount()
		nodeStats["emitted"] = node.emittedCount()

		executionStats.NodeStats[node.Name()] = nodeStats

		return nil
	})

	if err != nil {
		return executionStats, err
	}

	return executionStats, nil
}

// Return a graphviz .dot formatted byte array.
// Label edges with relavant execution information.
func (et *ExecutingTask) EDot(labels bool) []byte {

	var buf bytes.Buffer

	buf.Write([]byte("digraph "))
	buf.Write([]byte(et.Task.Name))
	buf.Write([]byte(" {\n"))
	// Write graph attributes
	unit := "points"
	if et.Task.Type == BatchTask {
		unit = "batches"
	}
	if labels {
		buf.Write([]byte(
			fmt.Sprintf("graph [label=\"Throughput: %0.2f %s/s\"];\n",
				et.getThroughput(),
				unit,
			),
		))
	} else {
		buf.Write([]byte(
			fmt.Sprintf("graph [throughput=\"%0.2f %s/s\"];\n",
				et.getThroughput(),
				unit,
			),
		))
	}

	et.walk(func(n Node) error {
		n.edot(&buf, labels)
		return nil
	})
	buf.Write([]byte("}"))

	return buf.Bytes()
}

// Return the current throughput value.
func (et *ExecutingTask) getThroughput() float64 {
	et.tmu.RLock()
	defer et.tmu.RUnlock()
	return et.throughput
}

func (et *ExecutingTask) calcThroughput() {
	defer et.wg.Done()
	var previous int64
	last := time.Now()
	ticker := time.NewTicker(time.Second)
	for {
		select {
		case <-ticker.C:
			current := et.source.collectedCount()
			now := time.Now()
			elapsed := float64(now.Sub(last)) / float64(time.Second)

			et.tmu.Lock()
			et.throughput = float64(current-previous) / elapsed
			et.tmu.Unlock()

			last = now
			previous = current

		case <-et.stopping:
			return
		}
	}
}

// Create a  node from a given pipeline node.
func (et *ExecutingTask) createNode(p pipeline.Node, l *log.Logger) (n Node, err error) {
	switch t := p.(type) {
	case *pipeline.StreamNode:
		n, err = newStreamNode(et, t, l)
	case *pipeline.SourceStreamNode:
		n, err = newSourceStreamNode(et, t, l)
	case *pipeline.SourceBatchNode:
		n, err = newSourceBatchNode(et, t, l)
	case *pipeline.BatchNode:
		n, err = newBatchNode(et, t, l)
	case *pipeline.WindowNode:
		n, err = newWindowNode(et, t, l)
	case *pipeline.HTTPOutNode:
		n, err = newHTTPOutNode(et, t, l)
	case *pipeline.InfluxDBOutNode:
		n, err = newInfluxDBOutNode(et, t, l)
	case *pipeline.MapNode:
		n, err = newMapNode(et, t, l)
	case *pipeline.ReduceNode:
		n, err = newReduceNode(et, t, l)
	case *pipeline.AlertNode:
		n, err = newAlertNode(et, t, l)
	case *pipeline.GroupByNode:
		n, err = newGroupByNode(et, t, l)
	case *pipeline.UnionNode:
		n, err = newUnionNode(et, t, l)
	case *pipeline.JoinNode:
		n, err = newJoinNode(et, t, l)
	case *pipeline.EvalNode:
		n, err = newEvalNode(et, t, l)
	case *pipeline.WhereNode:
		n, err = newWhereNode(et, t, l)
	case *pipeline.SampleNode:
		n, err = newSampleNode(et, t, l)
	case *pipeline.DerivativeNode:
		n, err = newDerivativeNode(et, t, l)
	case *pipeline.UDFNode:
		n, err = newUDFNode(et, t, l)
	case *pipeline.StatsNode:
		n, err = newStatsNode(et, t, l)
	case *pipeline.ShiftNode:
		n, err = newShiftNode(et, t, l)
	case *pipeline.NoOpNode:
		n, err = newNoOpNode(et, t, l)
	case *pipeline.InfluxQLNode:
		n, err = newInfluxQLNode(et, t, l)
	case *pipeline.LogNode:
		n, err = newLogNode(et, t, l)
	default:
		return nil, fmt.Errorf("unknown pipeline node type %T", p)
	}
	if err == nil && n != nil {
		n.init()
	}
	return n, err
}

type TaskSnapshot struct {
	NodeSnapshots map[string][]byte
}

func (et *ExecutingTask) Snapshot() (*TaskSnapshot, error) {
	snapshot := &TaskSnapshot{
		NodeSnapshots: make(map[string][]byte),
	}
	err := et.walk(func(n Node) error {
		data, err := n.snapshot()
		if err != nil {
			return err
		}
		snapshot.NodeSnapshots[n.Name()] = data
		return nil
	})
	if err != nil {
		return nil, err
	}
	return snapshot, nil
}

func (et *ExecutingTask) runSnapshotter() {
	defer et.wg.Done()
	// Wait random duration to splay snapshot events across interval
	select {
	case <-time.After(time.Duration(rand.Float64() * float64(et.Task.SnapshotInterval))):
	case <-et.stopping:
		return
	}
	ticker := time.NewTicker(et.Task.SnapshotInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			snapshot, err := et.Snapshot()
			if err != nil {
				et.logger.Println("E! failed to snapshot task", et.Task.Name, err)
				break
			}
			size := 0
			for _, data := range snapshot.NodeSnapshots {
				size += len(data)
			}
			// Only save the snapshot if it has content
			if size > 0 {
				err = et.tm.TaskStore.SaveSnapshot(et.Task.Name, snapshot)
				if err != nil {
					et.logger.Println("E! failed to save task snapshot", et.Task.Name, err)
				}
			}
		case <-et.stopping:
			return
		}
	}
}
<commit_msg>Closing calcThroughput ticker<commit_after>package kapacitor

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"math/rand"
	"sync"
	"time"

	"github.com/influxdata/kapacitor/pipeline"
)

// The type of a task
type TaskType int

const (
	StreamTask TaskType = iota
	BatchTask
)

func (t TaskType) String() string {
	switch t {
	case StreamTask:
		return "stream"
	case BatchTask:
		return "batch"
	default:
		return "unknown"
	}
}

type DBRP struct {
	Database        string `json:"db"`
	RetentionPolicy string `json:"rp"`
}

func CreateDBRPMap(dbrps []DBRP) map[DBRP]bool {
	dbMap := make(map[DBRP]bool, len(dbrps))
	for _, dbrp := range dbrps {
		dbMap[dbrp] = true
	}
	return dbMap
}

func (d DBRP) String() string {
	return fmt.Sprintf("%q.%q", d.Database, d.RetentionPolicy)
}

// The complete definition of a task, its name, pipeline and type.
type Task struct {
	Name             string
	Pipeline         *pipeline.Pipeline
	Type             TaskType
	DBRPs            []DBRP
	SnapshotInterval time.Duration
}

func (t *Task) Dot() []byte {
	return t.Pipeline.Dot(t.Name)
}

// ----------------------------------
// ExecutingTask

// A task that is ready for execution.
type ExecutingTask struct {
	tm      *TaskMaster
	Task    *Task
	source  Node
	outputs map[string]Output
	// node lookup from pipeline.ID -> Node
	lookup   map[pipeline.ID]Node
	nodes    []Node
	stopping chan struct{}
	wg       sync.WaitGroup
	logger   *log.Logger

	// Mutex for throughput var
	tmu        sync.RWMutex
	throughput float64
}

// Create a new  task from a defined kapacitor.
func NewExecutingTask(tm *TaskMaster, t *Task) (*ExecutingTask, error) {
	l := tm.LogService.NewLogger(fmt.Sprintf("[task:%s] ", t.Name), log.LstdFlags)
	et := &ExecutingTask{
		tm:      tm,
		Task:    t,
		outputs: make(map[string]Output),
		lookup:  make(map[pipeline.ID]Node),
		logger:  l,
	}
	err := et.link()
	if err != nil {
		return nil, err
	}
	return et, nil
}

// walks the entire pipeline applying function f.
func (et *ExecutingTask) walk(f func(n Node) error) error {
	for _, n := range et.nodes {
		err := f(n)
		if err != nil {
			return err
		}
	}
	return nil
}

// walks the entire pipeline in reverse order applying function f.
func (et *ExecutingTask) rwalk(f func(n Node) error) error {
	for i := len(et.nodes) - 1; i >= 0; i-- {
		err := f(et.nodes[i])
		if err != nil {
			return err
		}
	}
	return nil
}

// Link all the nodes together based on the task pipeline.
func (et *ExecutingTask) link() error {

	// Walk Pipeline and create equivalent executing nodes
	err := et.Task.Pipeline.Walk(func(n pipeline.Node) error {
		l := et.tm.LogService.NewLogger(
			fmt.Sprintf("[%s:%s] ", et.Task.Name, n.Name()),
			log.LstdFlags,
		)
		en, err := et.createNode(n, l)
		if err != nil {
			return err
		}
		et.lookup[n.ID()] = en
		// Save the walk order
		et.nodes = append(et.nodes, en)
		// Duplicate the Edges
		for _, p := range n.Parents() {
			ep := et.lookup[p.ID()]
			err := ep.linkChild(en)
			if err != nil {
				return err
			}
		}
		return err
	})
	if err != nil {
		return err
	}

	// The first node is always the source node
	et.source = et.nodes[0]
	return nil
}

// Start the task.
func (et *ExecutingTask) start(ins []*Edge, snapshot *TaskSnapshot) error {

	for _, in := range ins {
		et.source.addParentEdge(in)
	}
	validSnapshot := false
	if snapshot != nil {
		err := et.walk(func(n Node) error {
			_, ok := snapshot.NodeSnapshots[n.Name()]
			if !ok {
				return fmt.Errorf("task pipeline changed not using snapshot")
			}
			return nil
		})
		validSnapshot = err == nil
	}

	err := et.walk(func(n Node) error {
		if validSnapshot {
			n.start(snapshot.NodeSnapshots[n.Name()])
		} else {
			n.start(nil)
		}
		return nil
	})
	if err != nil {
		return err
	}
	et.stopping = make(chan struct{})
	if et.Task.SnapshotInterval > 0 {
		et.wg.Add(1)
		go et.runSnapshotter()
	}
	// Start calcThroughput
	et.wg.Add(1)
	go et.calcThroughput()
	return nil
}

func (et *ExecutingTask) stop() (err error) {
	close(et.stopping)
	et.walk(func(n Node) error {
		n.stop()
		e := n.Wait()
		if e != nil {
			err = e
		}
		return nil
	})
	et.wg.Wait()
	return
}

var ErrWrongTaskType = errors.New("wrong task type")

// Instruct source batch node to start querying and sending batches of data
func (et *ExecutingTask) StartBatching() error {
	if et.Task.Type != BatchTask {
		return ErrWrongTaskType
	}

	batcher := et.source.(*SourceBatchNode)

	err := et.checkDBRPs(batcher)
	if err != nil {
		batcher.Abort()
		return err
	}

	batcher.Start()
	return nil
}

func (et *ExecutingTask) BatchCount() (int, error) {
	if et.Task.Type != BatchTask {
		return 0, ErrWrongTaskType
	}

	batcher := et.source.(*SourceBatchNode)
	return batcher.Count(), nil
}

// Get the next `num` batch queries that the batcher will run starting at time `start`.
func (et *ExecutingTask) BatchQueries(start, stop time.Time) ([][]string, error) {
	if et.Task.Type != BatchTask {
		return nil, ErrWrongTaskType
	}

	batcher := et.source.(*SourceBatchNode)

	err := et.checkDBRPs(batcher)
	if err != nil {
		return nil, err
	}

	return batcher.Queries(start, stop), nil
}

// Check that the task allows access to DBRPs
func (et *ExecutingTask) checkDBRPs(batcher *SourceBatchNode) error {
	dbMap := CreateDBRPMap(et.Task.DBRPs)
	dbrps, err := batcher.DBRPs()
	if err != nil {
		return err
	}
	for _, dbrp := range dbrps {
		if !dbMap[dbrp] {
			return fmt.Errorf("batch query is not allowed to request data from %v", dbrp)
		}
	}
	return nil
}

// Stop all stats nodes
func (et *ExecutingTask) StopStats() {
	et.walk(func(n Node) error {
		if s, ok := n.(*StatsNode); ok {
			s.stopStats()
		}
		return nil
	})
}

// Wait till the task finishes and return any error
func (et *ExecutingTask) Wait() error {
	return et.rwalk(func(n Node) error {
		return n.Wait()
	})
}

// Get a named output.
func (et *ExecutingTask) GetOutput(name string) (Output, error) {
	if o, ok := et.outputs[name]; ok {
		return o, nil
	} else {
		return nil, fmt.Errorf("unknown output %s", name)
	}
}

// Register a named output.
func (et *ExecutingTask) registerOutput(name string, o Output) {
	et.outputs[name] = o
}

type ExecutionStats struct {
	TaskStats map[string]interface{}
	NodeStats map[string]map[string]interface{}
}

func (et *ExecutingTask) ExecutionStats() (ExecutionStats, error) {
	executionStats := ExecutionStats{
		TaskStats: make(map[string]interface{}),
		NodeStats: make(map[string]map[string]interface{}),
	}

	// Fill the task stats
	executionStats.TaskStats["throughput"] = et.getThroughput()

	// Fill the nodes stats
	err := et.walk(func(node Node) error {
		nodeStats := node.stats()

		// Add collected and emitted
		nodeStats["collected"] = node.collectedCount()
		nodeStats["emitted"] = node.emittedCount()

		executionStats.NodeStats[node.Name()] = nodeStats

		return nil
	})

	if err != nil {
		return executionStats, err
	}

	return executionStats, nil
}

// Return a graphviz .dot formatted byte array.
// Label edges with relavant execution information.
func (et *ExecutingTask) EDot(labels bool) []byte {

	var buf bytes.Buffer

	buf.Write([]byte("digraph "))
	buf.Write([]byte(et.Task.Name))
	buf.Write([]byte(" {\n"))
	// Write graph attributes
	unit := "points"
	if et.Task.Type == BatchTask {
		unit = "batches"
	}
	if labels {
		buf.Write([]byte(
			fmt.Sprintf("graph [label=\"Throughput: %0.2f %s/s\"];\n",
				et.getThroughput(),
				unit,
			),
		))
	} else {
		buf.Write([]byte(
			fmt.Sprintf("graph [throughput=\"%0.2f %s/s\"];\n",
				et.getThroughput(),
				unit,
			),
		))
	}

	et.walk(func(n Node) error {
		n.edot(&buf, labels)
		return nil
	})
	buf.Write([]byte("}"))

	return buf.Bytes()
}

// Return the current throughput value.
func (et *ExecutingTask) getThroughput() float64 {
	et.tmu.RLock()
	defer et.tmu.RUnlock()
	return et.throughput
}

func (et *ExecutingTask) calcThroughput() {
	defer et.wg.Done()
	var previous int64
	last := time.Now()
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			current := et.source.collectedCount()
			now := time.Now()
			elapsed := float64(now.Sub(last)) / float64(time.Second)

			et.tmu.Lock()
			et.throughput = float64(current-previous) / elapsed
			et.tmu.Unlock()

			last = now
			previous = current

		case <-et.stopping:
			return
		}
	}
}

// Create a  node from a given pipeline node.
func (et *ExecutingTask) createNode(p pipeline.Node, l *log.Logger) (n Node, err error) {
	switch t := p.(type) {
	case *pipeline.StreamNode:
		n, err = newStreamNode(et, t, l)
	case *pipeline.SourceStreamNode:
		n, err = newSourceStreamNode(et, t, l)
	case *pipeline.SourceBatchNode:
		n, err = newSourceBatchNode(et, t, l)
	case *pipeline.BatchNode:
		n, err = newBatchNode(et, t, l)
	case *pipeline.WindowNode:
		n, err = newWindowNode(et, t, l)
	case *pipeline.HTTPOutNode:
		n, err = newHTTPOutNode(et, t, l)
	case *pipeline.InfluxDBOutNode:
		n, err = newInfluxDBOutNode(et, t, l)
	case *pipeline.MapNode:
		n, err = newMapNode(et, t, l)
	case *pipeline.ReduceNode:
		n, err = newReduceNode(et, t, l)
	case *pipeline.AlertNode:
		n, err = newAlertNode(et, t, l)
	case *pipeline.GroupByNode:
		n, err = newGroupByNode(et, t, l)
	case *pipeline.UnionNode:
		n, err = newUnionNode(et, t, l)
	case *pipeline.JoinNode:
		n, err = newJoinNode(et, t, l)
	case *pipeline.EvalNode:
		n, err = newEvalNode(et, t, l)
	case *pipeline.WhereNode:
		n, err = newWhereNode(et, t, l)
	case *pipeline.SampleNode:
		n, err = newSampleNode(et, t, l)
	case *pipeline.DerivativeNode:
		n, err = newDerivativeNode(et, t, l)
	case *pipeline.UDFNode:
		n, err = newUDFNode(et, t, l)
	case *pipeline.StatsNode:
		n, err = newStatsNode(et, t, l)
	case *pipeline.ShiftNode:
		n, err = newShiftNode(et, t, l)
	case *pipeline.NoOpNode:
		n, err = newNoOpNode(et, t, l)
	case *pipeline.InfluxQLNode:
		n, err = newInfluxQLNode(et, t, l)
	case *pipeline.LogNode:
		n, err = newLogNode(et, t, l)
	default:
		return nil, fmt.Errorf("unknown pipeline node type %T", p)
	}
	if err == nil && n != nil {
		n.init()
	}
	return n, err
}

type TaskSnapshot struct {
	NodeSnapshots map[string][]byte
}

func (et *ExecutingTask) Snapshot() (*TaskSnapshot, error) {
	snapshot := &TaskSnapshot{
		NodeSnapshots: make(map[string][]byte),
	}
	err := et.walk(func(n Node) error {
		data, err := n.snapshot()
		if err != nil {
			return err
		}
		snapshot.NodeSnapshots[n.Name()] = data
		return nil
	})
	if err != nil {
		return nil, err
	}
	return snapshot, nil
}

func (et *ExecutingTask) runSnapshotter() {
	defer et.wg.Done()
	// Wait random duration to splay snapshot events across interval
	select {
	case <-time.After(time.Duration(rand.Float64() * float64(et.Task.SnapshotInterval))):
	case <-et.stopping:
		return
	}
	ticker := time.NewTicker(et.Task.SnapshotInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			snapshot, err := et.Snapshot()
			if err != nil {
				et.logger.Println("E! failed to snapshot task", et.Task.Name, err)
				break
			}
			size := 0
			for _, data := range snapshot.NodeSnapshots {
				size += len(data)
			}
			// Only save the snapshot if it has content
			if size > 0 {
				err = et.tm.TaskStore.SaveSnapshot(et.Task.Name, snapshot)
				if err != nil {
					et.logger.Println("E! failed to save task snapshot", et.Task.Name, err)
				}
			}
		case <-et.stopping:
			return
		}
	}
}
<|endoftext|>