<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"strconv"
	"strings"

	"github.com/alexflint/go-arg"
	"github.com/brentp/xopen"
	"github.com/gogetdata/ggd-utils"
)

var args struct {
	Path   string `arg:"positional"`
	Genome string `arg:"-g,required,help:a genome file of chromosome sizes and order"`
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var VCFCommentAfterHeader = errors.New("comment line after non-header #CHROM line in VCF")

type chromStartGetter func([]byte) ([]byte, int, error)

var get_vcf_chrom_start chromStartGetter = func(line []byte) ([]byte, int, error) {
	cpr := bytes.SplitN(line, []byte{'\t'}, 3)
	chrom := cpr[0]
	pos, err := strconv.Atoi(string(cpr[1]))
	return chrom, pos, err
}

var get_gff_chrom_start chromStartGetter = func(line []byte) ([]byte, int, error) {
	toks := bytes.SplitN(line, []byte{'\t'}, 5)
	start, err := strconv.Atoi(string(toks[3]))
	return toks[0], start, err
}

func main() {

	p := arg.MustParse(&args)
	if args.Path == "" || args.Genome == "" {
		p.Fail("must specify a path to check and a genome file")
	}
	gf, err := ggd_utils.ReadGenomeFile(args.Genome)
	if err != nil {
		log.Fatal(err)
	}

	if strings.HasSuffix(args.Path, ".vcf.gz") {
		checkVCF(args.Path, gf)
	} else if strings.HasSuffix(args.Path, ".bed.gz") {
		checkTab(args.Path, gf, get_vcf_chrom_start)
	} else {
		found := false
		for _, suff := range []string{"gff", "gtf", "gff3", "gff2"} {
			if strings.HasSuffix(args.Path, suff) || strings.HasSuffix(args.Path, suff+".gz") {
				found = true
				checkTab(args.Path, gf, get_gff_chrom_start)
				break
			}
		}
		if !found {
			log.Fatalf("Don't know how to check this type of file: %s\n", args.Path)
		}
	}
}

func checkLine(iline int, line []byte, lastChrom []byte, lastStart int,
	get_chrom_start func([]byte) ([]byte, int, error),
	cmp func(a, b string) bool) ([]byte, int, error) {

	chrom, start, err := get_chrom_start(line)
	if err != nil {
		return chrom, start, err
	}
	if !bytes.Equal(chrom, lastChrom) {
		if len(lastChrom) != 0 {
			if !cmp(string(lastChrom), string(chrom)) {
				return chrom, start, fmt.Errorf("chromosomes not in specified sort order: %s, %s at line %d\nuse gsort (https://github.com/brentp/gsort/) to order according to the genome file", lastChrom, chrom, iline)
			}
		}
		lastChrom = chrom
		lastStart = start
	} else {
		if start < lastStart {
			return chrom, start, fmt.Errorf("positions not sorted: %d => %d at line %d\n", lastStart, start, iline)
		}
		lastStart = start
	}
	if start < 0 {
		return chrom, start, fmt.Errorf("negative position at line: %d (%d)", iline, start)
	}

	return chrom, start, nil
}

func checkTab(path string, gf *ggd_utils.GenomeFile, getter chromStartGetter) {
	if !(xopen.Exists(path+".tbi") || xopen.Exists(path+".csi")) {
		log.Fatalf("BED: %s should have a .tbi", path)
	}
	rdr, err := xopen.Ropen(path)
	iline := 1
	check(err)
	afterHeader := false
	lastChrom := []byte("")
	lastPos := -1
	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		if len(line) == 0 {
			continue
		}
		if line[0] == '#' || bytes.HasPrefix(line, []byte("track")) || bytes.HasPrefix(line, []byte("browser")) {
			if afterHeader {
				log.Fatalf("found comment/header line after header at line: %d", iline)
			}
		} else {
			afterHeader = true
			lastChrom, lastPos, err = checkLine(iline, line, lastChrom, lastPos, getter, gf.Less)
			if err != nil {
				log.Fatal(err)
			}
			if gf.Lengths[string(lastChrom)] < lastPos {
				log.Fatalf("position: %d beyond end of chromosome %s", lastPos, lastChrom)
			}
		}
		iline += 1
	}
}

func checkVCF(path string, gf *ggd_utils.GenomeFile) {
	if !(xopen.Exists(path+".tbi") || xopen.Exists(path+".csi")) {
		log.Fatal("VCF should have a .tbi or .csi")
	}

	rdr, err := xopen.Ropen(path)
	check(err)
	afterHeader := false
	iline := 1
	lastChrom := []byte("")
	lastPos := -1

	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		if iline == 1 {
			if !bytes.HasPrefix(line, []byte("##fileformat=VCF")) {
				log.Println(string(line))
				log.Fatal("VCF header line '##fileformat=VCF... not found")
			}
		}
		if line[0] == '#' {
			if bytes.HasPrefix(line, []byte("#CHROM\t")) {
				afterHeader = true
			} else if afterHeader {
				log.Fatal(VCFCommentAfterHeader, iline)
			}
		} else {
			if !afterHeader {
				log.Fatal("VCF header line '#CHROM\t ... not found")
			}
			lastChrom, lastPos, err = checkLine(iline, line, lastChrom, lastPos, get_vcf_chrom_start, gf.Less)
			if err != nil {
				log.Fatal(err)
			}
			if gf.Lengths[string(lastChrom)] < lastPos {
				log.Fatalf("position: %d beyond end of chromosome %s", lastPos, lastChrom)
			}
		}
		iline += 1
	}
}
<commit_msg>check for EOF<commit_after>package main

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/alexflint/go-arg"
	"github.com/biogo/hts/bgzf"
	"github.com/brentp/xopen"
	"github.com/gogetdata/ggd-utils"
)

var args struct {
	Path   string `arg:"positional"`
	Genome string `arg:"-g,required,help:a genome file of chromosome sizes and order"`
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var VCFCommentAfterHeader = errors.New("comment line after non-header #CHROM line in VCF")

type chromStartGetter func([]byte) ([]byte, int, error)

var get_vcf_chrom_start chromStartGetter = func(line []byte) ([]byte, int, error) {
	cpr := bytes.SplitN(line, []byte{'\t'}, 3)
	chrom := cpr[0]
	pos, err := strconv.Atoi(string(cpr[1]))
	return chrom, pos, err
}

var get_gff_chrom_start chromStartGetter = func(line []byte) ([]byte, int, error) {
	toks := bytes.SplitN(line, []byte{'\t'}, 5)
	start, err := strconv.Atoi(string(toks[3]))
	return toks[0], start, err
}

func main() {

	p := arg.MustParse(&args)
	if args.Path == "" || args.Genome == "" {
		p.Fail("must specify a path to check and a genome file")
	}
	gf, err := ggd_utils.ReadGenomeFile(args.Genome)
	if err != nil {
		log.Fatal(err)
	}

	if strings.HasSuffix(args.Path, ".vcf.gz") {
		checkVCF(args.Path, gf)
	} else if strings.HasSuffix(args.Path, ".bed.gz") {
		checkTab(args.Path, gf, get_vcf_chrom_start)
	} else {
		found := false
		for _, suff := range []string{"gff", "gtf", "gff3", "gff2"} {
			if strings.HasSuffix(args.Path, suff) || strings.HasSuffix(args.Path, suff+".gz") {
				found = true
				checkTab(args.Path, gf, get_gff_chrom_start)
				break
			}
		}
		if !found {
			log.Fatalf("Don't know how to check this type of file: %s\n", args.Path)
		}
	}
}

func checkLine(iline int, line []byte, lastChrom []byte, lastStart int,
	get_chrom_start func([]byte) ([]byte, int, error),
	cmp func(a, b string) bool) ([]byte, int, error) {

	chrom, start, err := get_chrom_start(line)
	if err != nil {
		return chrom, start, err
	}
	if !bytes.Equal(chrom, lastChrom) {
		if len(lastChrom) != 0 {
			if !cmp(string(lastChrom), string(chrom)) {
				return chrom, start, fmt.Errorf("chromosomes not in specified sort order: %s, %s at line %d\nuse gsort (https://github.com/brentp/gsort/) to order according to the genome file", lastChrom, chrom, iline)
			}
		}
		lastChrom = chrom
		lastStart = start
	} else {
		if start < lastStart {
			return chrom, start, fmt.Errorf("positions not sorted: %d => %d at line %d\n", lastStart, start, iline)
		}
		lastStart = start
	}
	if start < 0 {
		return chrom, start, fmt.Errorf("negative position at line: %d (%d)", iline, start)
	}

	return chrom, start, nil
}

func checkTab(path string, gf *ggd_utils.GenomeFile, getter chromStartGetter) {
	if !(xopen.Exists(path+".tbi") || xopen.Exists(path+".csi")) {
		log.Fatalf("BED: %s should have a .tbi", path)
	}
	fh, err := os.Open(path)
	if err != nil {
		log.Fatalf("BED: unable to open file: %s", path)
	}
	if ok, _ := bgzf.HasEOF(fh); !ok {
		log.Fatal("missing EOF")
	}

	bgz, err := bgzf.NewReader(fh, 1)
	defer bgz.Close()
	defer fh.Close()

	rdr := bufio.NewReader(bgz)
	iline := 1
	check(err)
	afterHeader := false
	lastChrom := []byte("")
	lastPos := -1
	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		if len(line) == 0 {
			continue
		}
		if line[0] == '#' || bytes.HasPrefix(line, []byte("track")) || bytes.HasPrefix(line, []byte("browser")) {
			if afterHeader {
				log.Fatalf("found comment/header line after header at line: %d", iline)
			}
		} else {
			afterHeader = true
			lastChrom, lastPos, err = checkLine(iline, line, lastChrom, lastPos, getter, gf.Less)
			if err != nil {
				log.Fatal(err)
			}
			if gf.Lengths[string(lastChrom)] < lastPos {
				log.Fatalf("position: %d beyond end of chromosome %s", lastPos, lastChrom)
			}
		}
		iline += 1
	}
}

func checkVCF(path string, gf *ggd_utils.GenomeFile) {
	if !(xopen.Exists(path+".tbi") || xopen.Exists(path+".csi")) {
		log.Fatal("VCF should have a .tbi or .csi")
	}

	fh, err := os.Open(path)
	check(err)
	defer fh.Close()
	if ok, _ := bgzf.HasEOF(fh); !ok {
		log.Fatalf("missing EOF in %s", path)
	}

	bgz, err := bgzf.NewReader(fh, 1)
	check(err)
	defer bgz.Close()
	rdr := bufio.NewReader(bgz)

	afterHeader := false
	iline := 1
	lastChrom := []byte("")
	lastPos := -1

	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		if iline == 1 {
			if !bytes.HasPrefix(line, []byte("##fileformat=VCF")) {
				log.Println(string(line))
				log.Fatal("VCF header line '##fileformat=VCF... not found")
			}
		}
		if line[0] == '#' {
			if bytes.HasPrefix(line, []byte("#CHROM\t")) {
				afterHeader = true
			} else if afterHeader {
				log.Fatal(VCFCommentAfterHeader, iline)
			}
		} else {
			if !afterHeader {
				log.Fatal("VCF header line '#CHROM\t ... not found")
			}
			lastChrom, lastPos, err = checkLine(iline, line, lastChrom, lastPos, get_vcf_chrom_start, gf.Less)
			if err != nil {
				log.Fatal(err)
			}
			if gf.Lengths[string(lastChrom)] < lastPos {
				log.Fatalf("position: %d beyond end of chromosome %s", lastPos, lastChrom)
			}
		}
		iline += 1
	}
}
<|endoftext|>


<commit_before>// These tests verify the test running logic.

package gocheck_test

import (
	. "launchpad.net/gocheck"
	"time"
)

