<commit_before>// Copyright 2017, Kerby Shedden and the Muscato contributors.

// muscato_screen is an initial screening step used by Muscato to
// identify candidate matches of a set of reads into a set of target
// gene sequences.  The results of the screen may contain false
// positives, but will not contain any false negatives.
//
// The approach is to use a Bloom filter to sketch the reads based on
// the subsequences that appear at defined offsets within the reads.
// For example, if position 10 is an offset and we are looking at
// subequences of width 15, then the read subsequences from position
// 10 through position 25 are entered into a Bloom filter.  Then, we
// scan through every target gene looking for matches to the Bloom
// filter.  When a match occurs, the match position (in the target)
// and flanking sequences are saved for subequent checking against the
// full read sequence.
//
// A simple entropy check is used to avoid considering subsequences
// that could match large numbers of reads or genes (and hence would
// be uninformative).  Currently, this check is based on the number of
// distinct dinucleotide subsequences in the window (e.g. in the
// 15-mer in the example above).
//
// The results are saved in files named bmatch*.txt.sz, where * is the
// window number.
//
// The format of the bmatch files is:
//
// (window sequence) (left tail) (right tail) (gene id) (position)

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"path"
	"strings"
	"sync"

	"github.com/chmduquesne/rollinghash"
	"github.com/chmduquesne/rollinghash/buzhash32"
	"github.com/golang-collections/go-datastructures/bitarray"
	"github.com/golang/snappy"
	"github.com/kshedden/seqmatch/utils"
)

const (
	// Number of goroutines, should probably scale with the number
	// of aailable cores.
	concurrency int = 100
)

var (
	// A log
	logger *log.Logger

	// Configuration information
	config *utils.Config

	// All working files are stored here
	tmpdir string

	// Bitarrays that back the Bloom filters
	smp []bitarray.BitArray

	// Tables to produce independent running hashes
	tables [][256]uint32

	// Communicate results back to driver
	hitchan chan rec

	// Semaphore for limiting goroutines
	limit chan bool

	// Line length for output
	bufsize int
)

// genTables generates base hash functions for a collection of rolling hashes.
func genTables() {
	tables = make([][256]uint32, config.NumHash)
	for j := 0; j < config.NumHash; j++ {
		mp := make(map[uint32]bool)
		for i := 0; i < 256; i++ {
			for {
				x := uint32(rand.Int63())
				if !mp[x] {
					tables[j][i] = x
					mp[x] = true
					break
				}
			}
		}
	}
}

// buildBloom constructs bloom filters for each window
func buildBloom() {

	logger.Printf("Building Bloom sketch of read collection...")

	hashes := make([]rollinghash.Hash32, config.NumHash)
	for j := range hashes {
		hashes[j] = buzhash32.NewFromUint32Array(tables[j])
	}

	fname := path.Join(tmpdir, "reads_sorted.txt.sz")
	fid, err := os.Open(fname)
	if err != nil {
		logger.Print(err)
		panic(err)
	}
	defer fid.Close()
	snr := snappy.NewReader(fid)
	scanner := bufio.NewScanner(snr)
	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)

	// Workspace for sequence diversity checker
	wk := make([]int, 25)

	var j int
	for ; scanner.Scan(); j++ {

		if j%1000000 == 0 {
			logger.Printf("%d\n", j)
		}

		line := scanner.Bytes()
		seq := bytes.Fields(line)[0]

		for k := 0; k < len(config.Windows); k++ {
			q1 := config.Windows[k]
			q2 := q1 + config.WindowWidth
			if q2 > len(seq) {
				continue
			}
			seqw := seq[q1:q2]

			// Check entropy
			if utils.CountDinuc(seqw, wk) < config.MinDinuc {
				continue
			}

			// Update the Bloom filter for this sequence
			for _, ha := range hashes {
				ha.Reset()
				_, err = ha.Write(seqw)
				if err != nil {
					panic(err)
				}
				x := uint64(ha.Sum32()) % config.BloomSize
				err := smp[k].SetBit(x)
				if err != nil {
					logger.Print(err)
					panic(err)
				}
			}
		}
	}

	if err := scanner.Err(); err != nil {
		msg := fmt.Sprintf("Problem reading reads_sorted.txt.sz on line %d\n", j)
		os.Stderr.WriteString(msg)
		logger.Print(err)
		panic(err)
	}

	logger.Printf("Done constructing Bloom filters")
}

type rec struct {
	mseq  string
	left  string
	right string
	win   int
	tnum  int
	pos   uint32
}

// checkWin returns the indices of the Bloom filters that match the
// current state of the hashes.
func checkWin(ix []int, iw []uint64, hashes []rollinghash.Hash32) []int {

	// Get the hash states
	for j, ha := range hashes {
		iw[j] = uint64(ha.Sum32()) % config.BloomSize
	}

	ix = ix[0:0]

	// Loop over Bloom filters
	for k, ba := range smp {

		// Determine if the Bloom filter matches
		g := true
		for j := range hashes {
			f, err := ba.GetBit(iw[j])
			if err != nil {
				logger.Print(err)
				panic(err)
			}
			if !f {
				g = false
				break
			}
		}
		if g {
			ix = append(ix, k)
		}
	}

	return ix
}

// process one target sequence, runs concurrently with main loop.
func processseq(seq []byte, genenum int) {

	defer func() { <-limit }()

	hashes := make([]rollinghash.Hash32, config.NumHash)
	for j := range hashes {
		hashes[j] = buzhash32.NewFromUint32Array(tables[j])
	}

	// Initialize the hashes with the first window.
	hlen := config.WindowWidth
	if len(seq) < hlen {
		// Not long enough even for one window.
		return
	}
	for j := range hashes {
		_, err := hashes[j].Write(seq[0:hlen])
		if err != nil {
			panic(err)
		}
	}

	ix := make([]int, len(smp))
	iw := make([]uint64, config.NumHash)

	// Check if the initial window is a match
	ix = checkWin(ix, iw, hashes)
	for _, i := range ix {

		q1 := config.Windows[i]
		if q1 != 0 {
			// The only way the read can fit is if the
			// window starts at the beginning of the read.
			continue
		}
		q2 := q1 + config.WindowWidth

		jz := 100 - q2
		if jz > len(seq) {
			jz = len(seq)
		}
		hitchan <- rec{
			mseq:  string(seq[0:hlen]),
			left:  "",
			right: string(seq[hlen:jz]),
			tnum:  genenum,
			win:   i,
			pos:   0,
		}
	}

	// Check the rest of the windows
	for j := hlen; j < len(seq); j++ {

		for _, ha := range hashes {
			ha.Roll(seq[j])
		}
		ix = checkWin(ix, iw, hashes)

		// Process a match
		for _, i := range ix {

			q1 := config.Windows[i]
			q2 := q1 + config.WindowWidth
			if j < q2-1 {
				// The read would not fit
				continue
			}

			// Matching sequence is jx:jy
			jx := j - hlen + 1
			jy := j + 1

			// Left tail is jw:jx
			jw := jx - q1

			// Right tail is jy:jz
			jz := jy + config.MaxReadLength - q2
			if jz > len(seq) {
				// May not be long enough to fit, but
				// we don't know until we merge.
				jz = len(seq)
			}

			if jw >= 0 {
				hitchan <- rec{
					mseq:  string(seq[jx:jy]),
					left:  string(seq[jw:jx]),
					right: string(seq[jy:jz]),
					tnum:  genenum,
					win:   i,
					pos:   uint32(j - hlen + 1),
				}
			}
		}
	}
}

// Retrieve the results and write to disk
func harvest(wg *sync.WaitGroup) {

	var wtrs []io.Writer
	var allwtrs []io.Closer
	for k := 0; k < len(config.Windows); k++ {
		f := fmt.Sprintf("bmatch_%d.txt.sz", k)
		outname := path.Join(tmpdir, f)
		out, err := os.Create(outname)
		if err != nil {
			logger.Print(err)
			panic(err)
		}
		wtr := snappy.NewBufferedWriter(out)
		wtrs = append(wtrs, wtr)
		allwtrs = append(allwtrs, wtr, out)
	}

	bb := bytes.Repeat([]byte(" "), bufsize)
	bb[bufsize-1] = byte('\n')

	for r := range hitchan {

		wtr := wtrs[r.win]

		n1, err1 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.mseq)))
		n2, err2 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.left)))
		n3, err3 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.right)))
		n4, err4 := wtr.Write([]byte(fmt.Sprintf("%011d\t", r.tnum)))
		n5, err5 := wtr.Write([]byte(fmt.Sprintf("%d", r.pos)))

		for _, err := range []error{err1, err2, err3, err4, err5} {
			if err != nil {
				logger.Print(err)
				panic("writing error")
			}
		}

		n := n1 + n2 + n3 + n4 + n5
		if n > bufsize {
			panic("output line is too long")
		}

		// The rest of the line is spaces, then newline.
		_, err := wtr.Write(bb[n:bufsize])
		if err != nil {
			logger.Print(err)
			panic(err)
		}
	}

	for _, wtr := range allwtrs {
		wtr.Close()
	}
	wg.Done()
	logger.Printf("Exiting harvest")
}

// search loops through the target sequences, checking each window
// within each target gene for possible matches to the read
// collection.
func search() {

	fid, err := os.Open(config.GeneFileName)
	if err != nil {
		logger.Print(err)
		panic(err)
	}
	defer fid.Close()
	snr := snappy.NewReader(fid)

	// Target file contains some very long lines
	scanner := bufio.NewScanner(snr)
	sbuf := make([]byte, 1024*1024)
	scanner.Buffer(sbuf, 1024*1024)

	hitchan = make(chan rec)
	limit = make(chan bool, concurrency)

	var wg sync.WaitGroup
	wg.Add(1)
	go harvest(&wg)

	var i int
	for ; scanner.Scan(); i++ {

		if i%1000000 == 0 {
			logger.Printf("%d\n", i)
		}

		line := scanner.Text() // need a copy here

		toks := strings.Split(line, "\t")
		seq := toks[0] // The sequence

		limit <- true
		go processseq([]byte(seq), i)
	}

	if err := scanner.Err(); err != nil {
		msg := fmt.Sprintf("Problem reading %s on line %d\n", config.GeneFileName, i)
		os.Stderr.WriteString(msg)
		logger.Print(err)
		panic(err)
	}

	for k := 0; k < concurrency; k++ {
		limit <- true
	}

	close(hitchan)
	wg.Wait()
	logger.Printf("done with search")
}

func setupLogger() {
	logname := path.Join(tmpdir, "muscato_screen.log")
	logfid, err := os.Create(logname)
	if err != nil {
		panic(err)
	}
	logger = log.New(logfid, "", log.Ltime)
}

func estimateFullness() {

	n := 1000
	logger.Printf("Bloom filter fill rates:\n")

	for j, ba := range smp {
		c := 0
		for k := 0; k < n; k++ {
			i := uint64(rand.Int63()) % config.BloomSize
			f, err := ba.GetBit(i)
			if err != nil {
				panic(err)
			}
			if f {
				c++
			}
		}
		logger.Printf("%3d %.3f\n", j, float64(c)/float64(n))
	}
}

func main() {

	if len(os.Args) != 3 {
		panic("wrong number of arguments")
	}

	config = utils.ReadConfig(os.Args[1])

	if config.TempDir == "" {
		tmpdir = os.Args[2]
	} else {
		tmpdir = config.TempDir
	}

	bufsize = config.MaxReadLength + 50

	setupLogger()
	genTables()

	smp = make([]bitarray.BitArray, len(config.Windows))
	for k := range smp {
		smp[k] = bitarray.NewBitArray(config.BloomSize)
	}

	buildBloom()
	estimateFullness()
	search()
}
<commit_msg>update comments<commit_after>// Copyright 2017, Kerby Shedden and the Muscato contributors.

// muscato_screen is an initial screening step used by Muscato to
// identify candidate matches of a set of reads into a set of target
// gene sequences.  The results of the screen may contain false
// positives, but will not contain any false negatives.
//
// The approach is to use a Bloom filter to sketch the reads based on
// the subsequences that appear at defined offsets within the reads.
// For example, if position 10 is an offset and we are looking at
// subequences of width 15, then the read subsequences from position
// 10 through position 25 are entered into a Bloom filter.  Then, we
// scan through every target gene looking for matches to the Bloom
// filter.  When a match occurs, the match position (in the target)
// and flanking sequences are saved for subequent checking against the
// full read sequence.
//
// A simple entropy check is used to avoid considering subsequences
// that could match large numbers of reads or genes (and hence would
// be uninformative).  Currently, this check is based on the number of
// distinct dinucleotide subsequences in the window (e.g. in the
// 15-mer in the example above).
//
// The results are saved in files named bmatch*.txt.sz, where * is the
// window number.
//
// The format of the bmatch files is:
//
// (window sequence) (left tail) (right tail) (gene id) (position)

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"path"
	"strings"
	"sync"

	"github.com/chmduquesne/rollinghash"
	"github.com/chmduquesne/rollinghash/buzhash32"
	"github.com/golang-collections/go-datastructures/bitarray"
	"github.com/golang/snappy"
	"github.com/kshedden/seqmatch/utils"
)

const (
	// Number of goroutines, should probably scale with the number
	// of aailable cores.
	concurrency int = 100
)

var (
	// A log
	logger *log.Logger

	// Configuration information
	config *utils.Config

	// All working files are stored here
	tmpdir string

	// Bitarrays that back the Bloom filters
	smp []bitarray.BitArray

	// Tables to produce independent running hashes
	tables [][256]uint32

	// Communicate results back to driver
	hitchan chan rec

	// Semaphore for limiting goroutines
	limit chan bool

	// Line length for output
	bufsize int
)

// genTables generates base hash functions for a collection of rolling hashes.
func genTables() {
	tables = make([][256]uint32, config.NumHash)
	for j := 0; j < config.NumHash; j++ {
		mp := make(map[uint32]bool)
		for i := 0; i < 256; i++ {
			for {
				x := uint32(rand.Int63())
				if !mp[x] {
					tables[j][i] = x
					mp[x] = true
					break
				}
			}
		}
	}
}

// buildBloom constructs bloom filters for each window
func buildBloom() {

	logger.Printf("Building Bloom sketch of read collection...")

	hashes := make([]rollinghash.Hash32, config.NumHash)
	for j := range hashes {
		hashes[j] = buzhash32.NewFromUint32Array(tables[j])
	}

	fname := path.Join(tmpdir, "reads_sorted.txt.sz")
	fid, err := os.Open(fname)
	if err != nil {
		logger.Print(err)
		panic(err)
	}
	defer fid.Close()
	snr := snappy.NewReader(fid)
	scanner := bufio.NewScanner(snr)
	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)

	// Workspace for sequence diversity checker
	wk := make([]int, 25)

	var j int
	for ; scanner.Scan(); j++ {

		if j%1000000 == 0 {
			logger.Printf("%d\n", j)
		}

		line := scanner.Bytes()
		seq := bytes.Fields(line)[0]

		for k := 0; k < len(config.Windows); k++ {
			q1 := config.Windows[k]
			q2 := q1 + config.WindowWidth
			if q2 > len(seq) {
				continue
			}
			seqw := seq[q1:q2]

			// Check entropy
			if utils.CountDinuc(seqw, wk) < config.MinDinuc {
				continue
			}

			// Update the Bloom filter for this sequence
			for _, ha := range hashes {
				ha.Reset()
				_, err = ha.Write(seqw)
				if err != nil {
					panic(err)
				}
				x := uint64(ha.Sum32()) % config.BloomSize
				err := smp[k].SetBit(x)
				if err != nil {
					logger.Print(err)
					panic(err)
				}
			}
		}
	}

	if err := scanner.Err(); err != nil {
		msg := fmt.Sprintf("Problem reading reads_sorted.txt.sz on line %d\n", j)
		os.Stderr.WriteString(msg)
		logger.Print(err)
		panic(err)
	}

	logger.Printf("Done constructing Bloom filters")
}

type rec struct {
	mseq  string
	left  string
	right string
	win   int
	tnum  int
	pos   uint32
}

// checkWin returns the indices of the Bloom filters that match the
// current state of the hashes.  iw is workspace and hashses contains
// the hashes that define the Bloom filters.
func checkWin(ix []int, iw []uint64, hashes []rollinghash.Hash32) []int {

	// Get the hash states
	for j, ha := range hashes {
		iw[j] = uint64(ha.Sum32()) % config.BloomSize
	}

	ix = ix[0:0]

	// Loop over Bloom filters
	for k, ba := range smp {

		// Determine if the Bloom filter matches
		g := true
		for j := range hashes {
			f, err := ba.GetBit(iw[j])
			if err != nil {
				logger.Print(err)
				panic(err)
			}
			if !f {
				g = false
				break
			}
		}
		if g {
			ix = append(ix, k)
		}
	}

	return ix
}

// process one target sequence, runs concurrently with main loop.
func processseq(seq []byte, genenum int) {

	defer func() { <-limit }()

	hashes := make([]rollinghash.Hash32, config.NumHash)
	for j := range hashes {
		hashes[j] = buzhash32.NewFromUint32Array(tables[j])
	}

	// Initialize the hashes with the first window.
	hlen := config.WindowWidth
	if len(seq) < hlen {
		// Not long enough even for one window.
		return
	}
	for j := range hashes {
		_, err := hashes[j].Write(seq[0:hlen])
		if err != nil {
			panic(err)
		}
	}

	ix := make([]int, len(smp))
	iw := make([]uint64, config.NumHash)

	// Check if the initial window is a match
	ix = checkWin(ix, iw, hashes)
	for _, i := range ix {

		q1 := config.Windows[i]
		if q1 != 0 {
			// The only way the read can fit is if the
			// window starts at the beginning of the read.
			continue
		}
		q2 := q1 + config.WindowWidth

		jz := 100 - q2
		if jz > len(seq) {
			jz = len(seq)
		}
		hitchan <- rec{
			mseq:  string(seq[0:hlen]),
			left:  "",
			right: string(seq[hlen:jz]),
			tnum:  genenum,
			win:   i,
			pos:   0,
		}
	}

	// Check the rest of the windows
	for j := hlen; j < len(seq); j++ {

		for _, ha := range hashes {
			ha.Roll(seq[j])
		}
		ix = checkWin(ix, iw, hashes)

		// Process a match
		for _, i := range ix {

			q1 := config.Windows[i]
			q2 := q1 + config.WindowWidth
			if j < q2-1 {
				// The read would not fit
				continue
			}

			// Matching sequence is jx:jy
			jx := j - hlen + 1
			jy := j + 1

			// Left tail is jw:jx
			jw := jx - q1

			// Right tail is jy:jz
			jz := jy + config.MaxReadLength - q2
			if jz > len(seq) {
				// May not be long enough to fit, but
				// we don't know until we merge.
				jz = len(seq)
			}

			if jw >= 0 {
				hitchan <- rec{
					mseq:  string(seq[jx:jy]),
					left:  string(seq[jw:jx]),
					right: string(seq[jy:jz]),
					tnum:  genenum,
					win:   i,
					pos:   uint32(j - hlen + 1),
				}
			}
		}
	}
}

// Retrieve the results and write to disk
func harvest(wg *sync.WaitGroup) {

	var wtrs []io.Writer
	var allwtrs []io.Closer
	for k := 0; k < len(config.Windows); k++ {
		f := fmt.Sprintf("bmatch_%d.txt.sz", k)
		outname := path.Join(tmpdir, f)
		out, err := os.Create(outname)
		if err != nil {
			logger.Print(err)
			panic(err)
		}
		wtr := snappy.NewBufferedWriter(out)
		wtrs = append(wtrs, wtr)
		allwtrs = append(allwtrs, wtr, out)
	}

	bb := bytes.Repeat([]byte(" "), bufsize)
	bb[bufsize-1] = byte('\n')

	for r := range hitchan {

		wtr := wtrs[r.win]

		n1, err1 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.mseq)))
		n2, err2 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.left)))
		n3, err3 := wtr.Write([]byte(fmt.Sprintf("%s\t", r.right)))
		n4, err4 := wtr.Write([]byte(fmt.Sprintf("%011d\t", r.tnum)))
		n5, err5 := wtr.Write([]byte(fmt.Sprintf("%d", r.pos)))

		for _, err := range []error{err1, err2, err3, err4, err5} {
			if err != nil {
				logger.Print(err)
				panic("writing error")
			}
		}

		n := n1 + n2 + n3 + n4 + n5
		if n > bufsize {
			panic("output line is too long")
		}

		// The rest of the line is spaces, then newline.
		_, err := wtr.Write(bb[n:bufsize])
		if err != nil {
			logger.Print(err)
			panic(err)
		}
	}

	for _, wtr := range allwtrs {
		wtr.Close()
	}
	wg.Done()
	logger.Printf("Exiting harvest")
}

// search loops through the target sequences, checking each window
// within each target gene for possible matches to the read
// collection.
func search() {

	fid, err := os.Open(config.GeneFileName)
	if err != nil {
		logger.Print(err)
		panic(err)
	}
	defer fid.Close()
	snr := snappy.NewReader(fid)

	// Target file contains some very long lines
	scanner := bufio.NewScanner(snr)
	sbuf := make([]byte, 1024*1024)
	scanner.Buffer(sbuf, 1024*1024)

	hitchan = make(chan rec)
	limit = make(chan bool, concurrency)

	var wg sync.WaitGroup
	wg.Add(1)
	go harvest(&wg)

	var i int
	for ; scanner.Scan(); i++ {

		if i%1000000 == 0 {
			logger.Printf("%d\n", i)
		}

		line := scanner.Text() // need a copy here

		toks := strings.Split(line, "\t")
		seq := toks[0] // The sequence

		limit <- true
		go processseq([]byte(seq), i)
	}

	if err := scanner.Err(); err != nil {
		msg := fmt.Sprintf("Problem reading %s on line %d\n", config.GeneFileName, i)
		os.Stderr.WriteString(msg)
		logger.Print(err)
		panic(err)
	}

	for k := 0; k < concurrency; k++ {
		limit <- true
	}

	close(hitchan)
	wg.Wait()
	logger.Printf("done with search")
}

func setupLogger() {
	logname := path.Join(tmpdir, "muscato_screen.log")
	logfid, err := os.Create(logname)
	if err != nil {
		panic(err)
	}
	logger = log.New(logfid, "", log.Ltime)
}

func estimateFullness() {

	n := 1000
	logger.Printf("Bloom filter fill rates:\n")

	for j, ba := range smp {
		c := 0
		for k := 0; k < n; k++ {
			i := uint64(rand.Int63()) % config.BloomSize
			f, err := ba.GetBit(i)
			if err != nil {
				panic(err)
			}
			if f {
				c++
			}
		}
		logger.Printf("%3d %.3f\n", j, float64(c)/float64(n))
	}
}

func main() {

	if len(os.Args) != 3 {
		panic("wrong number of arguments")
	}

	config = utils.ReadConfig(os.Args[1])

	if config.TempDir == "" {
		tmpdir = os.Args[2]
	} else {
		tmpdir = config.TempDir
	}

	bufsize = config.MaxReadLength + 50

	setupLogger()
	genTables()

	smp = make([]bitarray.BitArray, len(config.Windows))
	for k := range smp {
		smp[k] = bitarray.NewBitArray(config.BloomSize)
	}

	buildBloom()
	estimateFullness()
	search()
}
<|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 network