var benchmarkS = Suite(&BenchmarkS{})

type BenchmarkS struct{}

func (s *BenchmarkS) TestCountSuite(c *C) {
	suitesRun += 1
}

func (s *BenchmarkS) TestBasicTestTiming(c *C) {
	helper := FixtureHelper{sleepOn: "Test1", sleep: 1000000 * time.Nanosecond}
	output := String{}
	runConf := RunConf{Output: &output, Verbose: true}
	Run(&helper, &runConf)

	expected := "PASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.Test1\t0\\.001s\n" +
		"PASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.Test2\t0\\.000s\n"
	c.Assert(output.value, Matches, expected)
}

func (s *BenchmarkS) TestStreamTestTiming(c *C) {
	helper := FixtureHelper{sleepOn: "SetUpSuite", sleep: 1000000 * time.Nanosecond}
	output := String{}
	runConf := RunConf{Output: &output, Stream: true}
	Run(&helper, &runConf)

	expected := "(?s).*\nPASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.SetUpSuite\t *0\\.001s\n.*"
	c.Assert(output.value, Matches, expected)
}

func (s *BenchmarkS) TestBenchmark(c *C) {
	helper := FixtureHelper{sleep: 100000}
	output := String{}
	runConf := RunConf{
		Output: &output,
		Benchmark: true,
		BenchmarkTime: 10000000,
		Filter: "Benchmark1",
	}
	Run(&helper, &runConf)
	c.Check(helper.calls[0], Equals, "SetUpSuite")
	c.Check(helper.calls[1], Equals, "SetUpTest")
	c.Check(helper.calls[2], Equals, "Benchmark1")
	c.Check(helper.calls[3], Equals, "TearDownTest")
	c.Check(helper.calls[4], Equals, "SetUpTest")
	c.Check(helper.calls[5], Equals, "Benchmark1")
	c.Check(helper.calls[6], Equals, "TearDownTest")
	// ... and more.

	expected := "PASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.Benchmark1\t *100\t *[12][0-9]{5} ns/op\n"
	c.Assert(output.value, Matches, expected)
}

func (s *BenchmarkS) TestBenchmarkBytes(c *C) {
	helper := FixtureHelper{sleep: 100000}
	output := String{}
	runConf := RunConf{
		Output: &output,
		Benchmark: true,
		BenchmarkTime: 10000000,
		Filter: "Benchmark2",
	}
	Run(&helper, &runConf)

	expected := "PASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.Benchmark2\t *100\t *[12][0-9]{5} ns/op\t *[4-9]\\.[0-9]{2} MB/s\n"
	c.Assert(output.value, Matches, expected)
}
<commit_msg>benchmark_test: windows benchmark do not loop as much/fast<commit_after>// These tests verify the test running logic.

package gocheck_test

import (
	. "launchpad.net/gocheck"
	"time"
)

var benchmarkS = Suite(&BenchmarkS{})

type BenchmarkS struct{}

func (s *BenchmarkS) TestCountSuite(c *C) {
	suitesRun += 1
}

func (s *BenchmarkS) TestBasicTestTiming(c *C) {
	helper := FixtureHelper{sleepOn: "Test1", sleep: 1000000 * time.Nanosecond}
	output := String{}
	runConf := RunConf{Output: &output, Verbose: true}
	Run(&helper, &runConf)

	expected := "PASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.Test1\t0\\.001s\n" +
		"PASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.Test2\t0\\.000s\n"
	c.Assert(output.value, Matches, expected)
}

func (s *BenchmarkS) TestStreamTestTiming(c *C) {
	helper := FixtureHelper{sleepOn: "SetUpSuite", sleep: 1000000 * time.Nanosecond}
	output := String{}
	runConf := RunConf{Output: &output, Stream: true}
	Run(&helper, &runConf)

	expected := "(?s).*\nPASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.SetUpSuite\t *0\\.001s\n.*"
	c.Assert(output.value, Matches, expected)
}

func (s *BenchmarkS) TestBenchmark(c *C) {
	helper := FixtureHelper{sleep: 100000}
	output := String{}
	runConf := RunConf{
		Output: &output,
		Benchmark: true,
		BenchmarkTime: 10000000,
		Filter: "Benchmark1",
	}
	Run(&helper, &runConf)
	c.Check(helper.calls[0], Equals, "SetUpSuite")
	c.Check(helper.calls[1], Equals, "SetUpTest")
	c.Check(helper.calls[2], Equals, "Benchmark1")
	c.Check(helper.calls[3], Equals, "TearDownTest")
	c.Check(helper.calls[4], Equals, "SetUpTest")
	c.Check(helper.calls[5], Equals, "Benchmark1")
	c.Check(helper.calls[6], Equals, "TearDownTest")
	// ... and more.

	expected := "PASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.Benchmark1\t *(20|100)\t *[12][0-9]{5,7} ns/op\n"
	c.Assert(output.value, Matches, expected)
}

func (s *BenchmarkS) TestBenchmarkBytes(c *C) {
	helper := FixtureHelper{sleep: 100000}
	output := String{}
	runConf := RunConf{
		Output: &output,
		Benchmark: true,
		BenchmarkTime: 10000000,
		Filter: "Benchmark2",
	}
	Run(&helper, &runConf)

	expected := "PASS: gocheck_test\\.go:[0-9]+: FixtureHelper\\.Benchmark2\t *(20|100)\t *[12][0-9]{5,7} ns/op\t *[1-9]\\.[0-9]{2} MB/s\n"
	c.Assert(output.value, Matches, expected)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"
	"time"

	"google.golang.org/grpc"

	slackbot "github.com/BeepBoopHQ/go-slackbot"
	"github.com/nlopes/slack"
	"golang.org/x/net/context"

	"github.com/moul/kafka-gateway/gen/pb"
)

var (
	client kafkapb.KafkaServiceClient
	conn   *grpc.ClientConn
	ctx    context.Context
)

func init() {
	var err error
	conn, err = grpc.Dial(
		"127.0.0.1:9000",
		grpc.WithInsecure(),
		grpc.WithTimeout(time.Second),
	)
	if err != nil {
		panic(err)
	}
	ctx = context.Background()
	client = kafkapb.NewKafkaServiceClient(conn)
}

func main() {
	bot := slackbot.New(os.Getenv("SLACK_TOKEN"))
	bot.Hear("(?)kafka produce (.*)").MessageHandler(kafkaProduceHandler)
	bot.Run()
}

func kafkaProduceHandler(ctx context.Context, bot *slackbot.Bot, msg *slack.MessageEvent) {
	text := strings.Join(strings.Split(slackbot.StripDirectMention(msg.Text), " ")[2:], " ")
	ret, err := client.Producer(ctx, &kafkapb.ProducerRequest{
		Topic: "test",
		Key:   "test",
		Value: text,
	})
	if err != nil {
		bot.Reply(msg, fmt.Sprintf("Failed: %v", err), slackbot.WithTyping)
	} else {
		bot.Reply(msg, fmt.Sprintf("Done: %v", ret), slackbot.WithTyping)
	}
}
<commit_msg>Add a kafka consumer routine<commit_after>package main

import (
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"google.golang.org/grpc"

	slackbot "github.com/BeepBoopHQ/go-slackbot"
	"github.com/nlopes/slack"
	"golang.org/x/net/context"

	"github.com/moul/kafka-gateway/gen/pb"
)

var (
	client kafkapb.KafkaServiceClient
	conn   *grpc.ClientConn
	ctx    context.Context
	bot    *slackbot.Bot
)

func main() {
	var err error
	conn, err = grpc.Dial(
		"127.0.0.1:9000",
		grpc.WithInsecure(),
		grpc.WithTimeout(time.Second),
	)
	if err != nil {
		panic(err)
	}
	ctx = context.Background()
	client = kafkapb.NewKafkaServiceClient(conn)

	go kafkaConsumer()

	bot = slackbot.New(os.Getenv("SLACK_TOKEN"))
	bot.Hear("(?)kafka produce (.*)").MessageHandler(kafkaProduceHandler)
	bot.Run()
}

func kafkaConsumer() {
	stream, err := client.ConsumerStream(ctx, &kafkapb.ConsumerStreamRequest{
		Topics:   []string{"test"},
		ClientId: "slackbot",
	})
	if err != nil {
		panic(err)
	}

	for {
		message, err := stream.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			panic(err)
		}
		bot.RTM.SendMessage(bot.RTM.NewOutgoingMessage(fmt.Sprintf("<<< %v", message.Value), os.Getenv("SLACK_CHANNEL")))
	}
}

func kafkaProduceHandler(ctx context.Context, bot *slackbot.Bot, msg *slack.MessageEvent) {
	text := strings.Join(strings.Split(slackbot.StripDirectMention(msg.Text), " ")[2:], " ")
	ret, err := client.Producer(ctx, &kafkapb.ProducerRequest{
		Topic: "test",
		Key:   "test",
		Value: text,
	})
	if err != nil {
		bot.Reply(msg, fmt.Sprintf("Failed: %v", err), slackbot.WithTyping)
	} else {
		bot.Reply(msg, fmt.Sprintf("Done: %v", ret), slackbot.WithTyping)
	}
}
<|endoftext|>


<commit_before>package gobls

import (
	"bufio"
	"bytes"
	"fmt"
	"testing"
)

func benchmarkScanner(b *testing.B, lineLength int, makeScanner func(buf []byte) Scanner) {
	b.Helper()

	wanted := makeBuffer(1, lineLength)
	wanted = wanted[:len(wanted)-2] // trim final CRLF

	// NOTE: make buffer with line count set to b.N
	s := makeScanner(makeBuffer(b.N, lineLength))

	var count int

	b.ResetTimer()

	for s.Scan() {
		if got := s.Bytes(); !bytes.Equal(got, wanted) {
			b.Errorf("GOT: %#v; WANT: %#v", got, wanted)
		}
		count++
	}

	if got, want := s.Err(), error(nil); got != want {
		b.Errorf("GOT: %#v; WANT: %#v", got, want)
	}
	if got, want := count, b.N; got != want {
		b.Errorf("GOT: %#v; WANT: %#v", got, want)
	}
}

func BenchmarkScanner(b *testing.B) {
	for _, i := range []uint{6, 7, 8, 9, 10, 11, 12} {
		lineLength := 1 << i

		b.Run(fmt.Sprintf("%04d", lineLength), func(b *testing.B) {
			b.Run("bufio", func(b *testing.B) {
				benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
					return bufio.NewScanner(bytes.NewReader(buf))
				})
			})
			b.Run("reader", func(b *testing.B) {
				benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
					return NewScanner(bytes.NewReader(buf))
				})
			})
			b.Run("buffer", func(b *testing.B) {
				benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
					return NewBufferScanner(buf)
				})
			})
		})
	}

	b.Run("excessively long", func(b *testing.B) {
		const lineLength = excessivelyLongLineLength
		b.Run("bufio", func(b *testing.B) {
			benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
				return bufio.NewScanner(bytes.NewReader(buf))
			})
		})
		b.Run("reader", func(b *testing.B) {
			benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
				return NewScanner(bytes.NewReader(buf))
			})
		})
		b.Run("buffer", func(b *testing.B) {
			benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
				return NewBufferScanner(buf)
			})
		})
	})
}
<commit_msg>fix benchmark labels<commit_after>package gobls

import (
	"bufio"
	"bytes"
	"fmt"
	"testing"
)

func benchmarkScanner(b *testing.B, lineLength int, makeScanner func(buf []byte) Scanner) {
	b.Helper()

	wanted := makeBuffer(1, lineLength)
	wanted = wanted[:len(wanted)-2] // trim final CRLF

	// NOTE: make buffer with line count set to b.N
	s := makeScanner(makeBuffer(b.N, lineLength))

	var count int

	b.ResetTimer()

	for s.Scan() {
		if got := s.Bytes(); !bytes.Equal(got, wanted) {
			b.Errorf("GOT: %#v; WANT: %#v", got, wanted)
		}
		count++
	}

	if got, want := s.Err(), error(nil); got != want {
		b.Errorf("GOT: %#v; WANT: %#v", got, want)
	}
	if got, want := count, b.N; got != want {
		b.Errorf("GOT: %#v; WANT: %#v", got, want)
	}
}

func BenchmarkScanner(b *testing.B) {
	for _, i := range []uint{6, 7, 8, 9, 10, 11, 12} {
		lineLength := 1 << i

		b.Run(fmt.Sprintf("%04d", lineLength), func(b *testing.B) {
			b.Run("bufio", func(b *testing.B) {
				benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
					return bufio.NewScanner(bytes.NewReader(buf))
				})
			})
			b.Run("scanner", func(b *testing.B) {
				benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
					return NewScanner(bytes.NewReader(buf))
				})
			})
			b.Run("buffer", func(b *testing.B) {
				benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
					return NewBufferScanner(buf)
				})
			})
		})
	}

	b.Run("excessively long", func(b *testing.B) {
		const lineLength = excessivelyLongLineLength
		b.Run("bufio", func(b *testing.B) {
			benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
				return bufio.NewScanner(bytes.NewReader(buf))
			})
		})
		b.Run("scanner", func(b *testing.B) {
			benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
				return NewScanner(bytes.NewReader(buf))
			})
		})
		b.Run("buffer", func(b *testing.B) {
			benchmarkScanner(b, lineLength, func(buf []byte) Scanner {
				return NewBufferScanner(buf)
			})
		})
	})
}
<|endoftext|>


<commit_before>package create

import (
	"crypto/tls"
	"fmt"
	"os"

	"github.com/codegangsta/cli"
	"github.com/materials-commons/mcstore/cmd/pkg/opts"
	"github.com/materials-commons/mcstore/cmd/pkg/project"
	"github.com/materials-commons/mcstore/pkg/files"
	"github.com/parnurzeal/gorequest"
)

type createCommandArgs struct {
	projectName   string
	projectID     string
	directoryPath string
	n             int
}

var (
	// Command contains the arguments and function for the cli project create command.
	Command = cli.Command{
		Name:    "create",
		Aliases: []string{"cr", "c"},
		Usage:   "Create a new project",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "directory, dir, d",
				Usage: "The base directory for the project",
			},
			cli.StringFlag{
				Name:  "projectid, pid, p",
				Usage: "The projectid on Materials Commons",
			},
			cli.IntFlag{
				Name:  "parallel, n",
				Value: 3,
				Usage: "Number of simultaneous uploads to perform, defaults to 3",
			},
			cli.IntFlag{
				Name:  "parallel, n",
				Value: 3,
				Usage: "Number of simultaneous uploads to perform, defaults to 3",
			},
		},
		Action: createCLI,
	}

	args createCommandArgs
)

// createCLI implements the project create command.
func createCLI(c *cli.Context) {
	if err := validateArgs(c); err != nil {
		fmt.Println("Invalid arguments:", err)
		os.Exit(1)
	}

	proj, err := project.Create(args.directoryPath, args.projectName, args.projectID)
	if err != nil {
		fmt.Println("Unable to create project:", err)
		os.Exit(1)
	}

	fmt.Println("Indexing project...")
	indexProject(args.directoryPath, proj)
	fmt.Println("Done.")
}

// validate will validate the command line arguments. It will print a message
// and exit if there is a bad argument.
func validateArgs(c *cli.Context) error {
	if len(c.Args()) != 1 {
		return fmt.Errorf("You must supply a name for the project.")
	}
	args.projectName = c.Args()[0]

	args.projectID = c.String("projectid")
	if err := validateProject(args.projectID); err != nil {
		return err
	}

	args.directoryPath = c.String("directory")
	if err := validateDirectoryPath(args.directoryPath); err != nil {
		return err
	}

	args.n = opts.GetNumThreads(c)

	return nil
}

// validateProject ensures that the projectID exists and the user has access.
func validateProject(projectID string) error {
	if projectID == "" {
		return fmt.Errorf("You must specify a project id from Materials Commons.")
	}
	return nil
}

// validateDirectoryPath checks that the given directory path exists.
func validateDirectoryPath(path string) error {
	if path == "" {
		return fmt.Errorf("You must specify a local directory path where the project files are located.")
	}

	if _, err := os.Stat(path); err != nil {
		return fmt.Errorf("Invalid directory: %s", path)
	}
	return nil
}

// indexProject walks the directory tree and indexes each of the files found. Indexing
// can be performed in parallel.
func indexProject(path string, proj *project.MCProject) error {
	fn := func(done <-chan struct{}, entries <-chan files.TreeEntry, result chan<- string) {
		indexEntries(proj, done, entries, result)
	}
	walker := &files.PWalker{
		NumParallel: args.n,
		ProcessFn:   fn,
		ProcessDirs: true,
	}
	walker.PWalk(args.directoryPath)
	return nil
}

// indexEntries processes the entries sent along the entries channel. It also
// processes done channel events by exiting the go routine.
func indexEntries(proj *project.MCProject, done <-chan struct{}, entries <-chan files.TreeEntry, result chan<- string) {
	client := gorequest.New().TLSClientConfig(&tls.Config{InsecureSkipVerify: true})
	for entry := range entries {
		select {
		case result <- indexEntry(proj, entry, client):
		case <-done:
			return
		}
	}
}

func indexEntry(proj *project.MCProject, entry files.TreeEntry, client *gorequest.SuperAgent) string {
	switch {
	case entry.Finfo.IsDir():

	default:
	}
	return ""
}
<commit_msg>Remove projectid and duplicate parallel flags. Create doesn't have a project id as it is creating the project.<commit_after>package create

import (
	"crypto/tls"
	"fmt"
	"os"

	"github.com/codegangsta/cli"
	"github.com/materials-commons/mcstore/cmd/pkg/opts"
	"github.com/materials-commons/mcstore/cmd/pkg/project"
	"github.com/materials-commons/mcstore/pkg/files"
	"github.com/parnurzeal/gorequest"
)

type createCommandArgs struct {
	projectName   string
	projectID     string
	directoryPath string
	n             int
}

var (
	// Command contains the arguments and function for the cli project create command.
	Command = cli.Command{
		Name:    "create",
		Aliases: []string{"cr", "c"},
		Usage:   "Create a new project",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "directory, dir, d",
				Usage: "The base directory for the project",
			},
			cli.IntFlag{
				Name:  "parallel, n",
				Value: 3,
				Usage: "Number of simultaneous uploads to perform, defaults to 3",
			},
		},
		Action: createCLI,
	}

	args createCommandArgs
)

// createCLI implements the project create command.
func createCLI(c *cli.Context) {
	if err := validateArgs(c); err != nil {
		fmt.Println("Invalid arguments:", err)
		os.Exit(1)
	}

	proj, err := project.Create(args.directoryPath, args.projectName, args.projectID)
	if err != nil {
		fmt.Println("Unable to create project:", err)
		os.Exit(1)
	}

	fmt.Println("Indexing project...")
	indexProject(args.directoryPath, proj)
	fmt.Println("Done.")
}

// validate will validate the command line arguments. It will print a message
// and exit if there is a bad argument.
func validateArgs(c *cli.Context) error {
	if len(c.Args()) != 1 {
		return fmt.Errorf("You must supply a name for the project.")
	}
	args.projectName = c.Args()[0]

	if err := createProject(args.projectName); err != nil {
		return err
	}

	args.directoryPath = c.String("directory")
	if err := validateDirectoryPath(args.directoryPath); err != nil {
		return err
	}

	args.n = opts.GetNumThreads(c)

	return nil
}

// validateProject ensures that the projectID exists and the user has access.
func createProject(projectName string) error {

	return nil
}

// validateDirectoryPath checks that the given directory path exists.
func validateDirectoryPath(path string) error {
	if path == "" {
		return fmt.Errorf("You must specify a local directory path where the project files are located.")
	}

	if _, err := os.Stat(path); err != nil {
		return fmt.Errorf("Invalid directory: %s", path)
	}
	return nil
}

// indexProject walks the directory tree and indexes each of the files found. Indexing
// can be performed in parallel.
func indexProject(path string, proj *project.MCProject) error {
	fn := func(done <-chan struct{}, entries <-chan files.TreeEntry, result chan<- string) {
		indexEntries(proj, done, entries, result)
	}
	walker := &files.PWalker{
		NumParallel: args.n,
		ProcessFn:   fn,
		ProcessDirs: true,
	}
	walker.PWalk(args.directoryPath)
	return nil
}

// indexEntries processes the entries sent along the entries channel. It also
// processes done channel events by exiting the go routine.
func indexEntries(proj *project.MCProject, done <-chan struct{}, entries <-chan files.TreeEntry, result chan<- string) {
	client := gorequest.New().TLSClientConfig(&tls.Config{InsecureSkipVerify: true})
	for entry := range entries {
		select {
		case result <- indexEntry(proj, entry, client):
		case <-done:
			return
		}
	}
}

func indexEntry(proj *project.MCProject, entry files.TreeEntry, client *gorequest.SuperAgent) string {
	switch {
	case entry.Finfo.IsDir():

	default:
	}
	return ""
}
<|endoftext|>


<commit_before>package bencode

import (
	"io"
	"math/big"
	"reflect"
	"runtime"
	"sort"
	"strconv"
	"sync"

	"github.com/anacrolix/missinggo"
)

func isEmptyValue(v reflect.Value) bool {
	return missinggo.IsEmptyValue(v)
}

type Encoder struct {
	w interface {
		Flush() error
		io.Writer
		WriteString(string) (int, error)
	}
	scratch [64]byte
}

func (e *Encoder) Encode(v interface{}) (err error) {
	if v == nil {
		return
	}
	defer func() {
		if e := recover(); e != nil {
			if _, ok := e.(runtime.Error); ok {
				panic(e)
			}
			var ok bool
			err, ok = e.(error)
			if !ok {
				panic(e)
			}
		}
	}()
	e.reflectValue(reflect.ValueOf(v))
	return e.w.Flush()
}

type string_values []reflect.Value

func (sv string_values) Len() int           { return len(sv) }
func (sv string_values) Swap(i, j int)      { sv[i], sv[j] = sv[j], sv[i] }
func (sv string_values) Less(i, j int) bool { return sv.get(i) < sv.get(j) }
func (sv string_values) get(i int) string   { return sv[i].String() }

func (e *Encoder) write(s []byte) {
	_, err := e.w.Write(s)
	if err != nil {
		panic(err)
	}
}

func (e *Encoder) writeString(s string) {
	_, err := e.w.WriteString(s)
	if err != nil {
		panic(err)
	}
}

func (e *Encoder) reflectString(s string) {
	b := strconv.AppendInt(e.scratch[:0], int64(len(s)), 10)
	e.write(b)
	e.writeString(":")
	e.writeString(s)
}

func (e *Encoder) reflectByteSlice(s []byte) {
	b := strconv.AppendInt(e.scratch[:0], int64(len(s)), 10)
	e.write(b)
	e.writeString(":")
	e.write(s)
}

// returns true if the value implements Marshaler interface and marshaling was
// done successfully
func (e *Encoder) reflectMarshaler(v reflect.Value) bool {
	m, ok := v.Interface().(Marshaler)
	if !ok {
		// T doesn't work, try *T
		if v.Kind() != reflect.Ptr && v.CanAddr() {
			m, ok = v.Addr().Interface().(Marshaler)
			if ok {
				v = v.Addr()
			}
		}
	}
	if ok && (v.Kind() != reflect.Ptr || !v.IsNil()) {
		data, err := m.MarshalBencode()
		if err != nil {
			panic(&MarshalerError{v.Type(), err})
		}
		e.write(data)
		return true
	}

	return false
}