import (
	"context"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/onsi/ginkgo"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	api "k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/test/e2e/framework"
	e2enetwork "k8s.io/kubernetes/test/e2e/framework/network"
	e2epod "k8s.io/kubernetes/test/e2e/framework/pod"
	e2eresource "k8s.io/kubernetes/test/e2e/framework/resource"
	e2eservice "k8s.io/kubernetes/test/e2e/framework/service"
	e2etestfiles "k8s.io/kubernetes/test/e2e/framework/testfiles"
	"k8s.io/kubernetes/test/e2e/network/common"
)

const (
	dnsReadyTimeout = time.Minute

	// RespondingTimeout is how long to wait for a service to be responding.
	RespondingTimeout = 2 * time.Minute
)

const queryDNSPythonTemplate string = `
import socket
try:
	socket.gethostbyname('%s')
	print('ok')
except:
	print('err')`

var _ = common.SIGDescribe("ClusterDns [Feature:Example]", func() {
	f := framework.NewDefaultFramework("cluster-dns")

	var c clientset.Interface
	ginkgo.BeforeEach(func() {
		c = f.ClientSet
	})

	read := func(file string) string {
		data, err := e2etestfiles.Read(file)
		if err != nil {
			framework.Fail(err.Error())
		}
		return string(data)
	}

	ginkgo.It("should create pod that uses dns", func() {
		// contrary to the example, this test does not use contexts, for simplicity
		// namespaces are passed directly.
		// Also, for simplicity, we don't use yamls with namespaces, but we
		// create testing namespaces instead.

		backendRcName := "dns-backend"
		backendSvcName := "dns-backend"
		backendPodName := "dns-backend"
		frontendPodName := "dns-frontend"
		frontendPodContainerName := "dns-frontend"
		clusterDnsPath := "test/e2e/testing-manifests/cluster-dns"
		podOutput := "Hello World!"

		// we need two namespaces anyway, so let's forget about
		// the one created in BeforeEach and create two new ones.
		namespaces := []*v1.Namespace{nil, nil}
		for i := range namespaces {
			var err error
			namespaceName := fmt.Sprintf("dnsexample%d", i)
			namespaces[i], err = f.CreateNamespace(namespaceName, nil)
			framework.ExpectNoError(err, "failed to create namespace: %s", namespaceName)
		}

		for _, ns := range namespaces {
			framework.RunKubectlOrDieInput(ns.Name, read(filepath.Join(clusterDnsPath, "dns-backend-rc.yaml")), "create", "-f", "-")
		}

		for _, ns := range namespaces {
			framework.RunKubectlOrDieInput(ns.Name, read(filepath.Join(clusterDnsPath, "dns-backend-service.yaml")), "create", "-f", "-")
		}

		// wait for objects
		for _, ns := range namespaces {
			e2eresource.WaitForControlledPodsRunning(c, ns.Name, backendRcName, api.Kind("ReplicationController"))
			e2enetwork.WaitForService(c, ns.Name, backendSvcName, true, framework.Poll, framework.ServiceStartTimeout)
		}
		// it is not enough that pods are running because they may be set to running, but
		// the application itself may have not been initialized. Just query the application.
		for _, ns := range namespaces {
			label := labels.SelectorFromSet(labels.Set(map[string]string{"name": backendRcName}))
			options := metav1.ListOptions{LabelSelector: label.String()}
			pods, err := c.CoreV1().Pods(ns.Name).List(context.TODO(), options)
			framework.ExpectNoError(err, "failed to list pods in namespace: %s", ns.Name)
			err = e2epod.PodsResponding(c, ns.Name, backendPodName, false, pods)
			framework.ExpectNoError(err, "waiting for all pods to respond")
			framework.Logf("found %d backend pods responding in namespace %s", len(pods.Items), ns.Name)

			err = waitForServiceResponding(c, ns.Name, backendSvcName)
			framework.ExpectNoError(err, "waiting for the service to respond")
		}

		// Now another tricky part:
		// It may happen that the service name is not yet in DNS.
		// So if we start our pod, it will fail. We must make sure
		// the name is already resolvable. So let's try to query DNS from
		// the pod we have, until we find our service name.
		// This complicated code may be removed if the pod itself retried after
		// dns error or timeout.
		// This code is probably unnecessary, but let's stay on the safe side.
		label := labels.SelectorFromSet(labels.Set(map[string]string{"name": backendPodName}))
		options := metav1.ListOptions{LabelSelector: label.String()}
		pods, err := c.CoreV1().Pods(namespaces[0].Name).List(context.TODO(), options)

		if err != nil || pods == nil || len(pods.Items) == 0 {
			framework.Failf("no running pods found")
		}
		podName := pods.Items[0].Name

		queryDNS := fmt.Sprintf(queryDNSPythonTemplate, backendSvcName+"."+namespaces[0].Name)
		_, err = framework.LookForStringInPodExec(namespaces[0].Name, podName, []string{"python", "-c", queryDNS}, "ok", dnsReadyTimeout)
		framework.ExpectNoError(err, "waiting for output from pod exec")

		updatedPodYaml := strings.Replace(read(filepath.Join(clusterDnsPath, "dns-frontend-pod.yaml")), fmt.Sprintf("dns-backend.development.svc.%s", framework.TestContext.ClusterDNSDomain), fmt.Sprintf("dns-backend.%s.svc.%s", namespaces[0].Name, framework.TestContext.ClusterDNSDomain), 1)

		// create a pod in each namespace
		for _, ns := range namespaces {
			framework.RunKubectlOrDieInput(ns.Name, updatedPodYaml, "create", "-f", "-")
		}

		// wait until the pods have been scheduler, i.e. are not Pending anymore. Remember
		// that we cannot wait for the pods to be running because our pods terminate by themselves.
		for _, ns := range namespaces {
			err := e2epod.WaitForPodNotPending(c, ns.Name, frontendPodName)
			framework.ExpectNoError(err)
		}

		// wait for pods to print their result
		for _, ns := range namespaces {
			_, err := framework.LookForStringInLog(ns.Name, frontendPodName, frontendPodContainerName, podOutput, framework.PodStartTimeout)
			framework.ExpectNoError(err, "pod %s failed to print result in logs", frontendPodName)
		}
	})
})

// waitForServiceResponding waits for the service to be responding.
func waitForServiceResponding(c clientset.Interface, ns, name string) error {
	ginkgo.By(fmt.Sprintf("trying to dial the service %s.%s via the proxy", ns, name))

	return wait.PollImmediate(framework.Poll, RespondingTimeout, func() (done bool, err error) {
		proxyRequest, errProxy := e2eservice.GetServicesProxyRequest(c, c.CoreV1().RESTClient().Get())
		if errProxy != nil {
			framework.Logf("Failed to get services proxy request: %v:", errProxy)
			return false, nil
		}

		ctx, cancel := context.WithTimeout(context.Background(), framework.SingleCallTimeout)
		defer cancel()

		body, err := proxyRequest.Namespace(ns).
			Name(name).
			Do(ctx).
			Raw()
		if err != nil {
			if ctx.Err() != nil {
				framework.Failf("Failed to GET from service %s: %v", name, err)
				return true, err
			}
			framework.Logf("Failed to GET from service %s: %v:", name, err)
			return false, nil
		}
		got := string(body)
		if len(got) == 0 {
			framework.Logf("Service %s: expected non-empty response", name)
			return false, err // stop polling
		}
		framework.Logf("Service %s: found nonempty answer: %s", name, got)
		return true, nil
	})
}
<commit_msg>Remove couple of variables to simplify the code<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 network

import (
	"context"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/onsi/ginkgo"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	api "k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/test/e2e/framework"
	e2enetwork "k8s.io/kubernetes/test/e2e/framework/network"
	e2epod "k8s.io/kubernetes/test/e2e/framework/pod"
	e2eresource "k8s.io/kubernetes/test/e2e/framework/resource"
	e2eservice "k8s.io/kubernetes/test/e2e/framework/service"
	e2etestfiles "k8s.io/kubernetes/test/e2e/framework/testfiles"
	"k8s.io/kubernetes/test/e2e/network/common"
)

const (
	dnsReadyTimeout = time.Minute

	// RespondingTimeout is how long to wait for a service to be responding.
	RespondingTimeout = 2 * time.Minute
)

const queryDNSPythonTemplate string = `
import socket
try:
	socket.gethostbyname('%s')
	print('ok')
except:
	print('err')`

var _ = common.SIGDescribe("ClusterDns [Feature:Example]", func() {
	f := framework.NewDefaultFramework("cluster-dns")

	var c clientset.Interface
	ginkgo.BeforeEach(func() {
		c = f.ClientSet
	})

	read := func(file string) string {
		data, err := e2etestfiles.Read(file)
		if err != nil {
			framework.Fail(err.Error())
		}
		return string(data)
	}

	ginkgo.It("should create pod that uses dns", func() {
		// contrary to the example, this test does not use contexts, for simplicity
		// namespaces are passed directly.
		// Also, for simplicity, we don't use yamls with namespaces, but we
		// create testing namespaces instead.

		backendName := "dns-backend"
		frontendName := "dns-frontend"
		clusterDnsPath := "test/e2e/testing-manifests/cluster-dns"
		podOutput := "Hello World!"

		// we need two namespaces anyway, so let's forget about
		// the one created in BeforeEach and create two new ones.
		namespaces := []*v1.Namespace{nil, nil}
		for i := range namespaces {
			var err error
			namespaceName := fmt.Sprintf("dnsexample%d", i)
			namespaces[i], err = f.CreateNamespace(namespaceName, nil)
			framework.ExpectNoError(err, "failed to create namespace: %s", namespaceName)
		}

		for _, ns := range namespaces {
			framework.RunKubectlOrDieInput(ns.Name, read(filepath.Join(clusterDnsPath, "dns-backend-rc.yaml")), "create", "-f", "-")
		}

		for _, ns := range namespaces {
			framework.RunKubectlOrDieInput(ns.Name, read(filepath.Join(clusterDnsPath, "dns-backend-service.yaml")), "create", "-f", "-")
		}

		// wait for objects
		for _, ns := range namespaces {
			e2eresource.WaitForControlledPodsRunning(c, ns.Name, backendName, api.Kind("ReplicationController"))
			e2enetwork.WaitForService(c, ns.Name, backendName, true, framework.Poll, framework.ServiceStartTimeout)
		}
		// it is not enough that pods are running because they may be set to running, but
		// the application itself may have not been initialized. Just query the application.
		for _, ns := range namespaces {
			label := labels.SelectorFromSet(labels.Set(map[string]string{"name": backendName}))
			options := metav1.ListOptions{LabelSelector: label.String()}
			pods, err := c.CoreV1().Pods(ns.Name).List(context.TODO(), options)
			framework.ExpectNoError(err, "failed to list pods in namespace: %s", ns.Name)
			err = e2epod.PodsResponding(c, ns.Name, backendName, false, pods)
			framework.ExpectNoError(err, "waiting for all pods to respond")
			framework.Logf("found %d backend pods responding in namespace %s", len(pods.Items), ns.Name)

			err = waitForServiceResponding(c, ns.Name, backendName)
			framework.ExpectNoError(err, "waiting for the service to respond")
		}

		// Now another tricky part:
		// It may happen that the service name is not yet in DNS.
		// So if we start our pod, it will fail. We must make sure
		// the name is already resolvable. So let's try to query DNS from
		// the pod we have, until we find our service name.
		// This complicated code may be removed if the pod itself retried after
		// dns error or timeout.
		// This code is probably unnecessary, but let's stay on the safe side.
		label := labels.SelectorFromSet(labels.Set(map[string]string{"name": backendName}))
		options := metav1.ListOptions{LabelSelector: label.String()}
		pods, err := c.CoreV1().Pods(namespaces[0].Name).List(context.TODO(), options)

		if err != nil || pods == nil || len(pods.Items) == 0 {
			framework.Failf("no running pods found")
		}
		podName := pods.Items[0].Name

		queryDNS := fmt.Sprintf(queryDNSPythonTemplate, backendName+"."+namespaces[0].Name)
		_, err = framework.LookForStringInPodExec(namespaces[0].Name, podName, []string{"python", "-c", queryDNS}, "ok", dnsReadyTimeout)
		framework.ExpectNoError(err, "waiting for output from pod exec")

		updatedPodYaml := strings.Replace(read(filepath.Join(clusterDnsPath, "dns-frontend-pod.yaml")), fmt.Sprintf("dns-backend.development.svc.%s", framework.TestContext.ClusterDNSDomain), fmt.Sprintf("dns-backend.%s.svc.%s", namespaces[0].Name, framework.TestContext.ClusterDNSDomain), 1)

		// create a pod in each namespace
		for _, ns := range namespaces {
			framework.RunKubectlOrDieInput(ns.Name, updatedPodYaml, "create", "-f", "-")
		}

		// wait until the pods have been scheduler, i.e. are not Pending anymore. Remember
		// that we cannot wait for the pods to be running because our pods terminate by themselves.
		for _, ns := range namespaces {
			err := e2epod.WaitForPodNotPending(c, ns.Name, frontendName)
			framework.ExpectNoError(err)
		}

		// wait for pods to print their result
		for _, ns := range namespaces {
			_, err := framework.LookForStringInLog(ns.Name, frontendName, frontendName, podOutput, framework.PodStartTimeout)
			framework.ExpectNoError(err, "pod %s failed to print result in logs", frontendName)
		}
	})
})

// waitForServiceResponding waits for the service to be responding.
func waitForServiceResponding(c clientset.Interface, ns, name string) error {
	ginkgo.By(fmt.Sprintf("trying to dial the service %s.%s via the proxy", ns, name))

	return wait.PollImmediate(framework.Poll, RespondingTimeout, func() (done bool, err error) {
		proxyRequest, errProxy := e2eservice.GetServicesProxyRequest(c, c.CoreV1().RESTClient().Get())
		if errProxy != nil {
			framework.Logf("Failed to get services proxy request: %v:", errProxy)
			return false, nil
		}

		ctx, cancel := context.WithTimeout(context.Background(), framework.SingleCallTimeout)
		defer cancel()

		body, err := proxyRequest.Namespace(ns).
			Name(name).
			Do(ctx).
			Raw()
		if err != nil {
			if ctx.Err() != nil {
				framework.Failf("Failed to GET from service %s: %v", name, err)
				return true, err
			}
			framework.Logf("Failed to GET from service %s: %v:", name, err)
			return false, nil
		}
		got := string(body)
		if len(got) == 0 {
			framework.Logf("Service %s: expected non-empty response", name)
			return false, err // stop polling
		}
		framework.Logf("Service %s: found nonempty answer: %s", name, got)
		return true, nil
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 e2enode

import (
	"context"
	"os/exec"
	"strconv"
	"time"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/component-base/metrics/testutil"
	kubeletmetrics "k8s.io/kubernetes/pkg/kubelet/metrics"
	"k8s.io/kubernetes/test/e2e/framework"
	e2egpu "k8s.io/kubernetes/test/e2e/framework/gpu"
	e2emanifest "k8s.io/kubernetes/test/e2e/framework/manifest"
	e2emetrics "k8s.io/kubernetes/test/e2e/framework/metrics"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
)

// numberOfNVIDIAGPUs returns the number of GPUs advertised by a node
// This is based on the Device Plugin system and expected to run on a COS based node
// After the NVIDIA drivers were installed
// TODO make this generic and not linked to COS only
func numberOfNVIDIAGPUs(node *v1.Node) int64 {
	val, ok := node.Status.Capacity[e2egpu.NVIDIAGPUResourceName]
	if !ok {
		return 0
	}
	return val.Value()
}

// NVIDIADevicePlugin returns the official Google Device Plugin pod for NVIDIA GPU in GKE
func NVIDIADevicePlugin() *v1.Pod {
	ds, err := e2emanifest.DaemonSetFromURL(e2egpu.GPUDevicePluginDSYAML)
	framework.ExpectNoError(err)
	p := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "device-plugin-nvidia-gpu-" + string(uuid.NewUUID()),
			Namespace: metav1.NamespaceSystem,
		},
		Spec: ds.Spec.Template.Spec,
	}
	// Remove node affinity
	p.Spec.Affinity = nil
	return p
}

// Serial because the test restarts Kubelet
var _ = SIGDescribe("NVIDIA GPU Device Plugin [Feature:GPUDevicePlugin][NodeFeature:GPUDevicePlugin][Serial] [Disruptive]", func() {
	f := framework.NewDefaultFramework("device-plugin-gpus-errors")

	ginkgo.Context("DevicePlugin", func() {
		var devicePluginPod *v1.Pod
		var err error
		ginkgo.BeforeEach(func() {
			ginkgo.By("Ensuring that Nvidia GPUs exists on the node")
			if !checkIfNvidiaGPUsExistOnNode() {
				ginkgo.Skip("Nvidia GPUs do not exist on the node. Skipping test.")
			}

			ginkgo.By("Creating the Google Device Plugin pod for NVIDIA GPU in GKE")
			devicePluginPod, err = f.ClientSet.CoreV1().Pods(metav1.NamespaceSystem).Create(context.TODO(), NVIDIADevicePlugin(), metav1.CreateOptions{})
			framework.ExpectNoError(err)

			ginkgo.By("Waiting for GPUs to become available on the local node")
			gomega.Eventually(func() bool {
				return numberOfNVIDIAGPUs(getLocalNode(f)) > 0
			}, 5*time.Minute, framework.Poll).Should(gomega.BeTrue())

			if numberOfNVIDIAGPUs(getLocalNode(f)) < 2 {
				ginkgo.Skip("Not enough GPUs to execute this test (at least two needed)")
			}
		})

		ginkgo.AfterEach(func() {
			l, err := f.PodClient().List(context.TODO(), metav1.ListOptions{})
			framework.ExpectNoError(err)

			for _, p := range l.Items {
				if p.Namespace != f.Namespace.Name {
					continue
				}

				f.PodClient().Delete(context.TODO(), p.Name, metav1.DeleteOptions{})
			}
		})

		ginkgo.It("checks that when Kubelet restarts exclusive GPU assignation to pods is kept.", func() {
			ginkgo.By("Creating one GPU pod on a node with at least two GPUs")
			podRECMD := "devs=$(ls /dev/ | egrep '^nvidia[0-9]+$') && echo gpu devices: $devs"
			p1 := f.PodClient().CreateSync(makeBusyboxPod(e2egpu.NVIDIAGPUResourceName, podRECMD))

			deviceIDRE := "gpu devices: (nvidia[0-9]+)"
			devID1 := parseLog(f, p1.Name, p1.Name, deviceIDRE)
			p1, err := f.PodClient().Get(context.TODO(), p1.Name, metav1.GetOptions{})
			framework.ExpectNoError(err)

			ginkgo.By("Restarting Kubelet and waiting for the current running pod to restart")
			restartKubelet()

			ginkgo.By("Confirming that after a kubelet and pod restart, GPU assignment is kept")
			ensurePodContainerRestart(f, p1.Name, p1.Name)
			devIDRestart1 := parseLog(f, p1.Name, p1.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart1, devID1)

			ginkgo.By("Restarting Kubelet and creating another pod")
			restartKubelet()
			framework.WaitForAllNodesSchedulable(f.ClientSet, framework.TestContext.NodeSchedulableTimeout)
			gomega.Eventually(func() bool {
				return numberOfNVIDIAGPUs(getLocalNode(f)) > 0
			}, 5*time.Minute, framework.Poll).Should(gomega.BeTrue())
			p2 := f.PodClient().CreateSync(makeBusyboxPod(e2egpu.NVIDIAGPUResourceName, podRECMD))

			ginkgo.By("Checking that pods got a different GPU")
			devID2 := parseLog(f, p2.Name, p2.Name, deviceIDRE)

			framework.ExpectEqual(devID1, devID2)

			ginkgo.By("Deleting device plugin.")
			f.ClientSet.CoreV1().Pods(metav1.NamespaceSystem).Delete(context.TODO(), devicePluginPod.Name, metav1.DeleteOptions{})
			ginkgo.By("Waiting for GPUs to become unavailable on the local node")
			gomega.Eventually(func() bool {
				node, err := f.ClientSet.CoreV1().Nodes().Get(context.TODO(), framework.TestContext.NodeName, metav1.GetOptions{})
				framework.ExpectNoError(err)
				return numberOfNVIDIAGPUs(node) <= 0
			}, 10*time.Minute, framework.Poll).Should(gomega.BeTrue())
			ginkgo.By("Checking that scheduled pods can continue to run even after we delete device plugin.")
			ensurePodContainerRestart(f, p1.Name, p1.Name)
			devIDRestart1 = parseLog(f, p1.Name, p1.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart1, devID1)

			ensurePodContainerRestart(f, p2.Name, p2.Name)
			devIDRestart2 := parseLog(f, p2.Name, p2.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart2, devID2)
			ginkgo.By("Restarting Kubelet.")
			restartKubelet()
			ginkgo.By("Checking that scheduled pods can continue to run even after we delete device plugin and restart Kubelet.")
			ensurePodContainerRestart(f, p1.Name, p1.Name)
			devIDRestart1 = parseLog(f, p1.Name, p1.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart1, devID1)
			ensurePodContainerRestart(f, p2.Name, p2.Name)
			devIDRestart2 = parseLog(f, p2.Name, p2.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart2, devID2)
			logDevicePluginMetrics()

			// Cleanup
			f.PodClient().DeleteSync(p1.Name, metav1.DeleteOptions{}, framework.DefaultPodDeletionTimeout)
			f.PodClient().DeleteSync(p2.Name, metav1.DeleteOptions{}, framework.DefaultPodDeletionTimeout)
		})
	})
})

func checkIfNvidiaGPUsExistOnNode() bool {
	// Cannot use `lspci` because it is not installed on all distros by default.
	err := exec.Command("/bin/sh", "-c", "find /sys/devices/pci* -type f | grep vendor | xargs cat | grep 0x10de").Run()
	if err != nil {
		framework.Logf("check for nvidia GPUs failed. Got Error: %v", err)
		return false
	}
	return true
}

func logDevicePluginMetrics() {
	ms, err := e2emetrics.GrabKubeletMetricsWithoutProxy(framework.TestContext.NodeName+":10255", "/metrics")
	framework.ExpectNoError(err)
	for msKey, samples := range ms {
		switch msKey {
		case kubeletmetrics.KubeletSubsystem + "_" + kubeletmetrics.DevicePluginAllocationDurationKey:
			for _, sample := range samples {
				latency := sample.Value
				resource := string(sample.Metric["resource_name"])
				var quantile float64
				if val, ok := sample.Metric[testutil.QuantileLabel]; ok {
					var err error
					if quantile, err = strconv.ParseFloat(string(val), 64); err != nil {
						continue
					}
					framework.Logf("Metric: %v ResourceName: %v Quantile: %v Latency: %v", msKey, resource, quantile, latency)
				}
			}
		case kubeletmetrics.KubeletSubsystem + "_" + kubeletmetrics.DevicePluginRegistrationCountKey:
			for _, sample := range samples {
				resource := string(sample.Metric["resource_name"])
				count := sample.Value
				framework.Logf("Metric: %v ResourceName: %v Count: %v", msKey, resource, count)
			}
		}
	}
}
<commit_msg>Skip NVidia GPU test in node e2e CI jobs for containerd and other runtimes<commit_after>/*
Copyright 2017 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 e2enode

import (
	"context"
	"os/exec"
	"strconv"
	"time"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/component-base/metrics/testutil"
	kubeletmetrics "k8s.io/kubernetes/pkg/kubelet/metrics"
	"k8s.io/kubernetes/test/e2e/framework"
	e2egpu "k8s.io/kubernetes/test/e2e/framework/gpu"
	e2emanifest "k8s.io/kubernetes/test/e2e/framework/manifest"
	e2emetrics "k8s.io/kubernetes/test/e2e/framework/metrics"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
)

// numberOfNVIDIAGPUs returns the number of GPUs advertised by a node
// This is based on the Device Plugin system and expected to run on a COS based node
// After the NVIDIA drivers were installed
// TODO make this generic and not linked to COS only
func numberOfNVIDIAGPUs(node *v1.Node) int64 {
	val, ok := node.Status.Capacity[e2egpu.NVIDIAGPUResourceName]
	if !ok {
		return 0
	}
	return val.Value()
}

// NVIDIADevicePlugin returns the official Google Device Plugin pod for NVIDIA GPU in GKE
func NVIDIADevicePlugin() *v1.Pod {
	ds, err := e2emanifest.DaemonSetFromURL(e2egpu.GPUDevicePluginDSYAML)
	framework.ExpectNoError(err)
	p := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "device-plugin-nvidia-gpu-" + string(uuid.NewUUID()),
			Namespace: metav1.NamespaceSystem,
		},
		Spec: ds.Spec.Template.Spec,
	}
	// Remove node affinity
	p.Spec.Affinity = nil
	return p
}

// Serial because the test restarts Kubelet
var _ = SIGDescribe("NVIDIA GPU Device Plugin [Feature:GPUDevicePlugin][NodeFeature:GPUDevicePlugin][Serial] [Disruptive]", func() {
	f := framework.NewDefaultFramework("device-plugin-gpus-errors")

	ginkgo.Context("DevicePlugin", func() {
		var devicePluginPod *v1.Pod
		var err error
		ginkgo.BeforeEach(func() {
			ginkgo.By("Ensuring that Nvidia GPUs exists on the node")
			if !checkIfNvidiaGPUsExistOnNode() {
				ginkgo.Skip("Nvidia GPUs do not exist on the node. Skipping test.")
			}

			if framework.TestContext.ContainerRuntime != "docker" {
				ginkgo.Skip("Test works only with in-tree dockershim. Skipping test.")
			}

			ginkgo.By("Creating the Google Device Plugin pod for NVIDIA GPU in GKE")
			devicePluginPod, err = f.ClientSet.CoreV1().Pods(metav1.NamespaceSystem).Create(context.TODO(), NVIDIADevicePlugin(), metav1.CreateOptions{})
			framework.ExpectNoError(err)

			ginkgo.By("Waiting for GPUs to become available on the local node")
			gomega.Eventually(func() bool {
				return numberOfNVIDIAGPUs(getLocalNode(f)) > 0
			}, 5*time.Minute, framework.Poll).Should(gomega.BeTrue())

			if numberOfNVIDIAGPUs(getLocalNode(f)) < 2 {
				ginkgo.Skip("Not enough GPUs to execute this test (at least two needed)")
			}
		})

		ginkgo.AfterEach(func() {
			l, err := f.PodClient().List(context.TODO(), metav1.ListOptions{})
			framework.ExpectNoError(err)

			for _, p := range l.Items {
				if p.Namespace != f.Namespace.Name {
					continue
				}

				f.PodClient().Delete(context.TODO(), p.Name, metav1.DeleteOptions{})
			}
		})

		ginkgo.It("checks that when Kubelet restarts exclusive GPU assignation to pods is kept.", func() {
			ginkgo.By("Creating one GPU pod on a node with at least two GPUs")
			podRECMD := "devs=$(ls /dev/ | egrep '^nvidia[0-9]+$') && echo gpu devices: $devs"
			p1 := f.PodClient().CreateSync(makeBusyboxPod(e2egpu.NVIDIAGPUResourceName, podRECMD))

			deviceIDRE := "gpu devices: (nvidia[0-9]+)"
			devID1 := parseLog(f, p1.Name, p1.Name, deviceIDRE)
			p1, err := f.PodClient().Get(context.TODO(), p1.Name, metav1.GetOptions{})
			framework.ExpectNoError(err)

			ginkgo.By("Restarting Kubelet and waiting for the current running pod to restart")
			restartKubelet()

			ginkgo.By("Confirming that after a kubelet and pod restart, GPU assignment is kept")
			ensurePodContainerRestart(f, p1.Name, p1.Name)
			devIDRestart1 := parseLog(f, p1.Name, p1.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart1, devID1)

			ginkgo.By("Restarting Kubelet and creating another pod")
			restartKubelet()
			framework.WaitForAllNodesSchedulable(f.ClientSet, framework.TestContext.NodeSchedulableTimeout)
			gomega.Eventually(func() bool {
				return numberOfNVIDIAGPUs(getLocalNode(f)) > 0
			}, 5*time.Minute, framework.Poll).Should(gomega.BeTrue())
			p2 := f.PodClient().CreateSync(makeBusyboxPod(e2egpu.NVIDIAGPUResourceName, podRECMD))

			ginkgo.By("Checking that pods got a different GPU")
			devID2 := parseLog(f, p2.Name, p2.Name, deviceIDRE)

			framework.ExpectEqual(devID1, devID2)

			ginkgo.By("Deleting device plugin.")
			f.ClientSet.CoreV1().Pods(metav1.NamespaceSystem).Delete(context.TODO(), devicePluginPod.Name, metav1.DeleteOptions{})
			ginkgo.By("Waiting for GPUs to become unavailable on the local node")
			gomega.Eventually(func() bool {
				node, err := f.ClientSet.CoreV1().Nodes().Get(context.TODO(), framework.TestContext.NodeName, metav1.GetOptions{})
				framework.ExpectNoError(err)
				return numberOfNVIDIAGPUs(node) <= 0
			}, 10*time.Minute, framework.Poll).Should(gomega.BeTrue())
			ginkgo.By("Checking that scheduled pods can continue to run even after we delete device plugin.")
			ensurePodContainerRestart(f, p1.Name, p1.Name)
			devIDRestart1 = parseLog(f, p1.Name, p1.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart1, devID1)

			ensurePodContainerRestart(f, p2.Name, p2.Name)
			devIDRestart2 := parseLog(f, p2.Name, p2.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart2, devID2)
			ginkgo.By("Restarting Kubelet.")
			restartKubelet()
			ginkgo.By("Checking that scheduled pods can continue to run even after we delete device plugin and restart Kubelet.")
			ensurePodContainerRestart(f, p1.Name, p1.Name)
			devIDRestart1 = parseLog(f, p1.Name, p1.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart1, devID1)
			ensurePodContainerRestart(f, p2.Name, p2.Name)
			devIDRestart2 = parseLog(f, p2.Name, p2.Name, deviceIDRE)
			framework.ExpectEqual(devIDRestart2, devID2)
			logDevicePluginMetrics()

			// Cleanup
			f.PodClient().DeleteSync(p1.Name, metav1.DeleteOptions{}, framework.DefaultPodDeletionTimeout)
			f.PodClient().DeleteSync(p2.Name, metav1.DeleteOptions{}, framework.DefaultPodDeletionTimeout)
		})
	})
})

func checkIfNvidiaGPUsExistOnNode() bool {
	// Cannot use `lspci` because it is not installed on all distros by default.
	err := exec.Command("/bin/sh", "-c", "find /sys/devices/pci* -type f | grep vendor | xargs cat | grep 0x10de").Run()
	if err != nil {
		framework.Logf("check for nvidia GPUs failed. Got Error: %v", err)
		return false
	}
	return true
}

func logDevicePluginMetrics() {
	ms, err := e2emetrics.GrabKubeletMetricsWithoutProxy(framework.TestContext.NodeName+":10255", "/metrics")
	framework.ExpectNoError(err)
	for msKey, samples := range ms {
		switch msKey {
		case kubeletmetrics.KubeletSubsystem + "_" + kubeletmetrics.DevicePluginAllocationDurationKey:
			for _, sample := range samples {
				latency := sample.Value
				resource := string(sample.Metric["resource_name"])
				var quantile float64
				if val, ok := sample.Metric[testutil.QuantileLabel]; ok {
					var err error
					if quantile, err = strconv.ParseFloat(string(val), 64); err != nil {
						continue
					}
					framework.Logf("Metric: %v ResourceName: %v Quantile: %v Latency: %v", msKey, resource, quantile, latency)
				}
			}
		case kubeletmetrics.KubeletSubsystem + "_" + kubeletmetrics.DevicePluginRegistrationCountKey:
			for _, sample := range samples {
				resource := string(sample.Metric["resource_name"])
				count := sample.Value
				framework.Logf("Metric: %v ResourceName: %v Count: %v", msKey, resource, count)
			}
		}
	}
}
<|endoftext|>


<commit_before>package sweetiebot

import (
  "github.com/bwmarrin/discordgo"
  "strings"
  "strconv"
  "regexp"
)

// This module picks a random action to do whenever #manechat has been idle for several minutes (configurable)
type SpoilerModule struct {
  spoilerban *regexp.Regexp
  lastmsg int64 // Sanity rate limiter
}

func (w *SpoilerModule) Name() string {
  return "Spoiler"
}

func (w *SpoilerModule) Register(hooks *ModuleHooks) {
  w.lastmsg = 0
  w.UpdateRegex()
  hooks.OnMessageCreate = append(hooks.OnMessageCreate, w)
  hooks.OnMessageUpdate = append(hooks.OnMessageUpdate, w)
  hooks.OnCommand = append(hooks.OnCommand, w)
}

func (w *SpoilerModule) HasSpoiler(s *discordgo.Session, m *discordgo.Message) bool {
  cid := SBatoi(m.ChannelID)
  for _, v := range sb.config.SpoilChannels {
    if cid == v {
      return false // this is a spoiler channel so we don't monitor it
    }
  }
  if w.spoilerban != nil && w.spoilerban.MatchString(strings.ToLower(m.Content)) {
    s.ChannelMessageDelete(m.ChannelID, m.ID)
    if RateLimit(&w.lastmsg, sb.config.Maxspoiltime) {
      sb.SendMessage(m.ChannelID, "[](/sbtarget) ```POSTING SPOILERS IS A BANNABLE OFFENSE. All discussion about future episodes or seasons MUST be in #mylittlespoilers.```")
    }
    return true
  }
  return false
}

func (w *SpoilerModule) OnMessageCreate(s *discordgo.Session, m *discordgo.Message) {
  w.HasSpoiler(s, m)
}

func (w *SpoilerModule) OnMessageUpdate(s *discordgo.Session, m *discordgo.Message) {
  w.HasSpoiler(s, m)
}

func (w *SpoilerModule) OnCommand(s *discordgo.Session, m *discordgo.Message) bool {
  if UserHasRole(m.Author.ID, strconv.FormatUint(sb.config.AlertRole, 10)) { return false } // If we are a princess, always allow us to run this command, otherwise we can't unspoil things
  return w.HasSpoiler(s, m)
}

func (w *SpoilerModule) UpdateRegex() bool {
  if len(sb.config.Collections["spoiler"]) < 1 {
    w.spoilerban = nil
    return true
  }
  var err error
  w.spoilerban, err = regexp.Compile("(" + strings.Join(MapToSlice(sb.config.Collections["spoiler"]), "|") + ")")
  return err == nil
}<commit_msg>Change spoiler string for new emote<commit_after>package sweetiebot

import (
  "github.com/bwmarrin/discordgo"
  "strings"
  "strconv"
  "regexp"
)

// This module picks a random action to do whenever #manechat has been idle for several minutes (configurable)
type SpoilerModule struct {
  spoilerban *regexp.Regexp
  lastmsg int64 // Sanity rate limiter
}

func (w *SpoilerModule) Name() string {
  return "Spoiler"
}

func (w *SpoilerModule) Register(hooks *ModuleHooks) {
  w.lastmsg = 0
  w.UpdateRegex()
  hooks.OnMessageCreate = append(hooks.OnMessageCreate, w)
  hooks.OnMessageUpdate = append(hooks.OnMessageUpdate, w)
  hooks.OnCommand = append(hooks.OnCommand, w)
}

func (w *SpoilerModule) HasSpoiler(s *discordgo.Session, m *discordgo.Message) bool {
  cid := SBatoi(m.ChannelID)
  for _, v := range sb.config.SpoilChannels {
    if cid == v {
      return false // this is a spoiler channel so we don't monitor it
    }
  }
  if w.spoilerban != nil && w.spoilerban.MatchString(strings.ToLower(m.Content)) {
    s.ChannelMessageDelete(m.ChannelID, m.ID)
    if RateLimit(&w.lastmsg, sb.config.Maxspoiltime) {
      sb.SendMessage(m.ChannelID, "[](/nospoilers) ```NO SPOILERS! Posting spoilers is a bannable offense. All discussion about new and future content MUST be in #mylittlespoilers.```")
    }
    return true
  }
  return false
}

func (w *SpoilerModule) OnMessageCreate(s *discordgo.Session, m *discordgo.Message) {
  w.HasSpoiler(s, m)
}

func (w *SpoilerModule) OnMessageUpdate(s *discordgo.Session, m *discordgo.Message) {
  w.HasSpoiler(s, m)
}

func (w *SpoilerModule) OnCommand(s *discordgo.Session, m *discordgo.Message) bool {
  if UserHasRole(m.Author.ID, strconv.FormatUint(sb.config.AlertRole, 10)) { return false } // If we are a princess, always allow us to run this command, otherwise we can't unspoil things
  return w.HasSpoiler(s, m)
}

func (w *SpoilerModule) UpdateRegex() bool {
  if len(sb.config.Collections["spoiler"]) < 1 {
    w.spoilerban = nil
    return true
  }
  var err error
  w.spoilerban, err = regexp.Compile("(" + strings.Join(MapToSlice(sb.config.Collections["spoiler"]), "|") + ")")
  return err == nil
}<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative 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 main

import (
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"strconv"
	"sync"
	"time"

	"knative.dev/serving/test"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

// Algorithm from https://stackoverflow.com/a/21854246

// Only primes less than or equal to N will be generated
func primes(N int) []int {
	var x, y, n int
	nsqrt := math.Sqrt(float64(N))

	isPrime := make([]bool, N)

	for x = 1; float64(x) <= nsqrt; x++ {
		for y = 1; float64(y) <= nsqrt; y++ {
			n = 4*(x*x) + y*y
			if n <= N && (n%12 == 1 || n%12 == 5) {
				isPrime[n] = !isPrime[n]
			}
			n = 3*(x*x) + y*y
			if n <= N && n%12 == 7 {
				isPrime[n] = !isPrime[n]
			}
			n = 3*(x*x) - y*y
			if x > y && n <= N && n%12 == 11 {
				isPrime[n] = !isPrime[n]
			}
		}
	}

	for n = 5; float64(n) <= nsqrt; n++ {
		if isPrime[n] {
			for y = n * n; y < N; y += n * n {
				isPrime[y] = false
			}
		}
	}

	isPrime[2] = true
	isPrime[3] = true

	primes := make([]int, 0, 1270606)
	for x = 0; x < len(isPrime)-1; x++ {
		if isPrime[x] {
			primes = append(primes, x)
		}
	}

	// primes is now a slice that contains all primes numbers up to N
	return primes
}

func bloat(mb int) string {
	b := make([]byte, mb*1024*1024)
	for i := 0; i < len(b); i++ {
		b[i] = 1
	}
	return fmt.Sprintf("Allocated %v Mb of memory.\n", mb)
}

func prime(max int) string {
	p := primes(max)
	if len(p) == 0 {
		return fmt.Sprintf("There are no primes smaller than %d.\n", max)
	}
	return fmt.Sprintf("The largest prime less than %d is %d.\n", max, p[len(p)-1])
}

func sleep(d time.Duration) string {
	start := time.Now()
	time.Sleep(d)
	return fmt.Sprintf("Slept for %v.\n", time.Since(start))
}

func randSleep(randSleepTimeMean time.Duration, randSleepTimeStdDev int) string {
	start := time.Now()
	randRes := time.Duration(rand.NormFloat64()*float64(randSleepTimeStdDev))*time.Millisecond + randSleepTimeMean
	time.Sleep(randRes)
	return fmt.Sprintf("Randomly slept for %v.\n", time.Since(start))
}

func parseDurationParam(r *http.Request, param string) (time.Duration, bool, error) {
	value := r.URL.Query().Get(param)
	if value == "" {
		return 0, false, nil
	}
	d, err := time.ParseDuration(value)
	if err != nil {
		return 0, false, err
	}
	return d, true, nil
}

func parseIntParam(r *http.Request, param string) (int, bool, error) {
	value := r.URL.Query().Get(param)
	if value == "" {
		return 0, false, nil
	}
	i, err := strconv.Atoi(value)
	if err != nil {
		return 0, false, err
	}
	return i, true, nil
}

func handler(w http.ResponseWriter, r *http.Request) {
	// Validate inputs.
	var ms time.Duration
	msv, hasMs, err := parseIntParam(r, "sleep")
	if err != nil {
		// If it is a numeric error and it's parsing error, then
		// try to parse it as a duration
		if nerr, ok := err.(*strconv.NumError); ok && nerr.Err == strconv.ErrSyntax {
			ms, hasMs, err = parseDurationParam(r, "sleep")
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
		} else {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
	} else {
		ms = time.Duration(msv) * time.Millisecond
	}
	if ms < 0 {
		http.Error(w, "Negative query params are not supported", http.StatusBadRequest)
		return
	}
	mssd, hasMssd, err := parseIntParam(r, "sleep-stddev")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if mssd < 0 {
		http.Error(w, "Negative query params are not supported", http.StatusBadRequest)
		return
	}
	max, hasMax, err := parseIntParam(r, "prime")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if max < 0 {
		http.Error(w, "Negative query params are not supported", http.StatusBadRequest)
		return
	}
	mb, hasMb, err := parseIntParam(r, "bloat")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if mb < 0 {
		http.Error(w, "Negative durations are not supported", http.StatusBadRequest)
		return
	}
	// Consume time, cpu and memory in parallel.
	var wg sync.WaitGroup
	defer wg.Wait()
	if hasMs && !hasMssd && ms > 0 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Fprint(w, sleep(ms))
		}()
	}
	if hasMs && hasMssd && ms > 0 && mssd > 0 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Fprint(w, randSleep(ms, mssd))
		}()
	}
	if hasMax && max > 0 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Fprint(w, prime(max))
		}()
	}
	if hasMb && mb > 0 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Fprint(w, bloat(mb))
		}()
	}
}

func main() {
	test.ListenAndServeGracefully(":8080", handler)
}
<commit_msg>Microfixes to autoscale.go (#7594)<commit_after>/*
Copyright 2018 The Knative 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 main

import (
	"fmt"
	"math"
	"math/rand"
	"net/http"
	"strconv"
	"sync"
	"time"

	"knative.dev/serving/test"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

// Algorithm from https://stackoverflow.com/a/21854246

// Only primes less than or equal to N will be generated
func primes(N int) []int {
	var x, y, n int
	nsqrt := math.Sqrt(float64(N))

	isPrime := make([]bool, N)

	for x = 1; float64(x) <= nsqrt; x++ {
		for y = 1; float64(y) <= nsqrt; y++ {
			n = 4*(x*x) + y*y
			if n <= N && (n%12 == 1 || n%12 == 5) {
				isPrime[n] = !isPrime[n]
			}
			n = 3*(x*x) + y*y
			if n <= N && n%12 == 7 {
				isPrime[n] = !isPrime[n]
			}
			n = 3*(x*x) - y*y
			if x > y && n <= N && n%12 == 11 {
				isPrime[n] = !isPrime[n]
			}
		}
	}

	for n = 5; float64(n) <= nsqrt; n++ {
		if isPrime[n] {
			for y = n * n; y < N; y += n * n {
				isPrime[y] = false
			}
		}
	}

	isPrime[2] = true
	isPrime[3] = true

	primes := make([]int, 0, 1270606)
	for x = 0; x < len(isPrime)-1; x++ {
		if isPrime[x] {
			primes = append(primes, x)
		}
	}

	// primes is now a slice that contains all primes numbers up to N
	return primes
}

func bloat(mb int) string {
	b := make([]byte, mb*1024*1024)
	for i := 0; i < len(b); i++ {
		b[i] = 1
	}
	return fmt.Sprintf("Allocated %v Mb of memory.\n", mb)
}

func prime(max int) string {
	p := primes(max)
	if len(p) == 0 {
		return fmt.Sprintf("There are no primes smaller than %d.\n", max)
	}
	return fmt.Sprintf("The largest prime less than %d is %d.\n", max, p[len(p)-1])
}

func sleep(d time.Duration) string {
	start := time.Now()
	time.Sleep(d)
	return fmt.Sprintf("Slept for %v.\n", time.Since(start))
}

func randSleep(randSleepTimeMean time.Duration, randSleepTimeStdDev int) string {
	start := time.Now()
	randRes := time.Duration(rand.NormFloat64()*float64(randSleepTimeStdDev))*time.Millisecond + randSleepTimeMean
	time.Sleep(randRes)
	return fmt.Sprintf("Randomly slept for %v.\n", time.Since(start))
}

func parseDurationParam(r *http.Request, param string) (time.Duration, bool, error) {
	value := r.URL.Query().Get(param)
	if value == "" {
		return 0, false, nil
	}
	d, err := time.ParseDuration(value)
	if err != nil {
		return 0, false, err
	}
	return d, true, nil
}

func parseIntParam(r *http.Request, param string) (int, bool, error) {
	value := r.URL.Query().Get(param)
	if value == "" {
		return 0, false, nil
	}
	i, err := strconv.Atoi(value)
	if err != nil {
		return 0, false, err
	}
	return i, true, nil
}

func handler(w http.ResponseWriter, r *http.Request) {
	// Validate inputs.
	var ms time.Duration
	msv, hasMs, err := parseIntParam(r, "sleep")
	if err != nil {
		// If it is a numeric error and it's parsing error, then
		// try to parse it as a duration
		if nerr, ok := err.(*strconv.NumError); ok && nerr.Err == strconv.ErrSyntax {
			ms, hasMs, err = parseDurationParam(r, "sleep")
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
		} else {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
	} else {
		ms = time.Duration(msv) * time.Millisecond
	}
	if ms < 0 {
		http.Error(w, "Negative query params are not supported", http.StatusBadRequest)
		return
	}
	msSD, hasMsSD, err := parseIntParam(r, "sleep-stddev")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if msSD < 0 {
		http.Error(w, "Negative query params are not supported", http.StatusBadRequest)
		return
	}
	max, hasMax, err := parseIntParam(r, "prime")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if hasMax && max <= 0 {
		http.Error(w, "Non-positive query params are not supported", http.StatusBadRequest)
		return
	}
	mb, hasMb, err := parseIntParam(r, "bloat")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if hasMb && mb <= 0 {
		http.Error(w, "Non-positive durations are not supported", http.StatusBadRequest)
		return
	}
	// Consume time, cpu and memory in parallel.
	var wg sync.WaitGroup
	defer wg.Wait()
	if hasMs && !hasMsSD {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Fprint(w, sleep(ms))
		}()
	}
	if hasMs && hasMsSD {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Fprint(w, randSleep(ms, msSD))
		}()
	}
	if hasMax {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Fprint(w, prime(max))
		}()
	}
	if hasMb {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Fprint(w, bloat(mb))
		}()
	}
}

func main() {
	test.ListenAndServeGracefully(":8080", handler)
}
<|endoftext|>


<commit_before>// Package main provides ...
package main

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"flag"
	"os"
	"strconv"
	"sync"

	log "github.com/cihub/seelog"
	"github.com/jonaz/goenocean"
	"github.com/stampzilla/stampzilla-go/nodes/basenode"
	"github.com/stampzilla/stampzilla-go/protocol"
	"github.com/stampzilla/stampzilla-go/protocol/devices"
)

var VERSION string = "dev"
var BUILD_DATE string = ""

var state *State

func main() {

	node := protocol.NewNode("enocean")
	node.Version = VERSION
	node.BuildDate = BUILD_DATE

	flag.Parse()

	//Setup Config
	config := basenode.NewConfig()
	basenode.SetConfig(config)

	//Start communication with the server
	//serverSendChannel = make(chan interface{})
	//serverRecvChannel = make(chan protocol.Command)
	connection := basenode.Connect()
	go monitorState(node, connection)
	go serverRecv(connection)

	// Describe available actions
	node.AddAction("set", "Set", []string{"Devices.Id"})
	node.AddAction("toggle", "Toggle", []string{"Devices.Id"})
	node.AddAction("dim", "Dim", []string{"Devices.Id", "value"})

	// Describe available layouts
	node.AddLayout("1", "switch", "toggle", "Devices", []string{"on"}, "Switches")
	node.AddLayout("2", "slider", "dim", "Devices", []string{"dim"}, "Dimmers")
	node.AddLayout("3", "slider", "dim", "Devices", []string{"dim"}, "Specials")

	node.AddElement(&protocol.Element{
		Type: protocol.ElementTypeToggle,
		Name: "Lamp 0186ff7d",
		Command: &protocol.Command{
			Cmd:  "toggle",
			Args: []string{"0186ff7d"},
		},
		Feedback: `Devices["0186ff7d"].On`,
	})
	node.AddElement(&protocol.Element{
		Type: protocol.ElementTypeText,
		Name: "Lamp 0186ff7d power",
		//Command: &protocol.Command{
		//Cmd:  "toggle",
		//Args: []string{"0186ff7d"},
		//},
		Feedback: `Devices["0186ff7d"].PowerW`,
	})

	//Setup state
	state = NewState()
	state.Devices = readConfigFromFile()
	node.SetState(state)

	//TODO remove element generator and AddElement and AddLayout and AddAction. Devices will superseed that.
	elementGenerator := &ElementGenerator{}
	elementGenerator.State = state
	elementGenerator.Node = node
	elementGenerator.Run()

	for _, dev := range state.Devices {
		// TODO if RecvEEPs is f60201 then its a button and not lamp
		node.Devices().Add(&devices.Device{
			Type:   "lamp",
			Name:   dev.Name,
			Id:     dev.Id,
			Online: true,
			Node:   config.Uuid,
			StateMap: map[string]string{
				"On": "Devices[" + dev.Id + "]" + ".On",
			},
		})
	}

	checkDuplicateSenderIds()

	setupEnoceanCommunication(node, connection)
}

func monitorState(node *protocol.Node, connection basenode.Connection) {
	for s := range connection.State() {
		switch s {
		case basenode.ConnectionStateConnected:
			connection.Send(node)
		case basenode.ConnectionStateDisconnected:
		}
	}
}

func serverRecv(connection basenode.Connection) {
	for d := range connection.Receive() {
		processCommand(d)
	}
}

func checkDuplicateSenderIds() {
	for _, d := range state.Devices {
		id1 := d.Id()[3] & 0x7f
		for _, d1 := range state.Devices {
			if d.Id() == d1.Id() {
				continue
			}
			id2 := d1.Id()[3] & 0x7f
			if id2 == id1 {
				log.Error("DUPLICATE ID FOUND when generating senderIds for eltako devices")
			}
		}
	}
}

func processCommand(cmd protocol.Command) {
	log.Debug("INCOMING COMMAND", cmd)
	if len(cmd.Args) == 0 {
		log.Error("Missing device ID in arguments")
		return
	}

	device := state.DeviceByString(cmd.Args[0])
	if device == nil {
		log.Errorf("Device %s does not exist", device)
		return
	}
	switch cmd.Cmd {
	case "toggle":
		device.CmdToggle()
	case "on":
		device.CmdOn()
	case "off":
		device.CmdOff()
	case "dim":
		lvl, _ := strconv.Atoi(cmd.Args[1])
		device.CmdDim(lvl)
	case "learn":
		device.CmdLearn()
	}
}

var enoceanSend chan goenocean.Encoder

func setupEnoceanCommunication(node *protocol.Node, connection basenode.Connection) {

	enoceanSend = make(chan goenocean.Encoder, 100)
	recv := make(chan goenocean.Packet, 100)
	goenocean.Serial(enoceanSend, recv)

	getIDBase()
	reciever(node, connection, recv)
}

func getIDBase() {
	p := goenocean.NewPacket()
	p.SetPacketType(goenocean.PacketTypeCommonCommand)
	p.SetData([]byte{0x08})
	enoceanSend <- p
}

var usb300SenderId [4]byte

func reciever(node *protocol.Node, connection basenode.Connection, recv chan goenocean.Packet) {
	for p := range recv {
		if p.PacketType() == goenocean.PacketTypeResponse && len(p.Data()) == 5 {
			copy(usb300SenderId[:], p.Data()[1:4])
			log.Debugf("senderid: % x ( % x )", usb300SenderId, p.Data())
			continue
		}
		if p.SenderId() != [4]byte{0, 0, 0, 0} {
			incomingPacket(node, connection, p)
		}
	}
}

func incomingPacket(node *protocol.Node, connection basenode.Connection, p goenocean.Packet) {

	var d *Device
	if d = state.Device(p.SenderId()); d == nil {
		//Add unknown device
		d = state.AddDevice(p.SenderId(), "UNKNOWN", nil, false)
		//TODO add to devices list aswell? Maybe we need to configure it first?
		saveDevicesToFile()
		connection.Send(node)
	}

	log.Debug("Incoming packet")
	if t, ok := p.(goenocean.Telegram); ok {
		log.Debug("Packet is goenocean.Telegram")
		for _, deviceEep := range d.RecvEEPs {
			if deviceEep[0:2] != hex.EncodeToString([]byte{t.TelegramType()}) {
				log.Debug("Packet is wrong deviceEep ", deviceEep, t.TelegramType())
				continue
			}

			if h := handlers.getHandler(deviceEep); h != nil {
				h.Process(d, t)
				log.Info("Incoming packet processed from", d.IdString())
				//TODO add return bool in process and to send depending on that!
				connection.Send(node)
				return
			}
		}
	}

	//fmt.Println("Unknown packet")

}

var devFileMutex sync.Mutex

func saveDevicesToFile() {
	devFileMutex.Lock()
	defer devFileMutex.Unlock()
	configFile, err := os.Create("devices.json")
	if err != nil {
		log.Error("creating config file", err.Error())
	}
	var out bytes.Buffer
	b, err := json.MarshalIndent(state.Devices, "", "\t")
	if err != nil {
		log.Error("error marshal json", err)
	}
	json.Indent(&out, b, "", "\t")
	out.WriteTo(configFile)
}
func readConfigFromFile() map[string]*Device {
	devFileMutex.Lock()
	defer devFileMutex.Unlock()
	configFile, err := os.Open("devices.json")
	if err != nil {
		log.Error("opening config file", err.Error())
	}

	config := make(map[string]*Device)
	jsonParser := json.NewDecoder(configFile)
	if err = jsonParser.Decode(&config); err != nil {
		log.Error("parsing config file", err.Error())
	}

	return config
}
<commit_msg>Fix correct id type<commit_after>// Package main provides ...
package main

import (
	"bytes"
	"encoding/hex"
	"encoding/json"
	"flag"
	"os"
	"strconv"
	"sync"

	log "github.com/cihub/seelog"
	"github.com/jonaz/goenocean"
	"github.com/stampzilla/stampzilla-go/nodes/basenode"
	"github.com/stampzilla/stampzilla-go/protocol"
	"github.com/stampzilla/stampzilla-go/protocol/devices"
)

var VERSION string = "dev"
var BUILD_DATE string = ""

var state *State

func main() {

	node := protocol.NewNode("enocean")
	node.Version = VERSION
	node.BuildDate = BUILD_DATE

	flag.Parse()

	//Setup Config
	config := basenode.NewConfig()
	basenode.SetConfig(config)

	//Start communication with the server
	//serverSendChannel = make(chan interface{})
	//serverRecvChannel = make(chan protocol.Command)
	connection := basenode.Connect()
	go monitorState(node, connection)
	go serverRecv(connection)

	// Describe available actions
	node.AddAction("set", "Set", []string{"Devices.Id"})
	node.AddAction("toggle", "Toggle", []string{"Devices.Id"})
	node.AddAction("dim", "Dim", []string{"Devices.Id", "value"})

	// Describe available layouts
	node.AddLayout("1", "switch", "toggle", "Devices", []string{"on"}, "Switches")
	node.AddLayout("2", "slider", "dim", "Devices", []string{"dim"}, "Dimmers")
	node.AddLayout("3", "slider", "dim", "Devices", []string{"dim"}, "Specials")

	node.AddElement(&protocol.Element{
		Type: protocol.ElementTypeToggle,
		Name: "Lamp 0186ff7d",
		Command: &protocol.Command{
			Cmd:  "toggle",
			Args: []string{"0186ff7d"},
		},
		Feedback: `Devices["0186ff7d"].On`,
	})
	node.AddElement(&protocol.Element{
		Type: protocol.ElementTypeText,
		Name: "Lamp 0186ff7d power",
		//Command: &protocol.Command{
		//Cmd:  "toggle",
		//Args: []string{"0186ff7d"},
		//},
		Feedback: `Devices["0186ff7d"].PowerW`,
	})

	//Setup state
	state = NewState()
	state.Devices = readConfigFromFile()
	node.SetState(state)

	//TODO remove element generator and AddElement and AddLayout and AddAction. Devices will superseed that.
	elementGenerator := &ElementGenerator{}
	elementGenerator.State = state
	elementGenerator.Node = node
	elementGenerator.Run()

	for _, dev := range state.Devices {
		// TODO if RecvEEPs is f60201 then its a button and not lamp
		node.Devices().Add(&devices.Device{
			Type:   "lamp",
			Name:   dev.Name,
			Id:     dev.IdString(),
			Online: true,
			Node:   config.Uuid,
			StateMap: map[string]string{
				"On": "Devices[" + dev.IdString() + "]" + ".On",
			},
		})
	}

	checkDuplicateSenderIds()

	setupEnoceanCommunication(node, connection)
}

func monitorState(node *protocol.Node, connection basenode.Connection) {
	for s := range connection.State() {
		switch s {
		case basenode.ConnectionStateConnected:
			connection.Send(node)
		case basenode.ConnectionStateDisconnected:
		}
	}
}

func serverRecv(connection basenode.Connection) {
	for d := range connection.Receive() {
		processCommand(d)
	}
}

func checkDuplicateSenderIds() {
	for _, d := range state.Devices {
		id1 := d.Id()[3] & 0x7f
		for _, d1 := range state.Devices {
			if d.Id() == d1.Id() {
				continue
			}
			id2 := d1.Id()[3] & 0x7f
			if id2 == id1 {
				log.Error("DUPLICATE ID FOUND when generating senderIds for eltako devices")
			}
		}
	}
}

func processCommand(cmd protocol.Command) {
	log.Debug("INCOMING COMMAND", cmd)
	if len(cmd.Args) == 0 {
		log.Error("Missing device ID in arguments")
		return
	}

	device := state.DeviceByString(cmd.Args[0])
	if device == nil {
		log.Errorf("Device %s does not exist", device)
		return
	}
	switch cmd.Cmd {
	case "toggle":
		device.CmdToggle()
	case "on":
		device.CmdOn()
	case "off":
		device.CmdOff()
	case "dim":
		lvl, _ := strconv.Atoi(cmd.Args[1])
		device.CmdDim(lvl)
	case "learn":
		device.CmdLearn()
	}
}

var enoceanSend chan goenocean.Encoder

func setupEnoceanCommunication(node *protocol.Node, connection basenode.Connection) {

	enoceanSend = make(chan goenocean.Encoder, 100)
	recv := make(chan goenocean.Packet, 100)
	goenocean.Serial(enoceanSend, recv)

	getIDBase()
	reciever(node, connection, recv)
}

func getIDBase() {
	p := goenocean.NewPacket()
	p.SetPacketType(goenocean.PacketTypeCommonCommand)
	p.SetData([]byte{0x08})
	enoceanSend <- p
}

var usb300SenderId [4]byte

func reciever(node *protocol.Node, connection basenode.Connection, recv chan goenocean.Packet) {
	for p := range recv {
		if p.PacketType() == goenocean.PacketTypeResponse && len(p.Data()) == 5 {
			copy(usb300SenderId[:], p.Data()[1:4])
			log.Debugf("senderid: % x ( % x )", usb300SenderId, p.Data())
			continue
		}
		if p.SenderId() != [4]byte{0, 0, 0, 0} {
			incomingPacket(node, connection, p)
		}
	}
}

func incomingPacket(node *protocol.Node, connection basenode.Connection, p goenocean.Packet) {

	var d *Device
	if d = state.Device(p.SenderId()); d == nil {
		//Add unknown device
		d = state.AddDevice(p.SenderId(), "UNKNOWN", nil, false)
		//TODO add to devices list aswell? Maybe we need to configure it first?
		saveDevicesToFile()
		connection.Send(node)
	}

	log.Debug("Incoming packet")
	if t, ok := p.(goenocean.Telegram); ok {
		log.Debug("Packet is goenocean.Telegram")
		for _, deviceEep := range d.RecvEEPs {
			if deviceEep[0:2] != hex.EncodeToString([]byte{t.TelegramType()}) {
				log.Debug("Packet is wrong deviceEep ", deviceEep, t.TelegramType())
				continue
			}

			if h := handlers.getHandler(deviceEep); h != nil {
				h.Process(d, t)
				log.Info("Incoming packet processed from", d.IdString())
				//TODO add return bool in process and to send depending on that!
				connection.Send(node)
				return
			}
		}
	}

	//fmt.Println("Unknown packet")

}

var devFileMutex sync.Mutex

func saveDevicesToFile() {
	devFileMutex.Lock()
	defer devFileMutex.Unlock()
	configFile, err := os.Create("devices.json")
	if err != nil {
		log.Error("creating config file", err.Error())
	}
	var out bytes.Buffer
	b, err := json.MarshalIndent(state.Devices, "", "\t")
	if err != nil {
		log.Error("error marshal json", err)
	}
	json.Indent(&out, b, "", "\t")
	out.WriteTo(configFile)
}
func readConfigFromFile() map[string]*Device {
	devFileMutex.Lock()
	defer devFileMutex.Unlock()
	configFile, err := os.Open("devices.json")
	if err != nil {
		log.Error("opening config file", err.Error())
	}

	config := make(map[string]*Device)
	jsonParser := json.NewDecoder(configFile)
	if err = jsonParser.Decode(&config); err != nil {
		log.Error("parsing config file", err.Error())
	}

	return config
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/faiface/beep"
	"github.com/faiface/beep/mp3"
	"github.com/faiface/beep/speaker"
	"github.com/sirupsen/logrus"
	"github.com/stampzilla/stampzilla-go/nodes/stampzilla-server/models/devices"
)

type Player struct {
	Config PlayerConfig
	Name   string

	command chan string
	ctx     context.Context
	cancel  func()
	wg      sync.WaitGroup
}

type PlayerConfig struct {
	Shuffle  bool     `json:"shuffle"`
	Mode     string   `json:"mode"`
	Dir      string   `json:"dir"`
	Playlist []string `json:"playlist"`
}

var players = make(map[string]*Player, 0)

func startPlayers() {
	for name, config := range config.Players {
		player := &Player{
			Config:  config,
			Name:    name,
			command: make(chan string),
		}

		players[name] = player
		player.start()
	}
}

func restartPlayers() {
	for _, player := range players {
		player.stop()
	}
	players = make(map[string]*Player, 0)

	startPlayers()
}

func commandPlayer(player string, state bool) error {
	if p, ok := players[player]; ok {
		if state {
			p.command <- "play"
		} else {
			p.command <- "stop"
		}
		return nil
	}

	return fmt.Errorf("Player was not found")
}

func (player *Player) start() {
	player.ctx, player.cancel = context.WithCancel(context.Background())

	player.wg.Add(1)
	go player.Worker()
}

func (player *Player) stop() {
	player.cancel()
	player.wg.Wait()
}

func (player *Player) Worker() {
	defer player.wg.Done()
	defer logrus.Warnf("Worker %s EXIT", player.Name)

	dev := &devices.Device{
		Name:   "Player " + player.Name,
		ID:     devices.ID{ID: "player:" + player.Name},
		Online: true,
		Traits: []string{"OnOff"},
		State: devices.State{
			"on":   false,
			"file": "",
		},
	}
	n.AddOrUpdate(dev)

	for {
		// Not playing... (wait for shudown or command)
		select {
		case <-player.ctx.Done():
			return
		case cmd := <-player.command:
			if cmd != "play" {
				continue
			}
		}

		done, streamer, err := player.playFile("songs/jingle.mp3")
		if err != nil {
			logrus.Errorf("Playback failed: %s", err.Error())
		}

		dev.State["on"] = true
		n.AddOrUpdate(dev)

		// Playing.. (wait for shutdown, command or playback done)
	L:
		for {
			select {
			case <-player.ctx.Done():
				streamer.Close()
				return
			case cmd := <-player.command:
				if cmd != "stop" {
					continue
				}
				streamer.Close()
			case <-done:
				dev.State["on"] = false
				n.AddOrUpdate(dev)

				break L
			}
		}
		logrus.Info("Playback done")
	}
}

func (player *Player) playFile(file string) (chan struct{}, beep.StreamSeekCloser, error) {
	f, err := os.Open(file)
	if err != nil {
		return nil, nil, err
	}

	streamer, format, err := mp3.Decode(f)
	if err != nil {
		return nil, nil, err
	}

	sr := beep.SampleRate(44100)
	speaker.Init(sr, sr.N(time.Second/10))

	resampled := beep.Resample(4, format.SampleRate, sr, streamer)

	done := make(chan struct{})
	speaker.Play(beep.Seq(resampled, beep.Callback(func() {
		streamer.Close()
		done <- struct{}{}
	})))

	return done, streamer, nil
}
<commit_msg>Added some random and used more config<commit_after>package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/faiface/beep"
	"github.com/faiface/beep/mp3"
	"github.com/faiface/beep/speaker"
	"github.com/sirupsen/logrus"
	"github.com/stampzilla/stampzilla-go/nodes/stampzilla-server/models/devices"
)

type Player struct {
	Config PlayerConfig
	Name   string

	playlist []string
	command  chan string
	ctx      context.Context
	cancel   func()
	wg       sync.WaitGroup
}

type PlayerConfig struct {
	Shuffle  bool     `json:"shuffle"`
	Mode     string   `json:"mode"`
	Dir      string   `json:"dir"`
	Playlist []string `json:"playlist"`
}

var players = make(map[string]*Player, 0)

func startPlayers() {
	for name, config := range config.Players {
		player := &Player{
			Config:  config,
			Name:    name,
			command: make(chan string),
		}

		players[name] = player
		player.start()
	}
}

func restartPlayers() {
	for _, player := range players {
		player.stop()
	}
	players = make(map[string]*Player, 0)

	startPlayers()
}

func commandPlayer(player string, state bool) error {
	if p, ok := players[player]; ok {
		if state {
			p.command <- "play"
		} else {
			p.command <- "stop"
		}
		return nil
	}

	return fmt.Errorf("Player was not found")
}

func (player *Player) start() {
	player.ctx, player.cancel = context.WithCancel(context.Background())

	player.wg.Add(1)
	go player.Worker()
}

func (player *Player) stop() {
	player.cancel()
	player.wg.Wait()
}

func (player *Player) makePlaylist() []string {
	playlist := make([]string, 0)
	files, err := filepath.Glob(player.Config.Dir + "/*.mp3")
	if err != nil {
		log.Fatal(err)
	}

	if player.Config.Shuffle {
		r := rand.New(rand.NewSource(time.Now().Unix()))
		for _, i := range r.Perm(len(files)) {
			playlist = append(playlist, files[i])
		}
	} else {
		for _, f := range files {
			playlist = append(playlist, f)
		}
	}

	return playlist
}

func (player *Player) Worker() {
	defer player.wg.Done()
	defer logrus.Warnf("Worker %s EXIT", player.Name)

	dev := &devices.Device{
		Name:   "Player " + player.Name,
		ID:     devices.ID{ID: "player:" + player.Name},
		Online: true,
		Traits: []string{"OnOff"},
		State: devices.State{
			"on":   false,
			"file": "",
		},
	}
	n.AddOrUpdate(dev)

	for {
		if dev.State["on"] != true {
			// Not playing... (wait for shudown or command)
			select {
			case <-player.ctx.Done():
				return
			case cmd := <-player.command:
				if cmd != "play" {
					continue
				}
			}
		}

		if len(player.playlist) == 0 {
			logrus.Info("Making new playlist")
			player.playlist = player.makePlaylist()
		}
		if len(player.playlist) == 0 {
			logrus.Warn("Playlist is empty, skipping")
			dev.State["on"] = false
			n.AddOrUpdate(dev)
			continue
		}

		nextSong := player.playlist[0]
		player.playlist = player.playlist[1:]

		done, streamer, err := player.playFile(nextSong)
		if err != nil {
			logrus.Errorf("Playback failed: %s", err.Error())
			continue
		}

		dev.State["on"] = true
		n.AddOrUpdate(dev)

		// Playing.. (wait for shutdown, command or playback done)
	L:
		for {
			select {
			case <-player.ctx.Done():
				streamer.Close()
				return
			case cmd := <-player.command:
				if cmd != "stop" {
					continue
				}
				if streamer != nil {
					streamer.Close()
				}
			case <-done:
				if player.Config.Mode == "single" {
					dev.State["on"] = false
					n.AddOrUpdate(dev)
				}

				break L
			}
		}
		logrus.Info("Playback done")
	}
}

func (player *Player) playFile(file string) (chan struct{}, beep.StreamSeekCloser, error) {
	f, err := os.Open(file)
	if err != nil {
		return nil, nil, err
	}

	streamer, format, err := mp3.Decode(f)
	if err != nil {
		return nil, nil, err
	}

	sr := beep.SampleRate(44100)
	speaker.Init(sr, sr.N(time.Second/10))

	resampled := beep.Resample(4, format.SampleRate, sr, streamer)

	done := make(chan struct{})
	speaker.Play(beep.Seq(resampled, beep.Callback(func() {
		streamer.Close()
		done <- struct{}{}
	})))

	return done, streamer, nil
}
<|endoftext|>