func (e *Encoder) reflectValue(v reflect.Value) {

	if e.reflectMarshaler(v) {
		return
	}

	switch t := v.Interface().(type) {
	case big.Int:
		e.writeString("i")
		e.writeString(t.String())
		e.writeString("e")
		return
	}

	switch v.Kind() {
	case reflect.Bool:
		if v.Bool() {
			e.writeString("i1e")
		} else {
			e.writeString("i0e")
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		b := strconv.AppendInt(e.scratch[:0], v.Int(), 10)
		e.writeString("i")
		e.write(b)
		e.writeString("e")
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		b := strconv.AppendUint(e.scratch[:0], v.Uint(), 10)
		e.writeString("i")
		e.write(b)
		e.writeString("e")
	case reflect.String:
		e.reflectString(v.String())
	case reflect.Struct:
		e.writeString("d")
		for _, ef := range encodeFields(v.Type()) {
			field_value := v.Field(ef.i)
			if ef.omit_empty && isEmptyValue(field_value) {
				continue
			}
			e.reflectString(ef.tag)
			e.reflectValue(field_value)
		}
		e.writeString("e")
	case reflect.Map:
		if v.Type().Key().Kind() != reflect.String {
			panic(&MarshalTypeError{v.Type()})
		}
		if v.IsNil() {
			e.writeString("de")
			break
		}
		e.writeString("d")
		sv := string_values(v.MapKeys())
		sort.Sort(sv)
		for _, key := range sv {
			e.reflectString(key.String())
			e.reflectValue(v.MapIndex(key))
		}
		e.writeString("e")
	case reflect.Slice:
		if v.IsNil() {
			e.writeString("le")
			break
		}
		if v.Type().Elem().Kind() == reflect.Uint8 {
			s := v.Bytes()
			e.reflectByteSlice(s)
			break
		}
		fallthrough
	case reflect.Array:
		e.writeString("l")
		for i, n := 0, v.Len(); i < n; i++ {
			e.reflectValue(v.Index(i))
		}
		e.writeString("e")
	case reflect.Interface:
		e.reflectValue(v.Elem())
	case reflect.Ptr:
		if v.IsNil() {
			v = reflect.Zero(v.Type().Elem())
		} else {
			v = v.Elem()
		}
		e.reflectValue(v)
	default:
		panic(&MarshalTypeError{v.Type()})
	}
}

type encodeField struct {
	i          int
	tag        string
	omit_empty bool
}

type encodeFieldsSortType []encodeField

func (ef encodeFieldsSortType) Len() int           { return len(ef) }
func (ef encodeFieldsSortType) Swap(i, j int)      { ef[i], ef[j] = ef[j], ef[i] }
func (ef encodeFieldsSortType) Less(i, j int) bool { return ef[i].tag < ef[j].tag }

var (
	typeCacheLock     sync.RWMutex
	encodeFieldsCache = make(map[reflect.Type][]encodeField)
)

func encodeFields(t reflect.Type) []encodeField {
	typeCacheLock.RLock()
	fs, ok := encodeFieldsCache[t]
	typeCacheLock.RUnlock()
	if ok {
		return fs
	}

	typeCacheLock.Lock()
	defer typeCacheLock.Unlock()
	fs, ok = encodeFieldsCache[t]
	if ok {
		return fs
	}

	for i, n := 0, t.NumField(); i < n; i++ {
		f := t.Field(i)
		if f.PkgPath != "" {
			continue
		}
		if f.Anonymous {
			continue
		}
		var ef encodeField
		ef.i = i
		ef.tag = f.Name

		tv := getTag(f.Tag)
		if tv.Ignore() {
			continue
		}
		if tv.Key() != "" {
			ef.tag = tv.Key()
		}
		ef.omit_empty = tv.OmitEmpty()
		fs = append(fs, ef)
	}
	fss := encodeFieldsSortType(fs)
	sort.Sort(fss)
	encodeFieldsCache[t] = fs
	return fs
}
<commit_msg>bencode: Avoid Value.Interface call testing for big.Int<commit_after>package bencode

import (
	"io"
	"math/big"
	"reflect"
	"runtime"
	"sort"
	"strconv"
	"sync"

	"github.com/anacrolix/missinggo"
)

func isEmptyValue(v reflect.Value) bool {
	return missinggo.IsEmptyValue(v)
}

type Encoder struct {
	w interface {
		Flush() error
		io.Writer
		WriteString(string) (int, error)
	}
	scratch [64]byte
}

func (e *Encoder) Encode(v interface{}) (err error) {
	if v == nil {
		return
	}
	defer func() {
		if e := recover(); e != nil {
			if _, ok := e.(runtime.Error); ok {
				panic(e)
			}
			var ok bool
			err, ok = e.(error)
			if !ok {
				panic(e)
			}
		}
	}()
	e.reflectValue(reflect.ValueOf(v))
	return e.w.Flush()
}

type string_values []reflect.Value

func (sv string_values) Len() int           { return len(sv) }
func (sv string_values) Swap(i, j int)      { sv[i], sv[j] = sv[j], sv[i] }
func (sv string_values) Less(i, j int) bool { return sv.get(i) < sv.get(j) }
func (sv string_values) get(i int) string   { return sv[i].String() }

func (e *Encoder) write(s []byte) {
	_, err := e.w.Write(s)
	if err != nil {
		panic(err)
	}
}

func (e *Encoder) writeString(s string) {
	_, err := e.w.WriteString(s)
	if err != nil {
		panic(err)
	}
}

func (e *Encoder) reflectString(s string) {
	b := strconv.AppendInt(e.scratch[:0], int64(len(s)), 10)
	e.write(b)
	e.writeString(":")
	e.writeString(s)
}

func (e *Encoder) reflectByteSlice(s []byte) {
	b := strconv.AppendInt(e.scratch[:0], int64(len(s)), 10)
	e.write(b)
	e.writeString(":")
	e.write(s)
}

// returns true if the value implements Marshaler interface and marshaling was
// done successfully
func (e *Encoder) reflectMarshaler(v reflect.Value) bool {
	m, ok := v.Interface().(Marshaler)
	if !ok {
		// T doesn't work, try *T
		if v.Kind() != reflect.Ptr && v.CanAddr() {
			m, ok = v.Addr().Interface().(Marshaler)
			if ok {
				v = v.Addr()
			}
		}
	}
	if ok && (v.Kind() != reflect.Ptr || !v.IsNil()) {
		data, err := m.MarshalBencode()
		if err != nil {
			panic(&MarshalerError{v.Type(), err})
		}
		e.write(data)
		return true
	}

	return false
}

var bigIntType = reflect.TypeOf(big.Int{})

func (e *Encoder) reflectValue(v reflect.Value) {

	if e.reflectMarshaler(v) {
		return
	}

	if v.Type() == bigIntType {
		e.writeString("i")
		bi := v.Interface().(big.Int)
		e.writeString(bi.String())
		e.writeString("e")
		return
	}

	switch v.Kind() {
	case reflect.Bool:
		if v.Bool() {
			e.writeString("i1e")
		} else {
			e.writeString("i0e")
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		b := strconv.AppendInt(e.scratch[:0], v.Int(), 10)
		e.writeString("i")
		e.write(b)
		e.writeString("e")
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		b := strconv.AppendUint(e.scratch[:0], v.Uint(), 10)
		e.writeString("i")
		e.write(b)
		e.writeString("e")
	case reflect.String:
		e.reflectString(v.String())
	case reflect.Struct:
		e.writeString("d")
		for _, ef := range encodeFields(v.Type()) {
			field_value := v.Field(ef.i)
			if ef.omit_empty && isEmptyValue(field_value) {
				continue
			}
			e.reflectString(ef.tag)
			e.reflectValue(field_value)
		}
		e.writeString("e")
	case reflect.Map:
		if v.Type().Key().Kind() != reflect.String {
			panic(&MarshalTypeError{v.Type()})
		}
		if v.IsNil() {
			e.writeString("de")
			break
		}
		e.writeString("d")
		sv := string_values(v.MapKeys())
		sort.Sort(sv)
		for _, key := range sv {
			e.reflectString(key.String())
			e.reflectValue(v.MapIndex(key))
		}
		e.writeString("e")
	case reflect.Slice:
		if v.IsNil() {
			e.writeString("le")
			break
		}
		if v.Type().Elem().Kind() == reflect.Uint8 {
			s := v.Bytes()
			e.reflectByteSlice(s)
			break
		}
		fallthrough
	case reflect.Array:
		e.writeString("l")
		for i, n := 0, v.Len(); i < n; i++ {
			e.reflectValue(v.Index(i))
		}
		e.writeString("e")
	case reflect.Interface:
		e.reflectValue(v.Elem())
	case reflect.Ptr:
		if v.IsNil() {
			v = reflect.Zero(v.Type().Elem())
		} else {
			v = v.Elem()
		}
		e.reflectValue(v)
	default:
		panic(&MarshalTypeError{v.Type()})
	}
}

type encodeField struct {
	i          int
	tag        string
	omit_empty bool
}

type encodeFieldsSortType []encodeField

func (ef encodeFieldsSortType) Len() int           { return len(ef) }
func (ef encodeFieldsSortType) Swap(i, j int)      { ef[i], ef[j] = ef[j], ef[i] }
func (ef encodeFieldsSortType) Less(i, j int) bool { return ef[i].tag < ef[j].tag }

var (
	typeCacheLock     sync.RWMutex
	encodeFieldsCache = make(map[reflect.Type][]encodeField)
)

func encodeFields(t reflect.Type) []encodeField {
	typeCacheLock.RLock()
	fs, ok := encodeFieldsCache[t]
	typeCacheLock.RUnlock()
	if ok {
		return fs
	}

	typeCacheLock.Lock()
	defer typeCacheLock.Unlock()
	fs, ok = encodeFieldsCache[t]
	if ok {
		return fs
	}

	for i, n := 0, t.NumField(); i < n; i++ {
		f := t.Field(i)
		if f.PkgPath != "" {
			continue
		}
		if f.Anonymous {
			continue
		}
		var ef encodeField
		ef.i = i
		ef.tag = f.Name

		tv := getTag(f.Tag)
		if tv.Ignore() {
			continue
		}
		if tv.Key() != "" {
			ef.tag = tv.Key()
		}
		ef.omit_empty = tv.OmitEmpty()
		fs = append(fs, ef)
	}
	fss := encodeFieldsSortType(fs)
	sort.Sort(fss)
	encodeFieldsCache[t] = fs
	return fs
}
<|endoftext|>


<commit_before>package main

const stmtCheckConfigList = `
SELECT configuration_id,
       repository_id,
       bucket_id,
	   configuration_name
FROM   soma.check_configurations
WHERE  repository_id = $1::uuid
AND    NOT deleted;`

const stmtCheckConfigShowBase = `
SELECT configuration_id,
       repository_id,
       bucket_id,
       configuration_name,
       configuration_object,
       configuration_object_type,
       configuration_active,
       inheritance_enabled,
       children_only,
       capability_id,
       interval,
       enabled,
       external_id
FROM   soma.check_configurations
WHERE  configuration_id = $1::uuid;`

const stmtCheckConfigShowThreshold = `
SELECT scc.configuration_id,
       sct.predicate,
	   sct.threshold,
	   sct.notification_level,
	   snl.level_shortname,
	   snl.level_numeric
FROM   soma.check_configurations scc
JOIN   soma.configuration_thresholds sct
ON     scc.configuration_id = sct.configuration_id
JOIN   soma.notification_levels snl
ON     sct.notification_level = snl.level_name
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrCustom = `
SELECT scc.configuration_id,
	   sccp.custom_property_id,
	   sccp.repository_id,
	   sccp.property_value,
	   scp.custom_property
FROM   soma.check_configurations scc
JOIN   soma.constraints_custom_property sccp
ON     scc.configuration_id = sccp.configuration_id
JOIN   soma.custom_properties scp
ON     sccp.custom_property_id = scp.custom_property_id
AND    sccp.repository_id = scp.repository_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrSystem = `
SELECT scc.configuration_id,
       scsp.system_property,
	   scsp.property_value
FROM   soma.check_configurations scc
JOIN   soma.constraints_system_property scsp
ON     scc.configuration_id = scsp.configuration_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrNative = `
SELECT scc.configuration_id,
       scnp.native_property,
	   scnp.property_value
FROM   soma.check_configurations scc
JOIN   soma.constraints_native_property scnp
ON     scc.configuration_id = scnp.configuration_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrService = `
SELECT scc.configuration_id,
       scsvp.organizational_team_id,
	   scsvp.service_property
FROM   soma.check_configurations scc
JOIN   soma.constraints_service_property scsvp
ON     scc.configuration_id = scsvp.configuration_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrAttribute = `
SELECT scc.configuration_id,
       scsa.service_property_attribute,
	   scsa.attribute_value
FROM   soma.check_configurations scc
JOIN   soma.constraints_service_attribute scsa
ON     scc.configuration_id = scsa.configuration_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrOncall = `
SELECT scc.configuration_id,
       scop.oncall_duty_id,
	   iodt.oncall_duty_name,
	   iodt.oncall_duty_phone_number
FROM   soma.check_configurations scc
JOIN   soma.constraints_oncall_property scop
ON     scc.configuration_id = scop.configuration_id
JOIN   inventory.oncall_duty_teams iodt
ON     scop.oncall_duty_id = iodt.oncall_duty_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigInstanceInfo = `
SELECT sci.check_instance_id,
       sc.object_id,
       sc.object_type,
       scic.status,
       scic.next_status
FROM   soma.check_configurations scc
JOIN   soma.check_instances sci
  ON   scc.configuration_id = sci.check_configuration_id
JOIN   soma.checks sc
  ON   sci.check_id = sc.check_id
JOIN   soma.check_instance_configurations scic
  ON   sci.current_instance_config_id = scic.check_instance_config_id
WHERE  scc.configuration_id = $1::uuid;`

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Do not display deleted check instances in details<commit_after>package main

const stmtCheckConfigList = `
SELECT configuration_id,
       repository_id,
       bucket_id,
	   configuration_name
FROM   soma.check_configurations
WHERE  repository_id = $1::uuid
AND    NOT deleted;`

const stmtCheckConfigShowBase = `
SELECT configuration_id,
       repository_id,
       bucket_id,
       configuration_name,
       configuration_object,
       configuration_object_type,
       configuration_active,
       inheritance_enabled,
       children_only,
       capability_id,
       interval,
       enabled,
       external_id
FROM   soma.check_configurations
WHERE  configuration_id = $1::uuid;`

const stmtCheckConfigShowThreshold = `
SELECT scc.configuration_id,
       sct.predicate,
	   sct.threshold,
	   sct.notification_level,
	   snl.level_shortname,
	   snl.level_numeric
FROM   soma.check_configurations scc
JOIN   soma.configuration_thresholds sct
ON     scc.configuration_id = sct.configuration_id
JOIN   soma.notification_levels snl
ON     sct.notification_level = snl.level_name
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrCustom = `
SELECT scc.configuration_id,
	   sccp.custom_property_id,
	   sccp.repository_id,
	   sccp.property_value,
	   scp.custom_property
FROM   soma.check_configurations scc
JOIN   soma.constraints_custom_property sccp
ON     scc.configuration_id = sccp.configuration_id
JOIN   soma.custom_properties scp
ON     sccp.custom_property_id = scp.custom_property_id
AND    sccp.repository_id = scp.repository_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrSystem = `
SELECT scc.configuration_id,
       scsp.system_property,
	   scsp.property_value
FROM   soma.check_configurations scc
JOIN   soma.constraints_system_property scsp
ON     scc.configuration_id = scsp.configuration_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrNative = `
SELECT scc.configuration_id,
       scnp.native_property,
	   scnp.property_value
FROM   soma.check_configurations scc
JOIN   soma.constraints_native_property scnp
ON     scc.configuration_id = scnp.configuration_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrService = `
SELECT scc.configuration_id,
       scsvp.organizational_team_id,
	   scsvp.service_property
FROM   soma.check_configurations scc
JOIN   soma.constraints_service_property scsvp
ON     scc.configuration_id = scsvp.configuration_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrAttribute = `
SELECT scc.configuration_id,
       scsa.service_property_attribute,
	   scsa.attribute_value
FROM   soma.check_configurations scc
JOIN   soma.constraints_service_attribute scsa
ON     scc.configuration_id = scsa.configuration_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigShowConstrOncall = `
SELECT scc.configuration_id,
       scop.oncall_duty_id,
	   iodt.oncall_duty_name,
	   iodt.oncall_duty_phone_number
FROM   soma.check_configurations scc
JOIN   soma.constraints_oncall_property scop
ON     scc.configuration_id = scop.configuration_id
JOIN   inventory.oncall_duty_teams iodt
ON     scop.oncall_duty_id = iodt.oncall_duty_id
WHERE  scc.configuration_id = $1::uuid;`

const stmtCheckConfigInstanceInfo = `
SELECT sci.check_instance_id,
       sc.object_id,
       sc.object_type,
       scic.status,
       scic.next_status
FROM   soma.check_configurations scc
JOIN   soma.check_instances sci
  ON   scc.configuration_id = sci.check_configuration_id
JOIN   soma.checks sc
  ON   sci.check_id = sc.check_id
JOIN   soma.check_instance_configurations scic
  ON   sci.current_instance_config_id = scic.check_instance_config_id
WHERE  scc.configuration_id = $1::uuid
  AND  scc.status != 'awaiting_deletion';`

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package main_test

import (
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-cf-experimental/switchboard/cmd/switchboard/fakes"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"github.com/tedsuo/ifrit/grouper"
)

func switchboardRunner(args ...string) ifrit.Runner {
	command := exec.Command(switchboardBinPath, args...)
	runner := ginkgomon.New(ginkgomon.Config{
		Command:    command,
		Name:       fmt.Sprintf("switchboard"),
		StartCheck: "started",
	})
	return runner
}

type Response struct {
	BackendPort     uint
	HealthcheckPort uint
	Message         string
}

func sendData(conn net.Conn, data string) (Response, error) {
	conn.Write([]byte(data))
	buffer := make([]byte, 1024)
	n, err := conn.Read(buffer)
	if err != nil {
		return Response{}, err
	} else {
		response := Response{}
		err := json.Unmarshal(buffer[:n], &response)
		if err != nil {
			return Response{}, err
		}
		return response, nil
	}
}

var _ = Describe("Switchboard", func() {
	var process ifrit.Process
	var initialActiveHealthcheckPort uint
	var initialInactiveBackendPort uint
	var healthcheckRunner1, healthcheckRunner2 *fakes.FakeHealthcheck

	BeforeEach(func() {
		healthcheckRunner1 = fakes.NewFakeHealthcheck(dummyHealthcheckPort)
		healthcheckRunner2 = fakes.NewFakeHealthcheck(dummyHealthcheckPort2)

		group := grouper.NewParallel(os.Kill, grouper.Members{
			grouper.Member{"backend-1", fakes.NewFakeBackend(backendPort, dummyHealthcheckPort)},
			grouper.Member{"backend-2", fakes.NewFakeBackend(backendPort2, dummyHealthcheckPort2)},
			grouper.Member{"healthcheck-1", healthcheckRunner1},
			grouper.Member{"healthcheck-2", healthcheckRunner2},
			grouper.Member{"switchboard", switchboardRunner(
				fmt.Sprintf("-config=%s", proxyConfigFile),
				fmt.Sprintf("-pidFile=%s", pidFile),
			)},
		})
		process = ifrit.Invoke(group)

		var err error
		var conn net.Conn
		Eventually(func() error {
			conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
			return err
		}).ShouldNot(HaveOccurred())

		response, err := sendData(conn, "detect active")
		Expect(err).NotTo(HaveOccurred())

		initialActiveHealthcheckPort = response.HealthcheckPort
		if response.BackendPort == backendPort {
			initialInactiveBackendPort = backendPort2
		} else {
			initialInactiveBackendPort = backendPort
		}
	})

	AfterEach(func() {
		ginkgomon.Kill(process)
	})

	Context("when there are multiple concurrent clients", func() {
		var conn1, conn2, conn3 net.Conn
		var data1, data2, data3 Response

		It("proxies all the connections to the backend", func() {
			done1 := make(chan interface{})
			go func() {
				defer GinkgoRecover()
				defer close(done1)

				var err error
				Eventually(func() error {
					conn1, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data1, err = sendData(conn1, "test1")
				Expect(err).ToNot(HaveOccurred())
			}()

			done2 := make(chan interface{})
			go func() {
				defer GinkgoRecover()
				defer close(done2)

				var err error
				Eventually(func() error {
					conn2, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data2, err = sendData(conn2, "test2")
				Expect(err).ToNot(HaveOccurred())
			}()

			done3 := make(chan interface{})
			go func() {
				defer GinkgoRecover()
				defer close(done3)

				var err error
				Eventually(func() error {
					conn3, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data3, err = sendData(conn3, "test3")
				Expect(err).ToNot(HaveOccurred())
			}()

			<-done1
			<-done2
			<-done3

			Expect(data1.Message).Should(Equal("test1"))
			Expect(data2.Message).Should(Equal("test2"))
			Expect(data3.Message).Should(Equal("test3"))
		})
	})

	Context("when other clients disconnect", func() {
		var conn net.Conn
		var connToDisconnect net.Conn

		It("maintains a long-lived connection when other clients disconnect", func() {
			Eventually(func() error {
				var err error
				conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
				return err
			}).ShouldNot(HaveOccurred())

			Eventually(func() error {
				var err error
				connToDisconnect, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
				return err
			}).ShouldNot(HaveOccurred())

			dataBeforeDisconnect, err := sendData(conn, "data before disconnect")
			Expect(err).ToNot(HaveOccurred())
			Expect(dataBeforeDisconnect.Message).Should(Equal("data before disconnect"))

			connToDisconnect.Close()

			dataAfterDisconnect, err := sendData(conn, "data after disconnect")
			Expect(err).ToNot(HaveOccurred())
			Expect(dataAfterDisconnect.Message).Should(Equal("data after disconnect"))
		})
	})

	Context("when the healthcheck succeeds", func() {
		It("checks health again after the specified interval", func() {
			var client net.Conn
			Eventually(func() error {
				var err error
				client, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
				return err
			}).ShouldNot(HaveOccurred())

			data, err := sendData(client, "data around first healthcheck")
			Expect(err).NotTo(HaveOccurred())
			Expect(data.Message).To(Equal("data around first healthcheck"))

			Consistently(func() error {
				_, err = sendData(client, "data around subsequent healthcheck")
				return err
			}, 3*time.Second).ShouldNot(HaveOccurred())
		})
	})

	Context("when the cluster is down", func() {
		Context("when the healthcheck reports a 503", func() {
			It("disconnects client connections", func() {
				var conn net.Conn
				Eventually(func() error {
					var err error
					conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				dataWhileHealthy, err := sendData(conn, "data while healthy")
				Expect(err).ToNot(HaveOccurred())
				Expect(dataWhileHealthy.Message).To(Equal("data while healthy"))

				if initialActiveHealthcheckPort == dummyHealthcheckPort {
					healthcheckRunner1.SetStatusCode(http.StatusServiceUnavailable)
				} else {
					healthcheckRunner2.SetStatusCode(http.StatusServiceUnavailable)
				}

				Eventually(func() error {
					_, err := sendData(conn, "data when unhealthy")
					return err
				}, 2*time.Second).Should(HaveOccurred())
			})
		})

		Context("when a backend goes down", func() {
			var conn net.Conn
			var data Response

			BeforeEach(func() {
				Eventually(func() (err error) {
					conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data, err := sendData(conn, "data before hang")
				Expect(err).ShouldNot(HaveOccurred())
				Expect(data.Message).Should(Equal("data before hang"))

				if initialActiveHealthcheckPort == dummyHealthcheckPort {
					healthcheckRunner1.SetHang(true)
				} else {
					healthcheckRunner2.SetHang(true)
				}
			})

			It("disconnects existing client connections", func(done Done) {
				defer close(done)

				Eventually(func() error {
					_, err := sendData(conn, "data after hang")
					return err
				}, proxyConfig.HealthcheckTimeout()*10).Should(HaveOccurred())
			}, 5)

			It("proxies new connections to another backend", func(done Done) {
				defer close(done)

				time.Sleep(3 * proxyConfig.HealthcheckTimeout()) // wait for failover

				var err error
				Eventually(func() error {
					conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data, err = sendData(conn, "test")
				Expect(err).ToNot(HaveOccurred())
				Expect(data.Message).To(Equal("test"))
				Expect(data.BackendPort).To(Equal(initialInactiveBackendPort))
			}, 5)
		})

		Context("when all backends are down", func() {
			BeforeEach(func() {
				healthcheckRunner1.SetHang(true)
				healthcheckRunner2.SetHang(true)
			})

			It("rejects any new connections that are attempted", func(done Done) {
				defer close(done)

				time.Sleep(3 * proxyConfig.HealthcheckTimeout()) // wait for failover

				var conn net.Conn
				Eventually(func() (err error) {
					conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}, 1*time.Second).ShouldNot(HaveOccurred())

				Eventually(func() error {
					_, err := sendData(conn, "write that should fail")
					return err
				}, proxyConfig.HealthcheckTimeout()*4, 200*time.Millisecond).Should(HaveOccurred())

			}, 20)
		})
	})
})
<commit_msg>Inline creation of switchboard runner<commit_after>package main_test

import (
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/pivotal-cf-experimental/switchboard/cmd/switchboard/fakes"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"github.com/tedsuo/ifrit/grouper"
)

type Response struct {
	BackendPort     uint
	HealthcheckPort uint
	Message         string
}

func sendData(conn net.Conn, data string) (Response, error) {
	conn.Write([]byte(data))
	buffer := make([]byte, 1024)
	n, err := conn.Read(buffer)
	if err != nil {
		return Response{}, err
	} else {
		response := Response{}
		err := json.Unmarshal(buffer[:n], &response)
		if err != nil {
			return Response{}, err
		}
		return response, nil
	}
}

var _ = Describe("Switchboard", func() {
	var process ifrit.Process
	var initialActiveHealthcheckPort uint
	var initialInactiveBackendPort uint
	var healthcheckRunner1, healthcheckRunner2 *fakes.FakeHealthcheck

	BeforeEach(func() {
		backendRunner1 := fakes.NewFakeBackend(backendPort, dummyHealthcheckPort)
		backendRunner2 := fakes.NewFakeBackend(backendPort2, dummyHealthcheckPort2)
		healthcheckRunner1 = fakes.NewFakeHealthcheck(dummyHealthcheckPort)
		healthcheckRunner2 = fakes.NewFakeHealthcheck(dummyHealthcheckPort2)

		switchboardRunner := ginkgomon.New(ginkgomon.Config{
			Command: exec.Command(
				switchboardBinPath,
				fmt.Sprintf("-config=%s", proxyConfigFile),
				fmt.Sprintf("-pidFile=%s", pidFile),
			),
			Name:       fmt.Sprintf("switchboard"),
			StartCheck: "started",
		})

		group := grouper.NewParallel(os.Kill, grouper.Members{
			grouper.Member{"backend-1", backendRunner1},
			grouper.Member{"backend-2", backendRunner2},
			grouper.Member{"healthcheck-1", healthcheckRunner1},
			grouper.Member{"healthcheck-2", healthcheckRunner2},
			grouper.Member{"switchboard", switchboardRunner},
		})
		process = ifrit.Invoke(group)

		var err error
		var conn net.Conn
		Eventually(func() error {
			conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
			return err
		}).ShouldNot(HaveOccurred())

		response, err := sendData(conn, "detect active")
		Expect(err).NotTo(HaveOccurred())

		initialActiveHealthcheckPort = response.HealthcheckPort
		if response.BackendPort == backendPort {
			initialInactiveBackendPort = backendPort2
		} else {
			initialInactiveBackendPort = backendPort
		}
	})

	AfterEach(func() {
		ginkgomon.Kill(process)
	})

	Context("when there are multiple concurrent clients", func() {
		var conn1, conn2, conn3 net.Conn
		var data1, data2, data3 Response

		It("proxies all the connections to the backend", func() {
			done1 := make(chan interface{})
			go func() {
				defer GinkgoRecover()
				defer close(done1)

				var err error
				Eventually(func() error {
					conn1, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data1, err = sendData(conn1, "test1")
				Expect(err).ToNot(HaveOccurred())
			}()

			done2 := make(chan interface{})
			go func() {
				defer GinkgoRecover()
				defer close(done2)

				var err error
				Eventually(func() error {
					conn2, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data2, err = sendData(conn2, "test2")
				Expect(err).ToNot(HaveOccurred())
			}()

			done3 := make(chan interface{})
			go func() {
				defer GinkgoRecover()
				defer close(done3)

				var err error
				Eventually(func() error {
					conn3, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data3, err = sendData(conn3, "test3")
				Expect(err).ToNot(HaveOccurred())
			}()

			<-done1
			<-done2
			<-done3

			Expect(data1.Message).Should(Equal("test1"))
			Expect(data2.Message).Should(Equal("test2"))
			Expect(data3.Message).Should(Equal("test3"))
		})
	})

	Context("when other clients disconnect", func() {
		var conn net.Conn
		var connToDisconnect net.Conn

		It("maintains a long-lived connection when other clients disconnect", func() {
			Eventually(func() error {
				var err error
				conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
				return err
			}).ShouldNot(HaveOccurred())

			Eventually(func() error {
				var err error
				connToDisconnect, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
				return err
			}).ShouldNot(HaveOccurred())

			dataBeforeDisconnect, err := sendData(conn, "data before disconnect")
			Expect(err).ToNot(HaveOccurred())
			Expect(dataBeforeDisconnect.Message).Should(Equal("data before disconnect"))

			connToDisconnect.Close()

			dataAfterDisconnect, err := sendData(conn, "data after disconnect")
			Expect(err).ToNot(HaveOccurred())
			Expect(dataAfterDisconnect.Message).Should(Equal("data after disconnect"))
		})
	})

	Context("when the healthcheck succeeds", func() {
		It("checks health again after the specified interval", func() {
			var client net.Conn
			Eventually(func() error {
				var err error
				client, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
				return err
			}).ShouldNot(HaveOccurred())

			data, err := sendData(client, "data around first healthcheck")
			Expect(err).NotTo(HaveOccurred())
			Expect(data.Message).To(Equal("data around first healthcheck"))

			Consistently(func() error {
				_, err = sendData(client, "data around subsequent healthcheck")
				return err
			}, 3*time.Second).ShouldNot(HaveOccurred())
		})
	})

	Context("when the cluster is down", func() {
		Context("when the healthcheck reports a 503", func() {
			It("disconnects client connections", func() {
				var conn net.Conn
				Eventually(func() error {
					var err error
					conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				dataWhileHealthy, err := sendData(conn, "data while healthy")
				Expect(err).ToNot(HaveOccurred())
				Expect(dataWhileHealthy.Message).To(Equal("data while healthy"))

				if initialActiveHealthcheckPort == dummyHealthcheckPort {
					healthcheckRunner1.SetStatusCode(http.StatusServiceUnavailable)
				} else {
					healthcheckRunner2.SetStatusCode(http.StatusServiceUnavailable)
				}

				Eventually(func() error {
					_, err := sendData(conn, "data when unhealthy")
					return err
				}, 2*time.Second).Should(HaveOccurred())
			})
		})

		Context("when a backend goes down", func() {
			var conn net.Conn
			var data Response

			BeforeEach(func() {
				Eventually(func() (err error) {
					conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data, err := sendData(conn, "data before hang")
				Expect(err).ShouldNot(HaveOccurred())
				Expect(data.Message).Should(Equal("data before hang"))

				if initialActiveHealthcheckPort == dummyHealthcheckPort {
					healthcheckRunner1.SetHang(true)
				} else {
					healthcheckRunner2.SetHang(true)
				}
			})

			It("disconnects existing client connections", func(done Done) {
				defer close(done)

				Eventually(func() error {
					_, err := sendData(conn, "data after hang")
					return err
				}, proxyConfig.HealthcheckTimeout()*10).Should(HaveOccurred())
			}, 5)

			It("proxies new connections to another backend", func(done Done) {
				defer close(done)

				time.Sleep(3 * proxyConfig.HealthcheckTimeout()) // wait for failover

				var err error
				Eventually(func() error {
					conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}).ShouldNot(HaveOccurred())

				data, err = sendData(conn, "test")
				Expect(err).ToNot(HaveOccurred())
				Expect(data.Message).To(Equal("test"))
				Expect(data.BackendPort).To(Equal(initialInactiveBackendPort))
			}, 5)
		})

		Context("when all backends are down", func() {
			BeforeEach(func() {
				healthcheckRunner1.SetHang(true)
				healthcheckRunner2.SetHang(true)
			})

			It("rejects any new connections that are attempted", func(done Done) {
				defer close(done)

				time.Sleep(3 * proxyConfig.HealthcheckTimeout()) // wait for failover

				var conn net.Conn
				Eventually(func() (err error) {
					conn, err = net.Dial("tcp", fmt.Sprintf("localhost:%d", switchboardPort))
					return err
				}, 1*time.Second).ShouldNot(HaveOccurred())

				Eventually(func() error {
					_, err := sendData(conn, "write that should fail")
					return err
				}, proxyConfig.HealthcheckTimeout()*4, 200*time.Millisecond).Should(HaveOccurred())

			}, 20)
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2017-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"time"

	"github.com/u-root/u-root/pkg/booter"
	"github.com/u-root/u-root/pkg/ipmi"
	"github.com/u-root/u-root/pkg/smbios"
)

var (
	allowInteractive = flag.Bool("i", true, "Allow user to interrupt boot process and run commands")
	doQuiet          = flag.Bool("q", false, "Disable verbose output")
	interval         = flag.Int("I", 1, "Interval in seconds before looping to the next boot command")
	noDefaultBoot    = flag.Bool("nodefault", false, "Do not attempt default boot entries if regular ones fail")
)

var defaultBootsequence = [][]string{
	{"fbnetboot", "-userclass", "linuxboot"},
	{"localboot", "-grub"},
}

// Product list for running IPMI OEM commands
var productList = [2]string{"Tioga Pass", "Mono Lake"}

var selRecorded bool

func isMatched(productName string) bool {
	for _, v := range productList {
		if strings.HasPrefix(productName, v) {
			return true
		}
	}
	return false
}

func getSystemProductName(si *smbios.Info) (string, error) {
	t1, err := si.GetSystemInfo()
	if err != nil {
		log.Printf("Error getting System Information: %v", err)
		return "", err
	}
	return t1.ProductName, nil
}

func getSystemFWVersion(si *smbios.Info) (string, error) {
	t0, err := si.GetBIOSInfo()
	if err != nil {
		log.Printf("Error getting BIOS Information: %v", err)
		return "", err
	}
	return t0.Version, nil
}

func checkCMOSClear(ipmi *ipmi.IPMI) error {
	if cmosclear, bootorder, err := ipmi.IsCMOSClearSet(); cmosclear {
		log.Printf("CMOS clear starts")
		if err = cmosClear(); err != nil {
			return err
		}
		// ToDo: Reset RW_VPD to default values
		if err = ipmi.ClearCMOSClearValidBits(bootorder); err != nil {
			return err
		}
		if err = reboot(); err != nil {
			return err
		}
	} else if err != nil {
		return err
	}

	return nil
}

func runIPMICommands() {
	ipmi, err := ipmi.Open(0)
	if err != nil {
		log.Printf("Failed to open ipmi device %v, watchdog may still be running", err)
		return
	}
	defer ipmi.Close()

	if err = ipmi.ShutoffWatchdog(); err != nil {
		log.Printf("Failed to stop watchdog %v.", err)
	} else {
		log.Printf("Watchdog is stopped.")
	}

	// Below IPMI commands would require SMBIOS data
	si, err := smbios.FromSysfs()
	if err != nil {
		log.Printf("Error reading SMBIOS info: %v", err)
		return
	}

	if fwVersion, err := getSystemFWVersion(si); err == nil {
		log.Printf("System firmware version: %s", fwVersion)
		if err = ipmi.SetSystemFWVersion(fwVersion); err != nil {
			log.Printf("Failed to set system firmware version to BMC %v.", err)
		}
	}

	if productName, err := getSystemProductName(si); err == nil {
		if isMatched(productName) {
			log.Printf("Running OEM IPMI commands.")
			if err = checkCMOSClear(ipmi); err != nil {
				log.Printf("IPMI CMOS clear err: %v", err)
			}
		} else {
			log.Printf("No product name is matched for OEM commands.")
		}
	}

}

func sendBootSEL(sequence string) {
	var bootErr ipmi.Event

	i, err := ipmi.Open(0)
	if err != nil {
		log.Printf("Failed to open ipmi device to send SEL %v", err)
		return
	}
	defer i.Close()

	switch sequence {
	case "fbnetboot":
		bootErr.RecordID = 0x0000
		bootErr.RecordType = ipmi.OEM_NTS_TYPE
		bootErr.OEMNontsDefinedData = [13]uint8{}

		if err := i.LogSystemEvent(&bootErr); err != nil {
			log.Printf("SEL recorded: %s fail\n", sequence)
		}
	default:
	}
}

func main() {
	flag.Parse()

	log.Print(`
                     ____            _                 _                 _
                    / ___| _   _ ___| |_ ___ _ __ ___ | |__   ___   ___ | |_
                    \___ \| | | / __| __/ _ \ '_ ` + "`" + ` _ \| '_ \ / _ \ / _ \| __|
                     ___) | |_| \__ \ ||  __/ | | | | | |_) | (_) | (_) | |_
                    |____/ \__, |___/\__\___|_| |_| |_|_.__/ \___/ \___/ \__|
                           |___/
`)
	runIPMICommands()
	sleepInterval := time.Duration(*interval) * time.Second
	if *allowInteractive {
		log.Printf("**************************************************************************")
		log.Print("Starting boot sequence, press CTRL-C within 5 seconds to drop into a shell")
		log.Printf("**************************************************************************")
		time.Sleep(5 * time.Second)
	} else {
		signal.Ignore()
	}

	// Get and show boot entries
	bootEntries := booter.GetBootEntries()
	log.Printf("BOOT ENTRIES:")
	for _, entry := range bootEntries {
		log.Printf("    %v) %+v", entry.Name, string(entry.Config))
	}
	for _, entry := range bootEntries {
		log.Printf("Trying boot entry %s: %s", entry.Name, string(entry.Config))
		if err := entry.Booter.Boot(); err != nil {
			log.Printf("Warning: failed to boot with configuration: %+v", entry)
			sendBootSEL(entry.Booter.TypeName())
		}
		if !*doQuiet {
			log.Printf("Sleeping %v before attempting next boot command", sleepInterval)
		}
		time.Sleep(sleepInterval)
	}

	// if boot entries failed, use the default boot sequence
	log.Printf("Boot entries failed")

	if !*noDefaultBoot {
		log.Print("Falling back to the default boot sequence")
		for {
			for _, bootcmd := range defaultBootsequence {
				if !*doQuiet {
					bootcmd = append(bootcmd, "-d")
				}
				log.Printf("Running boot command: %v", bootcmd)
				cmd := exec.Command(bootcmd[0], bootcmd[1:]...)
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
				if err := cmd.Run(); err != nil {
					log.Printf("Error executing %v: %v", cmd, err)
					if !selRecorded {
						sendBootSEL(bootcmd[0])
					}
				}
			}
			selRecorded = true

			if !*doQuiet {
				log.Printf("Sleeping %v before attempting next boot command", sleepInterval)
			}
			time.Sleep(sleepInterval)
		}
	}
}
<commit_msg>systemboot: Add OEM CMOS clear SEL and update fbnetboot SEL<commit_after>// Copyright 2017-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"time"

	"github.com/u-root/u-root/pkg/booter"
	"github.com/u-root/u-root/pkg/ipmi"
	"github.com/u-root/u-root/pkg/smbios"
)

var (
	allowInteractive = flag.Bool("i", true, "Allow user to interrupt boot process and run commands")
	doQuiet          = flag.Bool("q", false, "Disable verbose output")
	interval         = flag.Int("I", 1, "Interval in seconds before looping to the next boot command")
	noDefaultBoot    = flag.Bool("nodefault", false, "Do not attempt default boot entries if regular ones fail")
)

var defaultBootsequence = [][]string{
	{"fbnetboot", "-userclass", "linuxboot"},
	{"localboot", "-grub"},
}

// Product list for running IPMI OEM commands
var productList = [2]string{"Tioga Pass", "Mono Lake"}

var selRecorded bool

func isMatched(productName string) bool {
	for _, v := range productList {
		if strings.HasPrefix(productName, v) {
			return true
		}
	}
	return false
}

func getSystemProductName(si *smbios.Info) (string, error) {
	t1, err := si.GetSystemInfo()
	if err != nil {
		log.Printf("Error getting System Information: %v", err)
		return "", err
	}
	return t1.ProductName, nil
}

func getSystemFWVersion(si *smbios.Info) (string, error) {
	t0, err := si.GetBIOSInfo()
	if err != nil {
		log.Printf("Error getting BIOS Information: %v", err)
		return "", err
	}
	return t0.Version, nil
}

func checkCMOSClear(ipmi *ipmi.IPMI) error {
	if cmosclear, bootorder, err := ipmi.IsCMOSClearSet(); cmosclear {
		log.Printf("CMOS clear starts")
		if err = cmosClear(); err != nil {
			return err
		}
		// ToDo: Reset RW_VPD to default values
		if err = ipmi.ClearCMOSClearValidBits(bootorder); err != nil {
			return err
		}
		addSEL("cmosclear")
		if err = reboot(); err != nil {
			return err
		}
	} else if err != nil {
		return err
	}

	return nil
}

func runIPMICommands() {
	ipmi, err := ipmi.Open(0)
	if err != nil {
		log.Printf("Failed to open ipmi device %v, watchdog may still be running", err)
		return
	}
	defer ipmi.Close()

	if err = ipmi.ShutoffWatchdog(); err != nil {
		log.Printf("Failed to stop watchdog %v.", err)
	} else {
		log.Printf("Watchdog is stopped.")
	}

	// Below IPMI commands would require SMBIOS data
	si, err := smbios.FromSysfs()
	if err != nil {
		log.Printf("Error reading SMBIOS info: %v", err)
		return
	}

	if fwVersion, err := getSystemFWVersion(si); err == nil {
		log.Printf("System firmware version: %s", fwVersion)
		if err = ipmi.SetSystemFWVersion(fwVersion); err != nil {
			log.Printf("Failed to set system firmware version to BMC %v.", err)
		}
	}

	if productName, err := getSystemProductName(si); err == nil {
		if isMatched(productName) {
			log.Printf("Running OEM IPMI commands.")
			if err = checkCMOSClear(ipmi); err != nil {
				log.Printf("IPMI CMOS clear err: %v", err)
			}
		} else {
			log.Printf("No product name is matched for OEM commands.")
		}
	}

}

func addSEL(sequence string) {
	var bootErr ipmi.Event

	i, err := ipmi.Open(0)
	if err != nil {
		log.Printf("Failed to open ipmi device to send SEL %v", err)
		return
	}
	defer i.Close()

	switch sequence {
	case "fbnetboot":
		bootErr.RecordID = 0
		bootErr.RecordType = ipmi.OEM_NTS_TYPE
		bootErr.OEMNontsDefinedData[0] = 0x28
		bootErr.OEMNontsDefinedData[5] = 0xf0
		for idx := 6; idx < 13; idx++ {
			bootErr.OEMNontsDefinedData[idx] = 0xff
		}
		if err := i.LogSystemEvent(&bootErr); err != nil {
			log.Printf("SEL recorded: %s fail\n", sequence)
		}
	case "cmosclear":
		bootErr.RecordID = 0
		bootErr.RecordType = ipmi.OEM_NTS_TYPE
		bootErr.OEMNontsDefinedData[0] = 0x28
		bootErr.OEMNontsDefinedData[5] = 0xf1
		for idx := 6; idx < 13; idx++ {
			bootErr.OEMNontsDefinedData[idx] = 0xff
		}
		if err := i.LogSystemEvent(&bootErr); err != nil {
			log.Printf("SEL recorded: %s fail\n", sequence)
		}
	default:
	}
}

func main() {
	flag.Parse()

	log.Print(`
                     ____            _                 _                 _
                    / ___| _   _ ___| |_ ___ _ __ ___ | |__   ___   ___ | |_
                    \___ \| | | / __| __/ _ \ '_ ` + "`" + ` _ \| '_ \ / _ \ / _ \| __|
                     ___) | |_| \__ \ ||  __/ | | | | | |_) | (_) | (_) | |_
                    |____/ \__, |___/\__\___|_| |_| |_|_.__/ \___/ \___/ \__|
                           |___/
`)
	runIPMICommands()
	sleepInterval := time.Duration(*interval) * time.Second
	if *allowInteractive {
		log.Printf("**************************************************************************")
		log.Print("Starting boot sequence, press CTRL-C within 5 seconds to drop into a shell")
		log.Printf("**************************************************************************")
		time.Sleep(5 * time.Second)
	} else {
		signal.Ignore()
	}

	// Get and show boot entries
	bootEntries := booter.GetBootEntries()
	log.Printf("BOOT ENTRIES:")
	for _, entry := range bootEntries {
		log.Printf("    %v) %+v", entry.Name, string(entry.Config))
	}
	for _, entry := range bootEntries {
		log.Printf("Trying boot entry %s: %s", entry.Name, string(entry.Config))
		if err := entry.Booter.Boot(); err != nil {
			log.Printf("Warning: failed to boot with configuration: %+v", entry)
			addSEL(entry.Booter.TypeName())
		}
		if !*doQuiet {
			log.Printf("Sleeping %v before attempting next boot command", sleepInterval)
		}
		time.Sleep(sleepInterval)
	}

	// if boot entries failed, use the default boot sequence
	log.Printf("Boot entries failed")

	if !*noDefaultBoot {
		log.Print("Falling back to the default boot sequence")
		for {
			for _, bootcmd := range defaultBootsequence {
				if !*doQuiet {
					bootcmd = append(bootcmd, "-d")
				}
				log.Printf("Running boot command: %v", bootcmd)
				cmd := exec.Command(bootcmd[0], bootcmd[1:]...)
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr
				if err := cmd.Run(); err != nil {
					log.Printf("Error executing %v: %v", cmd, err)
					if !selRecorded {
						addSEL(bootcmd[0])
					}
				}
			}
			selRecorded = true

			if !*doQuiet {
				log.Printf("Sleeping %v before attempting next boot command", sleepInterval)
			}
			time.Sleep(sleepInterval)
		}
	}
}
<|endoftext|>


<commit_before>// returns data points from Win32_OperatingSystem
// https://msdn.microsoft.com/en-us/library/aa394239 - Win32_OperatingSystem class

package collector

import (
	"log"

	"github.com/StackExchange/wmi"
	"github.com/prometheus/client_golang/prometheus"
)

func init() {
	Factories["os"] = NewOSCollector
}

// A OSCollector is a Prometheus collector for WMI metrics
type OSCollector struct {
	PhysicalMemoryFreeBytes *prometheus.Desc
	PagingFreeBytes         *prometheus.Desc
	VirtualMemoryFreeBytes  *prometheus.Desc
	ProcessesLimit          *prometheus.Desc
	ProcessMemoryLimitBytes *prometheus.Desc
	Processes               *prometheus.Desc
	Users                   *prometheus.Desc
	PagingLimitBytes        *prometheus.Desc
	VirtualMemoryBytes      *prometheus.Desc
	VisibleMemoryBytes      *prometheus.Desc
}

// NewOSCollector ...
func NewOSCollector() (Collector, error) {
	const subsystem = "os"

	return &OSCollector{
		PagingLimitBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "paging_limit_bytes"),
			"OperatingSystem.SizeStoredInPagingFiles",
			nil,
			nil,
		),
		PagingFreeBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "paging_free_bytes"),
			"OperatingSystem.FreeSpaceInPagingFiles",
			nil,
			nil,
		),
		PhysicalMemoryFreeBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "physical_memory_free_bytes"),
			"OperatingSystem.FreePhysicalMemory",
			nil,
			nil,
		),
		Processes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "processes"),
			"OperatingSystem.NumberOfProcesses",
			nil,
			nil,
		),
		ProcessesLimit: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "processes_limit"),
			"OperatingSystem.MaxNumberOfProcesses",
			nil,
			nil,
		),
		ProcessMemoryLimitBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "process_memory_limix_bytes"),
			"OperatingSystem.MaxProcessMemorySize",
			nil,
			nil,
		),
		Users: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "users"),
			"OperatingSystem.NumberOfUsers",
			nil,
			nil,
		),
		VirtualMemoryBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "virtual_memory_bytes"),
			"OperatingSystem.TotalVirtualMemorySize",
			nil,
			nil,
		),
		VisibleMemoryBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "visible_memory_bytes"),
			"OperatingSystem.TotalVisibleMemorySize",
			nil,
			nil,
		),
		VirtualMemoryFreeBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "virtual_memory_free_bytes"),
			"OperatingSystem.FreeVirtualMemory",
			nil,
			nil,
		),
	}, nil
}

// Collect sends the metric values for each metric
// to the provided prometheus Metric channel.
func (c *OSCollector) Collect(ch chan<- prometheus.Metric) error {
	if desc, err := c.collect(ch); err != nil {
		log.Println("[ERROR] failed collecting os metrics:", desc, err)
		return err
	}
	return nil
}

type Win32_OperatingSystem struct {
	FreePhysicalMemory      uint64
	FreeSpaceInPagingFiles  uint64
	FreeVirtualMemory       uint64
	MaxNumberOfProcesses    uint32
	MaxProcessMemorySize    uint64
	NumberOfProcesses       uint32
	NumberOfUsers           uint32
	SizeStoredInPagingFiles uint64
	TotalVirtualMemorySize  uint64
	TotalVisibleMemorySize  uint64
}

func (c *OSCollector) collect(ch chan<- prometheus.Metric) (*prometheus.Desc, error) {
	var dst []Win32_OperatingSystem
	if err := wmi.Query(wmi.CreateQuery(&dst, ""), &dst); err != nil {
		return nil, err
	}

	ch <- prometheus.MustNewConstMetric(
		c.PhysicalMemoryFreeBytes,
		prometheus.GaugeValue,
		float64(dst[0].FreePhysicalMemory*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.PagingFreeBytes,
		prometheus.GaugeValue,
		float64(dst[0].FreeSpaceInPagingFiles*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.VirtualMemoryFreeBytes,
		prometheus.GaugeValue,
		float64(dst[0].FreeVirtualMemory*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.ProcessesLimit,
		prometheus.GaugeValue,
		float64(dst[0].MaxNumberOfProcesses),
	)

	ch <- prometheus.MustNewConstMetric(
		c.ProcessMemoryLimitBytes,
		prometheus.GaugeValue,
		float64(dst[0].MaxProcessMemorySize*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.Processes,
		prometheus.GaugeValue,
		float64(dst[0].NumberOfProcesses),
	)

	ch <- prometheus.MustNewConstMetric(
		c.Users,
		prometheus.GaugeValue,
		float64(dst[0].NumberOfUsers),
	)

	ch <- prometheus.MustNewConstMetric(
		c.PagingLimitBytes,
		prometheus.GaugeValue,
		float64(dst[0].SizeStoredInPagingFiles*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.VirtualMemoryBytes,
		prometheus.GaugeValue,
		float64(dst[0].TotalVirtualMemorySize*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.VisibleMemoryBytes,
		prometheus.GaugeValue,
		float64(dst[0].TotalVisibleMemorySize*1024), // KiB -> bytes
	)

	return nil, nil
}
<commit_msg> Added metrics to OS collector<commit_after>// returns data points from Win32_OperatingSystem
// https://msdn.microsoft.com/en-us/library/aa394239 - Win32_OperatingSystem class

package collector

import (
	"log"
	"time"

	"github.com/StackExchange/wmi"
	"github.com/prometheus/client_golang/prometheus"
)

func init() {
	Factories["os"] = NewOSCollector
}

// A OSCollector is a Prometheus collector for WMI metrics
type OSCollector struct {
	PhysicalMemoryFreeBytes *prometheus.Desc
	PagingFreeBytes         *prometheus.Desc
	VirtualMemoryFreeBytes  *prometheus.Desc
	ProcessesLimit          *prometheus.Desc
	ProcessMemoryLimitBytes *prometheus.Desc
	Processes               *prometheus.Desc
	Users                   *prometheus.Desc
	PagingLimitBytes        *prometheus.Desc
	VirtualMemoryBytes      *prometheus.Desc
	VisibleMemoryBytes      *prometheus.Desc
	Time			*prometheus.Desc
	Timezone		*prometheus.Desc
}

// NewOSCollector ...
func NewOSCollector() (Collector, error) {
	const subsystem = "os"

	return &OSCollector{
		PagingLimitBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "paging_limit_bytes"),
			"OperatingSystem.SizeStoredInPagingFiles",
			nil,
			nil,
		),
		PagingFreeBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "paging_free_bytes"),
			"OperatingSystem.FreeSpaceInPagingFiles",
			nil,
			nil,
		),
		PhysicalMemoryFreeBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "physical_memory_free_bytes"),
			"OperatingSystem.FreePhysicalMemory",
			nil,
			nil,
		),
		Time: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "time"),
			"OperatingSystem.LocalDateTime",
			nil,
			nil,
		),
		Timezone: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "timezone"),
			"OperatingSystem.LocalDateTime",
			[]string{"timezone"},
			nil,
		),
		Processes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "processes"),
			"OperatingSystem.NumberOfProcesses",
			nil,
			nil,
		),
		ProcessesLimit: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "processes_limit"),
			"OperatingSystem.MaxNumberOfProcesses",
			nil,
			nil,
		),
		ProcessMemoryLimitBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "process_memory_limix_bytes"),
			"OperatingSystem.MaxProcessMemorySize",
			nil,
			nil,
		),
		Users: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "users"),
			"OperatingSystem.NumberOfUsers",
			nil,
			nil,
		),
		VirtualMemoryBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "virtual_memory_bytes"),
			"OperatingSystem.TotalVirtualMemorySize",
			nil,
			nil,
		),
		VisibleMemoryBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "visible_memory_bytes"),
			"OperatingSystem.TotalVisibleMemorySize",
			nil,
			nil,
		),
		VirtualMemoryFreeBytes: prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, subsystem, "virtual_memory_free_bytes"),
			"OperatingSystem.FreeVirtualMemory",
			nil,
			nil,
		),
	}, nil
}

// Collect sends the metric values for each metric
// to the provided prometheus Metric channel.
func (c *OSCollector) Collect(ch chan<- prometheus.Metric) error {
	if desc, err := c.collect(ch); err != nil {
		log.Println("[ERROR] failed collecting os metrics:", desc, err)
		return err
	}
	return nil
}

type Win32_OperatingSystem struct {
	FreePhysicalMemory      uint64
	FreeSpaceInPagingFiles  uint64
	FreeVirtualMemory       uint64
	MaxNumberOfProcesses    uint32
	MaxProcessMemorySize    uint64
	NumberOfProcesses       uint32
	NumberOfUsers           uint32
	SizeStoredInPagingFiles uint64
	TotalVirtualMemorySize  uint64
	TotalVisibleMemorySize  uint64
	LocalDateTime		time.Time
}

func (c *OSCollector) collect(ch chan<- prometheus.Metric) (*prometheus.Desc, error) {
	var dst []Win32_OperatingSystem
	if err := wmi.Query(wmi.CreateQuery(&dst, ""), &dst); err != nil {
		return nil, err
	}

	ch <- prometheus.MustNewConstMetric(
		c.PhysicalMemoryFreeBytes,
		prometheus.GaugeValue,
		float64(dst[0].FreePhysicalMemory*1024), // KiB -> bytes
	)

	time := dst[0].LocalDateTime

	ch <- prometheus.MustNewConstMetric(
		c.Time,
		prometheus.GaugeValue,
		float64(time.Unix()),
	)

	timezoneName, _ := time.Zone()

	ch <- prometheus.MustNewConstMetric(
		c.Timezone,
		prometheus.GaugeValue,
		1.0,
		timezoneName,
	)

	ch <- prometheus.MustNewConstMetric(
		c.PagingFreeBytes,
		prometheus.GaugeValue,
		float64(dst[0].FreeSpaceInPagingFiles*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.VirtualMemoryFreeBytes,
		prometheus.GaugeValue,
		float64(dst[0].FreeVirtualMemory*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.ProcessesLimit,
		prometheus.GaugeValue,
		float64(dst[0].MaxNumberOfProcesses),
	)

	ch <- prometheus.MustNewConstMetric(
		c.ProcessMemoryLimitBytes,
		prometheus.GaugeValue,
		float64(dst[0].MaxProcessMemorySize*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.Processes,
		prometheus.GaugeValue,
		float64(dst[0].NumberOfProcesses),
	)

	ch <- prometheus.MustNewConstMetric(
		c.Users,
		prometheus.GaugeValue,
		float64(dst[0].NumberOfUsers),
	)

	ch <- prometheus.MustNewConstMetric(
		c.PagingLimitBytes,
		prometheus.GaugeValue,
		float64(dst[0].SizeStoredInPagingFiles*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.VirtualMemoryBytes,
		prometheus.GaugeValue,
		float64(dst[0].TotalVirtualMemorySize*1024), // KiB -> bytes
	)

	ch <- prometheus.MustNewConstMetric(
		c.VisibleMemoryBytes,
		prometheus.GaugeValue,
		float64(dst[0].TotalVisibleMemorySize*1024), // KiB -> bytes
	)

	return nil, nil
}
<|endoftext|>