<commit_before>// An example FTP server build on top of go-raval. graval handles the details
// of the FTP protocol, we just provide a basic in-memory persistence driver.
//
// If you're looking to create a custom graval driver, this example is a
// reasonable starting point. I suggest copying this file and changing the
// function bodies as required.
//
// USAGE:
//
//    go get github.com/yob/graval
//    go install github.com/yob/graval/graval-mem
//    ./bin/graval-mem
//
package main

import (
	"github.com/yob/graval"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"time"
)

const (
	fileOne = "This is the first file available for download.\n\nBy Jàmes"
	fileTwo = "This is file number two.\n\n2012-12-04"
)

// A minimal driver for graval that stores everything in memory. The authentication
// details are fixed and the user is unable to upload, delete or rename any files.
//
// This really just exists as a minimal demonstration of the interface graval
// drivers are required to implement.
type MemDriver struct{}

func (driver *MemDriver) Authenticate(user string, pass string) bool {
	return user == "test" && pass == "1234"
}
func (driver *MemDriver) Bytes(path string) (bytes int) {
	switch path {
	case "/one.txt":
		bytes = len(fileOne)
		break
	case "/files/two.txt":
		bytes = len(fileTwo)
		break
	default:
		bytes = -1
	}
	return
}
func (driver *MemDriver) ModifiedTime(path string) (time.Time, error) {
	return time.Now(), nil
}
func (driver *MemDriver) ChangeDir(path string) bool {
	return path == "/" || path == "/files"
}
func (driver *MemDriver) DirContents(path string) (files []os.FileInfo) {
	files = []os.FileInfo{}
	switch path {
	case "/":
		files = append(files, graval.NewDirItem("files"))
		files = append(files, graval.NewFileItem("one.txt", len(fileOne)))
	case "/files":
		files = append(files, graval.NewFileItem("two.txt", len(fileOne)))
	}
	return files
}

func (driver *MemDriver) DeleteDir(path string) bool {
	return false
}
func (driver *MemDriver) DeleteFile(path string) bool {
	return false
}
func (driver *MemDriver) Rename(fromPath string, toPath string) bool {
	return false
}
func (driver *MemDriver) MakeDir(path string) bool {
	return false
}
func (driver *MemDriver) GetFile(path string) (reader io.ReadCloser, err error) {
	switch path {
	case "/one.txt":
		reader = ioutil.NopCloser(strings.NewReader(fileOne))
	case "/files/two.txt":
		reader = ioutil.NopCloser(strings.NewReader(fileTwo))
	}
	return
}
func (driver *MemDriver) PutFile(destPath string, data io.Reader) bool {
	return false
}

// graval requires a factory that will create a new driver instance for each
// client connection. Generally the factory will be fairly minimal. This is
// a good place to read any required config for your driver.
type MemDriverFactory struct{}

func (factory *MemDriverFactory) NewDriver() (graval.FTPDriver, error) {
	return &MemDriver{}, nil
}

// it's alive!
func main() {
	factory := &MemDriverFactory{}
	ftpServer := graval.NewFTPServer(&graval.FTPServerOpts{Factory: factory})
	err := ftpServer.ListenAndServe()
	if err != nil {
		log.Print(err)
		log.Fatal("Error starting server!")
	}
}
<commit_msg>graval-mem: set a custom server name<commit_after>// An example FTP server build on top of go-raval. graval handles the details
// of the FTP protocol, we just provide a basic in-memory persistence driver.
//
// If you're looking to create a custom graval driver, this example is a
// reasonable starting point. I suggest copying this file and changing the
// function bodies as required.
//
// USAGE:
//
//    go get github.com/yob/graval
//    go install github.com/yob/graval/graval-mem
//    ./bin/graval-mem
//
package main

import (
	"github.com/yob/graval"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"time"
)

const (
	fileOne = "This is the first file available for download.\n\nBy Jàmes"
	fileTwo = "This is file number two.\n\n2012-12-04"
)

// A minimal driver for graval that stores everything in memory. The authentication
// details are fixed and the user is unable to upload, delete or rename any files.
//
// This really just exists as a minimal demonstration of the interface graval
// drivers are required to implement.
type MemDriver struct{}

func (driver *MemDriver) Authenticate(user string, pass string) bool {
	return user == "test" && pass == "1234"
}
func (driver *MemDriver) Bytes(path string) (bytes int) {
	switch path {
	case "/one.txt":
		bytes = len(fileOne)
		break
	case "/files/two.txt":
		bytes = len(fileTwo)
		break
	default:
		bytes = -1
	}
	return
}
func (driver *MemDriver) ModifiedTime(path string) (time.Time, error) {
	return time.Now(), nil
}
func (driver *MemDriver) ChangeDir(path string) bool {
	return path == "/" || path == "/files"
}
func (driver *MemDriver) DirContents(path string) (files []os.FileInfo) {
	files = []os.FileInfo{}
	switch path {
	case "/":
		files = append(files, graval.NewDirItem("files"))
		files = append(files, graval.NewFileItem("one.txt", len(fileOne)))
	case "/files":
		files = append(files, graval.NewFileItem("two.txt", len(fileOne)))
	}
	return files
}

func (driver *MemDriver) DeleteDir(path string) bool {
	return false
}
func (driver *MemDriver) DeleteFile(path string) bool {
	return false
}
func (driver *MemDriver) Rename(fromPath string, toPath string) bool {
	return false
}
func (driver *MemDriver) MakeDir(path string) bool {
	return false
}
func (driver *MemDriver) GetFile(path string) (reader io.ReadCloser, err error) {
	switch path {
	case "/one.txt":
		reader = ioutil.NopCloser(strings.NewReader(fileOne))
	case "/files/two.txt":
		reader = ioutil.NopCloser(strings.NewReader(fileTwo))
	}
	return
}
func (driver *MemDriver) PutFile(destPath string, data io.Reader) bool {
	return false
}

// graval requires a factory that will create a new driver instance for each
// client connection. Generally the factory will be fairly minimal. This is
// a good place to read any required config for your driver.
type MemDriverFactory struct{}

func (factory *MemDriverFactory) NewDriver() (graval.FTPDriver, error) {
	return &MemDriver{}, nil
}

// it's alive!
func main() {
	factory := &MemDriverFactory{}
	opts := &graval.FTPServerOpts{
		Factory:    factory,
		ServerName: "graval-mem, the in memory FTP server",
	}
	ftpServer := graval.NewFTPServer(opts)
	err := ftpServer.ListenAndServe()
	if err != nil {
		log.Print(err)
		log.Fatal("Error starting server!")
	}
}
<|endoftext|>


<commit_before>package ics

import (
	"bufio"
	"bytes"
	"errors"
	"io"
	"strings"
)

const (
	ianaTokenChars  = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-"
	invSafeChars    = "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f\",:;\x7f"
	invQSafeChars   = "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f\"\x7f"
	invValueChars   = "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f\x7f"
	paramDelim      = ";"
	paramValueDelim = "="
	nameValueDelim  = ":"
)

type token struct {
	typ  tokenType
	data string
}

type tokenType uint8

const (
	TokenError tokenType = iota
	TokenName
	TokenParamName
	TokenParamValue
	TokenParamQValue
	TokenValue
)

type stateFn func() (token, stateFn)

type parser struct {
	br                 *bufio.Reader
	buf                bytes.Buffer
	state              stateFn
	pos, lineNo, colNo int
	lastWidth          int
	wasNewLine         bool
	err                error
}

func newParser(r io.Reader) *parser {
	p := &parser{
		br: bufio.NewReader(r),
	}
	p.state = p.parseName
	return p
}

func (p *parser) GetToken() (token, error) {
	var t token
	p.buf.Reset()
	t, p.state = p.state()
	if p.err == io.EOF {
		p.state = p.errorFn
		if t.typ == TokenError {
			p.err = io.ErrUnexpectedEOF
		} else {
			return t, nil
		}
	}
	return t, p.err
}

func (p *parser) next() rune {
	if p.err != nil {
		return -1
	}
	r, s, err := p.br.ReadRune()
	if err != nil {
		p.lastWidth = 0
		p.err = err
		return -1
	}
	p.buf.WriteRune(r)
	p.pos += s
	p.lastWidth = s
	if s > 0 {
		if r == '\n' {
			p.lineNo++
			p.wasNewLine = true
			p.colNo = 0
		} else {
			p.colNo++
			p.wasNewLine = false
		}
	}
	return r
}

func (p *parser) backup() {
	if p.lastWidth > 0 {
		p.pos -= p.lastWidth
		if p.wasNewLine {
			p.lineNo--
		}
		p.colNo--
		p.br.UnreadRune()
		p.buf.Truncate(p.buf.Len() - p.lastWidth)
		p.lastWidth = 0
	}
}

func (p *parser) accept(valid string) bool {
	if strings.ContainsRune(valid, p.next()) {
		return true
	}
	p.backup()
	return false
}

func (p *parser) acceptRun(valid string) {
	for {
		r := p.next()
		if r == -1 {
			return
		}
		if !strings.ContainsRune(valid, r) {
			p.backup()
			return
		}
	}
}

func (p *parser) except(invalid string) bool {
	r := p.next()
	if r == -1 {
		return false
	}
	if !strings.ContainsRune(invalid, r) {
		return true
	}
	p.backup()
	return false
}

func (p *parser) exceptRun(invalid string) {
	for {
		r := p.next()
		if r == -1 {
			return
		}
		if strings.ContainsRune(invalid, r) {
			p.backup()
			return
		}
	}
}

func (p *parser) parseName() (token, stateFn) {
	p.acceptRun(ianaTokenChars)
	t := token{
		TokenName,
		p.buf.String(),
	}
	if p.buf.Len() == 0 {
		p.err = ErrNoName
	} else if p.accept(paramDelim) {
		return t, p.parseParamName
	} else if p.accept(nameValueDelim) {
		return t, p.parseValue
	} else if p.err == nil {
		p.err = ErrInvalidChar
	}
	return p.errorFn()
}

func (p *parser) parseParamName() (token, stateFn) {
	p.acceptRun(ianaTokenChars)
	t := token{
		TokenParamName,
		p.buf.String(),
	}
	if p.buf.Len() == 0 {
		p.err = ErrNoParamName
	} else if p.accept(paramValueDelim) {
		return t, p.parseParamValue
	} else if p.err == nil {
		p.err = ErrInvalidChar
	}
	return p.errorFn()
}

func (p *parser) parseParamValue() (token, stateFn) {
	var t token
	if p.accept("\"") {
		p.exceptRun(invQSafeChars)
		if !p.accept("\"") {
			p.err = ErrInvalidChar
			return p.errorFn()
		}
		t.typ = TokenParamQValue
		t.data = p.buf.String()
	} else {
		p.exceptRun(invSafeChars)
		t.typ = TokenParamValue
		t.data = p.buf.String()
	}
	if p.accept(paramDelim) {
		return t, p.parseParamName
	} else if p.accept(nameValueDelim) {
		return t, p.parseValue
	} else if p.err == nil {
		p.err = ErrInvalidChar
	}
	return p.errorFn()
}

func (p *parser) parseValue() (token, stateFn) {
	var toRet []byte
	for {
		p.exceptRun(invValueChars)
		if !p.accept("\r") || !p.accept("\n") {
			if p.err == nil {
				p.err = ErrInvalidChar
			}
			return p.errorFn()
		}
		toAdd := p.buf.Bytes()
		toRet = append(toRet, toAdd[:len(toAdd)-2]...)
		if !p.accept(" ") {
			break
		}
		p.buf.Reset()
	}
	return token{
		TokenValue,
		string(toRet),
	}, p.parseName
}

func (p *parser) errorFn() (token, stateFn) {
	return token{
		TokenError,
		p.err.Error(),
	}, p.errorFn
}

// Errors

var (
	ErrInvalidChar = errors.New("invalid character")
	ErrNoName      = errors.New("zero length name")
	ErrNoParamName = errors.New("zero length param name")
)
<commit_msg>Added a method to clear an line with an error and replaced unnamed constants<commit_after>package ics

import (
	"bufio"
	"bytes"
	"errors"
	"io"
	"strings"
)

const (
	ianaTokenChars  = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-"
	invSafeChars    = "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f\",:;\x7f"
	invQSafeChars   = "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f\"\x7f"
	invValueChars   = "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x0a\x0b\x0c\x0d\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f\x7f"
	paramDelim      = ";"
	paramValueDelim = "="
	nameValueDelim  = ":"
	crlf            = "\r\n"
	dquote          = "\""
)

type token struct {
	typ  tokenType
	data string
}

type tokenType uint8

const (
	TokenError tokenType = iota
	TokenName
	TokenParamName
	TokenParamValue
	TokenParamQValue
	TokenValue
)

type stateFn func() (token, stateFn)

type parser struct {
	br                 *bufio.Reader
	buf                bytes.Buffer
	state              stateFn
	pos, lineNo, colNo int
	lastWidth          int
	wasNewLine         bool
	err                error
}

func newParser(r io.Reader) *parser {
	p := &parser{
		br: bufio.NewReader(r),
	}
	p.state = p.parseName
	return p
}

func (p *parser) GetToken() (token, error) {
	var t token
	p.buf.Reset()
	t, p.state = p.state()
	if p.err == io.EOF {
		p.state = p.errorFn
		if t.typ == TokenError {
			p.err = io.ErrUnexpectedEOF
		} else {
			return t, nil
		}
	}
	return t, p.err
}

func (p *parser) ClearError() {
	if p.err == io.EOF || p.err == io.ErrUnexpectedEOF {
		return
	}
	p.err = nil
	p.state = p.clearLine
}

func (p *parser) next() rune {
	if p.err != nil {
		return -1
	}
	r, s, err := p.br.ReadRune()
	if err != nil {
		p.lastWidth = 0
		p.err = err
		return -1
	}
	p.buf.WriteRune(r)
	p.pos += s
	p.lastWidth = s
	if s > 0 {
		if r == '\n' {
			p.lineNo++
			p.wasNewLine = true
			p.colNo = 0
		} else {
			p.colNo++
			p.wasNewLine = false
		}
	}
	return r
}

func (p *parser) backup() {
	if p.lastWidth > 0 {
		p.pos -= p.lastWidth
		if p.wasNewLine {
			p.lineNo--
		}
		p.colNo--
		p.br.UnreadRune()
		p.buf.Truncate(p.buf.Len() - p.lastWidth)
		p.lastWidth = 0
	}
}

func (p *parser) accept(valid string) bool {
	if strings.ContainsRune(valid, p.next()) {
		return true
	}
	p.backup()
	return false
}

func (p *parser) acceptRun(valid string) {
	for {
		r := p.next()
		if r == -1 {
			return
		}
		if !strings.ContainsRune(valid, r) {
			p.backup()
			return
		}
	}
}

func (p *parser) except(invalid string) bool {
	r := p.next()
	if r == -1 {
		return false
	}
	if !strings.ContainsRune(invalid, r) {
		return true
	}
	p.backup()
	return false
}

func (p *parser) exceptRun(invalid string) {
	for {
		r := p.next()
		if r == -1 {
			return
		}
		if strings.ContainsRune(invalid, r) {
			p.backup()
			return
		}
	}
}

func (p *parser) parseName() (token, stateFn) {
	p.acceptRun(ianaTokenChars)
	t := token{
		TokenName,
		p.buf.String(),
	}
	if p.buf.Len() == 0 {
		p.err = ErrNoName
	} else if p.accept(paramDelim) {
		return t, p.parseParamName
	} else if p.accept(nameValueDelim) {
		return t, p.parseValue
	} else if p.err == nil {
		p.err = ErrInvalidChar
	}
	return p.errorFn()
}

func (p *parser) parseParamName() (token, stateFn) {
	p.acceptRun(ianaTokenChars)
	t := token{
		TokenParamName,
		p.buf.String(),
	}
	if p.buf.Len() == 0 {
		p.err = ErrNoParamName
	} else if p.accept(paramValueDelim) {
		return t, p.parseParamValue
	} else if p.err == nil {
		p.err = ErrInvalidChar
	}
	return p.errorFn()
}

func (p *parser) parseParamValue() (token, stateFn) {
	var t token
	if p.accept(dquote) {
		p.exceptRun(invQSafeChars)
		if !p.accept(dquote) {
			p.err = ErrInvalidChar
			return p.errorFn()
		}
		t.typ = TokenParamQValue
		t.data = p.buf.String()
	} else {
		p.exceptRun(invSafeChars)
		t.typ = TokenParamValue
		t.data = p.buf.String()
	}
	if p.accept(paramDelim) {
		return t, p.parseParamName
	} else if p.accept(nameValueDelim) {
		return t, p.parseValue
	} else if p.err == nil {
		p.err = ErrInvalidChar
	}
	return p.errorFn()
}

func (p *parser) parseValue() (token, stateFn) {
	var toRet []byte
	for {
		p.exceptRun(invValueChars)
		if !p.accept(crlf[:1]) || !p.accept(crlf[1:]) {
			if p.err == nil {
				p.err = ErrInvalidChar
			}
			return p.errorFn()
		}
		toAdd := p.buf.Bytes()
		toRet = append(toRet, toAdd[:len(toAdd)-2]...)
		if !p.accept(" ") {
			break
		}
		p.buf.Reset()
	}
	return token{
		TokenValue,
		string(toRet),
	}, p.parseName
}

func (p *parser) errorFn() (token, stateFn) {
	return token{
		TokenError,
		p.err.Error(),
	}, p.errorFn
}

func (p *parser) clearLine() (token, stateFn) {
	for {
		p.exceptRun(crlf[:1])
		if p.err != nil {
			return p.errorFn()
		} else if p.accept(crlf[:1]) && p.accept(crlf[1:]) {
			return p.parseName()
		}
	}
}

// Errors

var (
	ErrInvalidChar = errors.New("invalid character")
	ErrNoName      = errors.New("zero length name")
	ErrNoParamName = errors.New("zero length param name")
)
<|endoftext|>


<commit_before>// Copyright 2017 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package ddl_test

import (
	"bytes"
	"context"
	"sync/atomic"
	"time"

	. "github.com/pingcap/check"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/tidb/ddl"
	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/store/mockstore"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/tablecodec"
	"github.com/pingcap/tidb/util/testkit"
)

type testDDLTableSplitSuite struct{}

var _ = Suite(&testDDLTableSplitSuite{})

func (s *testDDLTableSplitSuite) TestTableSplit(c *C) {
	store, err := mockstore.NewMockTikvStore()
	c.Assert(err, IsNil)
	defer store.Close()
	session.SetSchemaLease(0)
	session.DisableStats4Test()
	atomic.StoreUint32(&ddl.EnableSplitTableRegion, 1)
	dom, err := session.BootstrapSession(store)
	c.Assert(err, IsNil)
	tk := testkit.NewTestKit(c, store)
	tk.MustExec("use test")
	tk.MustExec(`create table t_part (a int key) partition by range(a) (
		partition p0 values less than (10),
		partition p1 values less than (20)
	)`)
	defer dom.Close()
	atomic.StoreUint32(&ddl.EnableSplitTableRegion, 0)
	infoSchema := dom.InfoSchema()
	c.Assert(infoSchema, NotNil)
	t, err := infoSchema.TableByName(model.NewCIStr("mysql"), model.NewCIStr("tidb"))
	c.Assert(err, IsNil)
	checkRegionStartWithTableID(c, t.Meta().ID, store.(kvStore))

	t, err = infoSchema.TableByName(model.NewCIStr("test"), model.NewCIStr("t_part"))
	c.Assert(err, IsNil)
	pi := t.Meta().GetPartitionInfo()
	c.Assert(pi, NotNil)
	for _, def := range pi.Definitions {
		checkRegionStartWithTableID(c, def.ID, store.(kvStore))
	}
}

type kvStore interface {
	GetRegionCache() *tikv.RegionCache
}

func checkRegionStartWithTableID(c *C, id int64, store kvStore) {
	regionStartKey := tablecodec.EncodeTablePrefix(id)
	var loc *tikv.KeyLocation
	var err error
	for i := 0; i < 10; i++ {
		cache := store.GetRegionCache()
		loc, err = cache.LocateKey(tikv.NewBackoffer(context.Background(), 5000), regionStartKey)
		c.Assert(err, IsNil)

		// Region cache may be out of date, so we need to drop this expired region and load it again.
		cache.InvalidateCachedRegion(loc.Region)
		if bytes.Equal(loc.StartKey, []byte(regionStartKey)) {
			return
		}
		time.Sleep(3 * time.Millisecond)
	}
	c.Assert(loc.StartKey, BytesEquals, []byte(regionStartKey))
}
<commit_msg>ddl: make ddl test more stable (#11897)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package ddl_test

import (
	"context"
	"sync/atomic"

	. "github.com/pingcap/check"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/tidb/ddl"
	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/store/mockstore"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/tablecodec"
	"github.com/pingcap/tidb/util/testkit"
)

type testDDLTableSplitSuite struct{}

var _ = Suite(&testDDLTableSplitSuite{})

func (s *testDDLTableSplitSuite) TestTableSplit(c *C) {
	store, err := mockstore.NewMockTikvStore()
	c.Assert(err, IsNil)
	defer store.Close()
	session.SetSchemaLease(0)
	session.DisableStats4Test()
	atomic.StoreUint32(&ddl.EnableSplitTableRegion, 1)
	dom, err := session.BootstrapSession(store)
	c.Assert(err, IsNil)
	tk := testkit.NewTestKit(c, store)
	tk.MustExec("use test")
	// Synced split table region.
	tk.MustExec("set global tidb_scatter_region = 1")
	tk.MustExec(`create table t_part (a int key) partition by range(a) (
		partition p0 values less than (10),
		partition p1 values less than (20)
	)`)
	defer dom.Close()
	atomic.StoreUint32(&ddl.EnableSplitTableRegion, 0)
	infoSchema := dom.InfoSchema()
	c.Assert(infoSchema, NotNil)
	t, err := infoSchema.TableByName(model.NewCIStr("mysql"), model.NewCIStr("tidb"))
	c.Assert(err, IsNil)
	checkRegionStartWithTableID(c, t.Meta().ID, store.(kvStore))

	t, err = infoSchema.TableByName(model.NewCIStr("test"), model.NewCIStr("t_part"))
	c.Assert(err, IsNil)
	pi := t.Meta().GetPartitionInfo()
	c.Assert(pi, NotNil)
	for _, def := range pi.Definitions {
		checkRegionStartWithTableID(c, def.ID, store.(kvStore))
	}
}

type kvStore interface {
	GetRegionCache() *tikv.RegionCache
}

func checkRegionStartWithTableID(c *C, id int64, store kvStore) {
	regionStartKey := tablecodec.EncodeTablePrefix(id)
	var loc *tikv.KeyLocation
	var err error
	cache := store.GetRegionCache()
	loc, err = cache.LocateKey(tikv.NewBackoffer(context.Background(), 5000), regionStartKey)
	c.Assert(err, IsNil)
	// Region cache may be out of date, so we need to drop this expired region and load it again.
	cache.InvalidateCachedRegion(loc.Region)
	c.Assert([]byte(loc.StartKey), BytesEquals, []byte(regionStartKey))
}
<|endoftext|>


<commit_before>///////////////////////////////////////////////////////////////////////
// Copyright (C) 2016 VMware, Inc. All rights reserved.
// -- VMware Confidential
///////////////////////////////////////////////////////////////////////

package task

import (
	"context"
	"errors"
	"sync"
	"time"
)

type ParalStep struct {
	GenericStep
	aggContextFn func(context.Context, context.Context) context.Context
	aggErrorFn   func([]error) error
}

//NewParalStep creates a new parallel step whose substeps can be executed at the same time
//Note that the new Step should be of state StepStateNone, and all of its substeps should be of state StepStateNone too.
func NewParalStep(stepInfo StepInfo, contextfn func(context.Context, context.Context) context.Context,
	errorfn func([]error) error, steps ...Step) *ParalStep {
	//the caller is responsible to make sure stepOpts and all step's state are set to StateNone
	return &ParalStep{
		GenericStep:  GenericStep{StepInfo: stepInfo, Steps: steps},
		aggContextFn: contextfn,
		aggErrorFn:   errorfn,
	}
}

func (s *ParalStep) Success(reqCtx context.Context) {
	if s.successFn == nil {
		s.SetState(StateCompleted)
	} else {
		s.successFn(reqCtx, s)
	}
}

func (s *ParalStep) Fail(reqCtx context.Context, err error) {
	if s.failFn == nil {
		s.SetState(StateFailed)
	} else {
		s.failFn(reqCtx, s, err)
	}
}

func (s *ParalStep) Run(reqCtx context.Context) (context.Context, error) {

	s.State = StateProcessing
	var runError error

	ctxc := make(chan context.Context)
	errc := make(chan error)

	var resultCtx context.Context
	var resultErr error
	var cancelErr error

	go func(reqCtx context.Context) {
		ctx := reqCtx
		getCtx := false

		for r := range ctxc {
			ctx = s.aggContextFn(ctx, r)
			getCtx = true
		}
		if getCtx {
			resultCtx = ctx
		}
	}(reqCtx)

	go func() {
		var stepErrors []error

		for e := range errc {
			stepErrors = append(stepErrors, e)
		}
		if stepErrors != nil {
			resultErr = s.aggErrorFn(stepErrors)
		}
	}()

	go func(ctx context.Context) {
		select {
		case <-reqCtx.Done():
			cancelErr = errors.New("step " + s.Name + " canceled")
		}

	}(reqCtx)

	var wg1 sync.WaitGroup
	wg1.Add(len(s.Steps))
	for _, step := range s.Steps {
		ctx := reqCtx
		go func(step Step, ctx context.Context) {
			//for parallel step, we do not care about returned ctx from step.Run()

			ctx, err := step.Run(ctx)
			if err != nil {
				//what do we do? let the consumer decide I guess
				//do we still need the context from this run aggregated? - no

				errc <- err
			} else {
				ctxc <- ctx
			}
			wg1.Done()
		}(step, ctx)
	}

	wg1.Wait()

	//the problem of this logic is if we always wait , then cancel signal is not canceling anything?
	//even the example in https://blog.golang.org/pipelines is not actually canceling anything.
	//it just not reporting to the result if canceld.
	//it seems the only use is that we get the cancel signal, we rollback after wait finishes

	close(ctxc)
	close(errc)

	var errs []error
	//let us wait a little bit till value received from channels
	timeout := time.After(2 * time.Second)
	<-timeout

	// fmt.Printf("step %s processing result\n", s.Name)

	if cancelErr != nil {
		errs = append(errs, cancelErr)

		if resultErr != nil {
			errs = append(errs, resultErr)
		}
		_, rollbackError := s.Rollback(reqCtx)

		if rollbackError != nil {
			errs = append(errs, rollbackError)
		}
		runError = s.aggErrorFn(errs)
		return reqCtx, runError
	} else if resultErr != nil {
		errs = append(errs, resultErr)
		runError = s.aggErrorFn(errs)
		s.Fail(reqCtx, runError)

		return reqCtx, runError
	} else if resultCtx != nil {
		s.Success(resultCtx)
		return resultCtx, nil
	}

	return reqCtx, nil
}
<commit_msg>replace time.After() with wait groups on paral step<commit_after>///////////////////////////////////////////////////////////////////////
// Copyright (C) 2016 VMware, Inc. All rights reserved.
// -- VMware Confidential
///////////////////////////////////////////////////////////////////////

package task

import (
	"context"
	"errors"
	"sync"
)

type ParalStep struct {
	GenericStep
	aggContextFn func(context.Context, context.Context) context.Context
	aggErrorFn   func([]error) error
}

//NewParalStep creates a new parallel step whose substeps can be executed at the same time
//Note that the new Step should be of state StepStateNone, and all of its substeps should be of state StepStateNone too.
func NewParalStep(stepInfo StepInfo, contextfn func(context.Context, context.Context) context.Context,
	errorfn func([]error) error, steps ...Step) *ParalStep {
	//the caller is responsible to make sure stepOpts and all step's state are set to StateNone
	return &ParalStep{
		GenericStep:  GenericStep{StepInfo: stepInfo, Steps: steps},
		aggContextFn: contextfn,
		aggErrorFn:   errorfn,
	}
}

func (s *ParalStep) Success(reqCtx context.Context) {
	if s.successFn == nil {
		s.SetState(StateCompleted)
	} else {
		s.successFn(reqCtx, s)
	}
}

func (s *ParalStep) Fail(reqCtx context.Context, err error) {
	if s.failFn == nil {
		s.SetState(StateFailed)
	} else {
		s.failFn(reqCtx, s, err)
	}
}

func (s *ParalStep) Run(reqCtx context.Context) (context.Context, error) {

	s.State = StateProcessing
	var runError error

	ctxc := make(chan context.Context)
	errc := make(chan error)

	var resultCtx context.Context
	var resultErr error
	var cancelErr error

	var wgCtx sync.WaitGroup
	wgCtx.Add(1)

	go func(reqCtx context.Context) {
		ctx := reqCtx
		getCtx := false

		for r := range ctxc {
			ctx = s.aggContextFn(ctx, r)
			getCtx = true
		}
		if getCtx {
			resultCtx = ctx
		}
		wgCtx.Done()
	}(reqCtx)

	var wgErr sync.WaitGroup
	wgErr.Add(1)
	go func() {
		var stepErrors []error

		for e := range errc {
			stepErrors = append(stepErrors, e)
		}
		if stepErrors != nil {
			resultErr = s.aggErrorFn(stepErrors)
		}
		wgErr.Done()
	}()

	go func(ctx context.Context) {
		select {
		case <-reqCtx.Done():
			cancelErr = errors.New("step " + s.Name + " canceled")
		}

	}(reqCtx)

	var wg1 sync.WaitGroup
	wg1.Add(len(s.Steps))
	for _, step := range s.Steps {
		ctx := reqCtx
		go func(step Step, ctx context.Context) {
			//for parallel step, we do not care about returned ctx from step.Run()

			ctx, err := step.Run(ctx)
			if err != nil {
				//what do we do? let the consumer decide I guess
				//do we still need the context from this run aggregated? - no

				errc <- err
			} else {
				ctxc <- ctx
			}
			wg1.Done()
		}(step, ctx)
	}

	wg1.Wait()

	//the problem of this logic is if we always wait , then cancel signal is not canceling anything?
	//even the example in https://blog.golang.org/pipelines is not actually canceling anything.
	//it just not reporting to the result if canceld.
	//it seems the only use is that we get the cancel signal, we rollback after wait finishes

	close(ctxc)
	close(errc)

	wgCtx.Wait()
	wgErr.Wait()

	var errs []error

	// fmt.Printf("step %s processing result\n", s.Name)

	if cancelErr != nil {
		errs = append(errs, cancelErr)

		if resultErr != nil {
			errs = append(errs, resultErr)
		}
		_, rollbackError := s.Rollback(reqCtx)

		if rollbackError != nil {
			errs = append(errs, rollbackError)
		}
		runError = s.aggErrorFn(errs)
		return reqCtx, runError
	} else if resultErr != nil {
		errs = append(errs, resultErr)
		runError = s.aggErrorFn(errs)
		s.Fail(reqCtx, runError)

		return reqCtx, runError
	} else if resultCtx != nil {
		s.Success(resultCtx)
		return resultCtx, nil
	}

	return reqCtx, nil
}
<|endoftext|>


<commit_before>package mtproto

import (
	"context"

	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/crypto"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/predicate"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/task"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/pipe"
)

var (
	dcList = []net.Address{
		net.ParseAddress("149.154.175.50"),
		net.ParseAddress("149.154.167.51"),
		net.ParseAddress("149.154.175.100"),
		net.ParseAddress("149.154.167.91"),
		net.ParseAddress("149.154.171.5"),
	}
)

type Server struct {
	user    *protocol.User
	account *Account
}

func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
	if len(config.User) == 0 {
		return nil, newError("no user configured.")
	}

	user := config.User[0]
	rawAccount, err := config.User[0].GetTypedAccount()
	if err != nil {
		return nil, newError("invalid account").Base(err)
	}
	account, ok := rawAccount.(*Account)
	if !ok {
		return nil, newError("not a MTProto account")
	}

	return &Server{
		user:    user,
		account: account,
	}, nil
}

func (s *Server) Network() net.NetworkList {
	return net.NetworkList{
		Network: []net.Network{net.Network_TCP},
	}
}

func (s *Server) Process(ctx context.Context, network net.Network, conn internet.Connection, dispatcher core.Dispatcher) error {
	auth, err := ReadAuthentication(conn)
	if err != nil {
		return newError("failed to read authentication header").Base(err)
	}
	defer putAuthenticationObject(auth)

	auth.ApplySecret(s.account.Secret)

	decryptor := crypto.NewAesCTRStream(auth.DecodingKey[:], auth.DecodingNonce[:])
	decryptor.XORKeyStream(auth.Header[:], auth.Header[:])

	if !predicate.BytesAll(auth.Header[56:60], 0xef) {
		return newError("invalid connection type: ", auth.Header[56:60])
	}

	dcID := auth.DataCenterID()
	if dcID >= uint16(len(dcList)) {
		return newError("invalid data center id: ", dcID)
	}

	dest := net.Destination{
		Network: net.Network_TCP,
		Address: dcList[dcID],
		Port:    net.Port(443),
	}
	link, err := dispatcher.Dispatch(ctx, dest)
	if err != nil {
		return newError("failed to dispatch request to: ", dest).Base(err)
	}

	request := func() error {
		reader := buf.NewReader(crypto.NewCryptionReader(decryptor, conn))
		return buf.Copy(reader, link.Writer)
	}

	response := func() error {
		encryptor := crypto.NewAesCTRStream(auth.EncodingKey[:], auth.EncodingNonce[:])
		writer := buf.NewWriter(crypto.NewCryptionWriter(encryptor, conn))
		return buf.Copy(link.Reader, writer)
	}

	var responseDoneAndCloseWriter = task.Single(response, task.OnSuccess(task.Close(link.Writer)))
	if err := task.Run(task.WithContext(ctx), task.Parallel(request, responseDoneAndCloseWriter))(); err != nil {
		pipe.CloseError(link.Reader)
		pipe.CloseError(link.Writer)
		return newError("connection ends").Base(err)
	}

	return nil
}

func init() {
	common.Must(common.RegisterConfig((*ServerConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewServer(ctx, config.(*ServerConfig))
	}))
}
<commit_msg>apply policy in mtproto inbound<commit_after>package mtproto

import (
	"context"
	"time"

	"v2ray.com/core"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/crypto"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/predicate"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal"
	"v2ray.com/core/common/task"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/pipe"
)

var (
	dcList = []net.Address{
		net.ParseAddress("149.154.175.50"),
		net.ParseAddress("149.154.167.51"),
		net.ParseAddress("149.154.175.100"),
		net.ParseAddress("149.154.167.91"),
		net.ParseAddress("149.154.171.5"),
	}
)

type Server struct {
	user    *protocol.User
	account *Account
	policy  core.PolicyManager
}

func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
	if len(config.User) == 0 {
		return nil, newError("no user configured.")
	}

	user := config.User[0]
	rawAccount, err := config.User[0].GetTypedAccount()
	if err != nil {
		return nil, newError("invalid account").Base(err)
	}
	account, ok := rawAccount.(*Account)
	if !ok {
		return nil, newError("not a MTProto account")
	}

	v := core.MustFromContext(ctx)

	return &Server{
		user:    user,
		account: account,
		policy:  v.PolicyManager(),
	}, nil
}

func (s *Server) Network() net.NetworkList {
	return net.NetworkList{
		Network: []net.Network{net.Network_TCP},
	}
}

func (s *Server) Process(ctx context.Context, network net.Network, conn internet.Connection, dispatcher core.Dispatcher) error {
	sPolicy := s.policy.ForLevel(s.user.Level)

	if err := conn.SetDeadline(time.Now().Add(sPolicy.Timeouts.Handshake)); err != nil {
		newError("failed to set deadline").Base(err).WriteToLog(session.ExportIDToError(ctx))
	}
	auth, err := ReadAuthentication(conn)
	if err != nil {
		return newError("failed to read authentication header").Base(err)
	}
	defer putAuthenticationObject(auth)

	if err := conn.SetDeadline(time.Time{}); err != nil {
		newError("failed to clear deadline").Base(err).WriteToLog(session.ExportIDToError(ctx))
	}

	auth.ApplySecret(s.account.Secret)

	decryptor := crypto.NewAesCTRStream(auth.DecodingKey[:], auth.DecodingNonce[:])
	decryptor.XORKeyStream(auth.Header[:], auth.Header[:])

	if !predicate.BytesAll(auth.Header[56:60], 0xef) {
		return newError("invalid connection type: ", auth.Header[56:60])
	}

	dcID := auth.DataCenterID()
	if dcID >= uint16(len(dcList)) {
		return newError("invalid data center id: ", dcID)
	}

	dest := net.Destination{
		Network: net.Network_TCP,
		Address: dcList[dcID],
		Port:    net.Port(443),
	}

	ctx, cancel := context.WithCancel(ctx)
	timer := signal.CancelAfterInactivity(ctx, cancel, sPolicy.Timeouts.ConnectionIdle)
	ctx = core.ContextWithBufferPolicy(ctx, sPolicy.Buffer)

	link, err := dispatcher.Dispatch(ctx, dest)
	if err != nil {
		return newError("failed to dispatch request to: ", dest).Base(err)
	}

	request := func() error {
		defer timer.SetTimeout(sPolicy.Timeouts.DownlinkOnly)

		reader := buf.NewReader(crypto.NewCryptionReader(decryptor, conn))
		return buf.Copy(reader, link.Writer)
	}

	response := func() error {
		defer timer.SetTimeout(sPolicy.Timeouts.UplinkOnly)

		encryptor := crypto.NewAesCTRStream(auth.EncodingKey[:], auth.EncodingNonce[:])
		writer := buf.NewWriter(crypto.NewCryptionWriter(encryptor, conn))
		return buf.Copy(link.Reader, writer)
	}

	var responseDoneAndCloseWriter = task.Single(response, task.OnSuccess(task.Close(link.Writer)))
	if err := task.Run(task.WithContext(ctx), task.Parallel(request, responseDoneAndCloseWriter))(); err != nil {
		pipe.CloseError(link.Reader)
		pipe.CloseError(link.Writer)
		return newError("connection ends").Base(err)
	}

	return nil
}

func init() {
	common.Must(common.RegisterConfig((*ServerConfig)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return NewServer(ctx, config.(*ServerConfig))
	}))
}
<|endoftext|>


<commit_before>/* Copyright 2017 Google Inc.

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

    https://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 mock

import (
	"github.com/google/link022/generated/ocstruct"
	"github.com/openconfig/ygot/ygot"
)

var (
	apName = "test-pi-1"

	radiusServerGroupName = "radius-server-group"
	radiusServerAddr      = "1.1.1.1"

	// GuestWLANName is the name of a mock guest WLAN.
	GuestWLANName = "Guest-Emu"
	// AuthWLANName is the name of a mock authenticated WLAN.
	AuthWLANName = "Auth-Emu"
)

// GenerateConfig generates an AP wireless for test.
func GenerateConfig(addAuthWLAN bool) *ocstruct.Device {
	ap := &ocstruct.Device{
		Hostname: ygot.String(apName),
	}

	if addAuthWLAN {
		ap.System = systemInfo()
	}

	ap.Radios = radios()
	ap.Ssids = wlans(addAuthWLAN)
	return ap
}

// RadiusServer generates a mock RadiusServer configuration.
func RadiusServer() *ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server {
	return &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server{
		Address: ygot.String(radiusServerAddr),
		Radius: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server_Radius{
			Config: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server_Radius_Config{
				AuthPort:  ygot.Uint16(1812),
				SecretKey: ygot.String("radiuspwd"),
			},
		},
	}
}

func systemInfo() *ocstruct.OpenconfigOfficeAp_System {
	return &ocstruct.OpenconfigOfficeAp_System{
		Aaa: &ocstruct.OpenconfigOfficeAp_System_Aaa{
			ServerGroups: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups{
				ServerGroup: map[string]*ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup{
					radiusServerGroupName: {
						Name: ygot.String(radiusServerGroupName),
						Config: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Config{
							Type: ocstruct.OpenconfigAaaTypes_AAA_SERVER_TYPE_RADIUS,
						},
						Servers: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers{
							Server: map[string]*ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server{
								radiusServerAddr: RadiusServer(),
							},
						},
					},
				},
			},
		},
	}
}

func radios() *ocstruct.OpenconfigOfficeAp_Radios {
	radios := &ocstruct.OpenconfigOfficeAp_Radios{}
	radios.Radio = make(map[uint8]*ocstruct.OpenconfigOfficeAp_Radios_Radio)

	radioID := uint8(1)
	radios.Radio[radioID] = &ocstruct.OpenconfigOfficeAp_Radios_Radio{
		Id: ygot.Uint8(radioID),
		Config: &ocstruct.OpenconfigOfficeAp_Radios_Radio_Config{
			Id:                 ygot.Uint8(radioID),
			Enabled:            ygot.Bool(true),
			OperatingFrequency: ocstruct.OpenconfigWifiTypes_OPERATING_FREQUENCY_FREQ_2GHZ,
			TransmitPower:      ygot.Uint8(5),
			Channel:            ygot.Uint8(8),
			ChannelWidth:       ygot.Uint8(10),
			Scanning:           ygot.Bool(true),
			ScanningInterval:   ygot.Uint8(30),
		},
	}

	return radios
}

func wlans(addAuthWLAN bool) *ocstruct.OpenconfigOfficeAp_Ssids {
	wlans := &ocstruct.OpenconfigOfficeAp_Ssids{}
	wlans.Ssid = make(map[string]*ocstruct.OpenconfigOfficeAp_Ssids_Ssid)

	wlans.Ssid[GuestWLANName] = guestWLAN()

	// Add auth WLAN.
	if addAuthWLAN {
		wlans.Ssid[AuthWLANName] = authWLAN()
	}

	return wlans
}

func guestWLAN() *ocstruct.OpenconfigOfficeAp_Ssids_Ssid {
	return &ocstruct.OpenconfigOfficeAp_Ssids_Ssid{
		Name: ygot.String(GuestWLANName),
		Config: &ocstruct.OpenconfigOfficeAp_Ssids_Ssid_Config{
			AdvertiseApname:    ygot.Bool(false),
			BasicDataRates:     []ocstruct.E_OpenconfigWifiTypes_DATA_RATE{ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_11MB, ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_24MB},
			BroadcastFilter:    ygot.Bool(false),
			Csa:                ygot.Bool(false),
			DhcpRequired:       ygot.Bool(true),
			Dot11K:             ygot.Bool(false),
			Dva:                ygot.Bool(false),
			Enabled:            ygot.Bool(true),
			GtkTimeout:         ygot.Uint16(1000),
			Hidden:             ygot.Bool(false),
			MulticastFilter:    ygot.Bool(false),
			Name:               ygot.String(GuestWLANName),
			OperatingFrequency: ocstruct.OpenconfigWifiTypes_OPERATING_FREQUENCY_FREQ_2_5_GHZ,
			Opmode:             ocstruct.OpenconfigOfficeAp_Ssids_Ssid_Config_Opmode_OPEN,
			PtkTimeout:         ygot.Uint16(1000),
			SupportedDataRates: []ocstruct.E_OpenconfigWifiTypes_DATA_RATE{ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_11MB, ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_24MB},
			VlanId:             ygot.Uint16(666),
		},
	}
}

func authWLAN() *ocstruct.OpenconfigOfficeAp_Ssids_Ssid {
	return &ocstruct.OpenconfigOfficeAp_Ssids_Ssid{
		Name: ygot.String(AuthWLANName),
		Config: &ocstruct.OpenconfigOfficeAp_Ssids_Ssid_Config{
			AdvertiseApname:    ygot.Bool(false),
			BasicDataRates:     []ocstruct.E_OpenconfigWifiTypes_DATA_RATE{ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_11MB, ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_24MB},
			BroadcastFilter:    ygot.Bool(false),
			Csa:                ygot.Bool(false),
			DhcpRequired:       ygot.Bool(true),
			Dot11K:             ygot.Bool(false),
			Dva:                ygot.Bool(false),
			Enabled:            ygot.Bool(true),
			GtkTimeout:         ygot.Uint16(1000),
			Hidden:             ygot.Bool(false),
			MulticastFilter:    ygot.Bool(false),
			Name:               ygot.String(AuthWLANName),
			OperatingFrequency: ocstruct.OpenconfigWifiTypes_OPERATING_FREQUENCY_FREQ_2_5_GHZ,
			Opmode:             ocstruct.OpenconfigOfficeAp_Ssids_Ssid_Config_Opmode_WPA2_ENTERPRISE,
			ServerGroup:        ygot.String(radiusServerGroupName),
			PtkTimeout:         ygot.Uint16(1000),
			SupportedDataRates: []ocstruct.E_OpenconfigWifiTypes_DATA_RATE{ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_11MB, ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_24MB},
			VlanId:             ygot.Uint16(250),
		},
	}
}
<commit_msg>Fixed the validation error in OpenConfig example code.<commit_after>/* Copyright 2017 Google Inc.

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

    https://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 mock

import (
	"github.com/google/link022/generated/ocstruct"
	"github.com/openconfig/ygot/ygot"
)

var (
	apName = "test-pi-1"

	radiusServerGroupName = "radius-server-group"
	radiusServerAddr      = "1.1.1.1"

	// GuestWLANName is the name of a mock guest WLAN.
	GuestWLANName = "Guest-Emu"
	// AuthWLANName is the name of a mock authenticated WLAN.
	AuthWLANName = "Auth-Emu"
)

// GenerateConfig generates an AP wireless for test.
func GenerateConfig(addAuthWLAN bool) *ocstruct.Device {
	ap := &ocstruct.Device{
		Hostname: ygot.String(apName),
	}

	if addAuthWLAN {
		ap.System = systemInfo()
	}

	ap.Radios = radios()
	ap.Ssids = wlans(addAuthWLAN)
	return ap
}

// RadiusServer generates a mock RadiusServer configuration.
func RadiusServer() *ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server {
	return &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server{
		Address: ygot.String(radiusServerAddr),
		Radius: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server_Radius{
			Config: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server_Radius_Config{
				AuthPort:  ygot.Uint16(1812),
				SecretKey: ygot.String("radiuspwd"),
			},
		},
	}
}

func systemInfo() *ocstruct.OpenconfigOfficeAp_System {
	return &ocstruct.OpenconfigOfficeAp_System{
		Aaa: &ocstruct.OpenconfigOfficeAp_System_Aaa{
			ServerGroups: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups{
				ServerGroup: map[string]*ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup{
					radiusServerGroupName: {
						Name: ygot.String(radiusServerGroupName),
						Config: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Config{
							Name: ygot.String(radiusServerGroupName),
							Type: ocstruct.OpenconfigAaaTypes_AAA_SERVER_TYPE_RADIUS,
						},
						Servers: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers{
							Server: map[string]*ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server{
								radiusServerAddr: {
									Config: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server_Config{
										Address: ygot.String(radiusServerAddr),
									},
									Address: ygot.String(radiusServerAddr),
									Radius: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server_Radius{
										Config: &ocstruct.OpenconfigOfficeAp_System_Aaa_ServerGroups_ServerGroup_Servers_Server_Radius_Config{
											AuthPort:  ygot.Uint16(1812),
											SecretKey: ygot.String("radiuspwd"),
										},
									},
								},
							},
						},
					},
				},
			},
		},
	}
}

func radios() *ocstruct.OpenconfigOfficeAp_Radios {
	radios := &ocstruct.OpenconfigOfficeAp_Radios{}
	radios.Radio = make(map[uint8]*ocstruct.OpenconfigOfficeAp_Radios_Radio)

	radioID := uint8(1)
	radios.Radio[radioID] = &ocstruct.OpenconfigOfficeAp_Radios_Radio{
		Id: ygot.Uint8(radioID),
		Config: &ocstruct.OpenconfigOfficeAp_Radios_Radio_Config{
			Id:                 ygot.Uint8(radioID),
			Enabled:            ygot.Bool(true),
			OperatingFrequency: ocstruct.OpenconfigWifiTypes_OPERATING_FREQUENCY_FREQ_2GHZ,
			TransmitPower:      ygot.Uint8(5),
			Channel:            ygot.Uint8(8),
			ChannelWidth:       ygot.Uint8(10),
			Scanning:           ygot.Bool(true),
			ScanningInterval:   ygot.Uint8(30),
		},
	}

	return radios
}

func wlans(addAuthWLAN bool) *ocstruct.OpenconfigOfficeAp_Ssids {
	wlans := &ocstruct.OpenconfigOfficeAp_Ssids{}
	wlans.Ssid = make(map[string]*ocstruct.OpenconfigOfficeAp_Ssids_Ssid)

	wlans.Ssid[GuestWLANName] = guestWLAN()

	// Add auth WLAN.
	if addAuthWLAN {
		wlans.Ssid[AuthWLANName] = authWLAN()
	}

	return wlans
}

func guestWLAN() *ocstruct.OpenconfigOfficeAp_Ssids_Ssid {
	return &ocstruct.OpenconfigOfficeAp_Ssids_Ssid{
		Name: ygot.String(GuestWLANName),
		Config: &ocstruct.OpenconfigOfficeAp_Ssids_Ssid_Config{
			AdvertiseApname:    ygot.Bool(false),
			BasicDataRates:     []ocstruct.E_OpenconfigWifiTypes_DATA_RATE{ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_11MB, ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_24MB},
			BroadcastFilter:    ygot.Bool(false),
			Csa:                ygot.Bool(false),
			DhcpRequired:       ygot.Bool(true),
			Dot11K:             ygot.Bool(false),
			Dva:                ygot.Bool(false),
			Enabled:            ygot.Bool(true),
			GtkTimeout:         ygot.Uint16(1000),
			Hidden:             ygot.Bool(false),
			MulticastFilter:    ygot.Bool(false),
			Name:               ygot.String(GuestWLANName),
			OperatingFrequency: ocstruct.OpenconfigWifiTypes_OPERATING_FREQUENCY_FREQ_2_5_GHZ,
			Opmode:             ocstruct.OpenconfigOfficeAp_Ssids_Ssid_Config_Opmode_OPEN,
			PtkTimeout:         ygot.Uint16(1000),
			SupportedDataRates: []ocstruct.E_OpenconfigWifiTypes_DATA_RATE{ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_11MB, ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_24MB},
			VlanId:             ygot.Uint16(666),
		},
	}
}

func authWLAN() *ocstruct.OpenconfigOfficeAp_Ssids_Ssid {
	return &ocstruct.OpenconfigOfficeAp_Ssids_Ssid{
		Name: ygot.String(AuthWLANName),
		Config: &ocstruct.OpenconfigOfficeAp_Ssids_Ssid_Config{
			AdvertiseApname:    ygot.Bool(false),
			BasicDataRates:     []ocstruct.E_OpenconfigWifiTypes_DATA_RATE{ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_11MB, ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_24MB},
			BroadcastFilter:    ygot.Bool(false),
			Csa:                ygot.Bool(false),
			DhcpRequired:       ygot.Bool(true),
			Dot11K:             ygot.Bool(false),
			Dva:                ygot.Bool(false),
			Enabled:            ygot.Bool(true),
			GtkTimeout:         ygot.Uint16(1000),
			Hidden:             ygot.Bool(false),
			MulticastFilter:    ygot.Bool(false),
			Name:               ygot.String(AuthWLANName),
			OperatingFrequency: ocstruct.OpenconfigWifiTypes_OPERATING_FREQUENCY_FREQ_2_5_GHZ,
			Opmode:             ocstruct.OpenconfigOfficeAp_Ssids_Ssid_Config_Opmode_WPA2_ENTERPRISE,
			ServerGroup:        ygot.String(radiusServerGroupName),
			PtkTimeout:         ygot.Uint16(1000),
			SupportedDataRates: []ocstruct.E_OpenconfigWifiTypes_DATA_RATE{ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_11MB, ocstruct.OpenconfigWifiTypes_DATA_RATE_RATE_24MB},
			VlanId:             ygot.Uint16(250),
		},
	}
}
<|endoftext|>


<commit_before>package web

import (
	"flag"
	"os"
	"path"
	"path/filepath"
	"runtime/debug"
	"testing"

	"github.com/h2oai/steamY/lib/fs"
	"github.com/h2oai/steamY/master/az"
	"github.com/h2oai/steamY/master/data"
	"github.com/h2oai/steamY/srv/web"
)

const (
	superuser = "superuser"
)

var h2oFrames = []struct {
	name string
}{
	{"bin_hex"},
	{"mul_hex"},
	{"reg_hex"},
}

var h2oModels = []struct {
	name     string
	category string
}{
	{"bin_gbm", "Binomial"},
	{"bin_glm", "Binomial"},
	{"bin_dpl", "Binomial"},
	{"mul_gbm", "Multinomial"},
	{"mul_dpl", "Multinomial"},
	{"reg_gbm", "Regression"},
	{"reg_glm", "Regression"},
	{"reg_dpl", "Regression"},
}

type test struct {
	t   *testing.T
	svc web.Service
	dir az.Directory
	su  az.Principal
}

var clusterAddress, workingDirectory, compilationServiceAddress string

func init() {
	flag.StringVar(&clusterAddress, "cluster-address", "localhost:54321", "Where the h2o cluster can be reached.")
	flag.StringVar(&workingDirectory, "working-directory", "", "Where the var folder will be located.")
	flag.StringVar(&compilationServiceAddress, "compilation-service-address", ":8080", "Where to find the compilation service.")
}

func newTest(t *testing.T) *test {
	dbOpts := driverDBOpts{
		data.Connection {DbName:"steam", User:"steam", SSLMode:"disable"}
		superuser,
		superuser,
	}

	// Truncate database tables

	if err := data.Destroy(dbOpts.Name, dbOpts.Username, dbOpts.SSLMode); err != nil {
		t.Fatalf("Failed truncating database: %s", err)
	}

	// Determine current directory

	wd, err := filepath.Abs(filepath.Dir(workingDirectory + "/"))
	if err != nil {
		t.Fatalf("Failed determining current directory: %s", err)
	}

	// Delete any remnant models in models directory

	if err := os.RemoveAll(path.Join(wd, "var/master/model")); err != nil {
		t.Fatalf("Failed removing old model directory: %v", err)
	}

	// Create service instance

	opts := driverOpts{
		path.Join(wd, fs.VarDir, "master"),
		":9001",
		compilationServiceAddress,
		"",
		[2]int{1025, 65535},
		driverYarnOpts{false, "", ""},
		dbOpts,
	}
	svc, dir, err := newService(opts)
	if err != nil {
		t.Fatal(err)
	}

	// Lookup superuser

	su, err := dir.Lookup(superuser)
	if err != nil {
		t.Fatal(err)
	}

	return &test{t, svc, dir, su}
}

func (t *test) log(args ...interface{}) {
	t.t.Log(args...)
}

func (t *test) fail(format string, args ...interface{}) {
	debug.PrintStack()
	t.t.Fatalf(format, args...)
}

func (t *test) nil(arg interface{}) {
	if arg != nil {
		t.fail("unexpected: %s", arg)
	}
}

func (t *test) notnil(arg interface{}) {
	if arg == nil {
		t.fail("unexpected nil: %s", arg)
	}
}

func (t *test) ok(condition bool, format string, args ...interface{}) {
	if !condition {
		t.fail(format, args...)
	}
}

func buildPermissionMap(t *test) map[string]int64 {
	permissions, err := t.svc.GetAllPermissions(t.su)
	t.nil(err)

	permissionMap := make(map[string]int64)
	for _, permission := range permissions {
		permissionMap[permission.Code] = permission.Id
	}

	return permissionMap
}

func buildEntityTypeMap(t *test) map[string]int64 {
	entityTypes, err := t.svc.GetAllEntityTypes(t.su)
	t.nil(err)

	entityTypeMap := make(map[string]int64)
	for _, entityType := range entityTypes {
		entityTypeMap[entityType.Name] = entityType.Id
	}
	return entityTypeMap
}
<commit_msg>STEAM-414 Fix typo<commit_after>package web

import (
	"flag"
	"os"
	"path"
	"path/filepath"
	"runtime/debug"
	"testing"

	"github.com/h2oai/steamY/lib/fs"
	"github.com/h2oai/steamY/master/az"
	"github.com/h2oai/steamY/master/data"
	"github.com/h2oai/steamY/srv/web"
)

const (
	superuser = "superuser"
)

var h2oFrames = []struct {
	name string
}{
	{"bin_hex"},
	{"mul_hex"},
	{"reg_hex"},
}

var h2oModels = []struct {
	name     string
	category string
}{
	{"bin_gbm", "Binomial"},
	{"bin_glm", "Binomial"},
	{"bin_dpl", "Binomial"},
	{"mul_gbm", "Multinomial"},
	{"mul_dpl", "Multinomial"},
	{"reg_gbm", "Regression"},
	{"reg_glm", "Regression"},
	{"reg_dpl", "Regression"},
}

type test struct {
	t   *testing.T
	svc web.Service
	dir az.Directory
	su  az.Principal
}

var clusterAddress, workingDirectory, compilationServiceAddress string

func init() {
	flag.StringVar(&clusterAddress, "cluster-address", "localhost:54321", "Where the h2o cluster can be reached.")
	flag.StringVar(&workingDirectory, "working-directory", "", "Where the var folder will be located.")
	flag.StringVar(&compilationServiceAddress, "compilation-service-address", ":8080", "Where to find the compilation service.")
}

func newTest(t *testing.T) *test {
	dbOpts := driverDBOpts{
		data.Connection{DbName: "steam", User: "steam", SSLMode: "disable"},
		superuser,
		superuser,
	}

	// Truncate database tables

	if err := data.Destroy(dbOpts.Name, dbOpts.Username, dbOpts.SSLMode); err != nil {
		t.Fatalf("Failed truncating database: %s", err)
	}

	// Determine current directory

	wd, err := filepath.Abs(filepath.Dir(workingDirectory + "/"))
	if err != nil {
		t.Fatalf("Failed determining current directory: %s", err)
	}

	// Delete any remnant models in models directory

	if err := os.RemoveAll(path.Join(wd, "var/master/model")); err != nil {
		t.Fatalf("Failed removing old model directory: %v", err)
	}

	// Create service instance

	opts := driverOpts{
		path.Join(wd, fs.VarDir, "master"),
		":9001",
		compilationServiceAddress,
		"",
		[2]int{1025, 65535},
		driverYarnOpts{false, "", ""},
		dbOpts,
	}
	svc, dir, err := newService(opts)
	if err != nil {
		t.Fatal(err)
	}

	// Lookup superuser

	su, err := dir.Lookup(superuser)
	if err != nil {
		t.Fatal(err)
	}

	return &test{t, svc, dir, su}
}

func (t *test) log(args ...interface{}) {
	t.t.Log(args...)
}

func (t *test) fail(format string, args ...interface{}) {
	debug.PrintStack()
	t.t.Fatalf(format, args...)
}

func (t *test) nil(arg interface{}) {
	if arg != nil {
		t.fail("unexpected: %s", arg)
	}
}

func (t *test) notnil(arg interface{}) {
	if arg == nil {
		t.fail("unexpected nil: %s", arg)
	}
}

func (t *test) ok(condition bool, format string, args ...interface{}) {
	if !condition {
		t.fail(format, args...)
	}
}

func buildPermissionMap(t *test) map[string]int64 {
	permissions, err := t.svc.GetAllPermissions(t.su)
	t.nil(err)

	permissionMap := make(map[string]int64)
	for _, permission := range permissions {
		permissionMap[permission.Code] = permission.Id
	}

	return permissionMap
}

func buildEntityTypeMap(t *test) map[string]int64 {
	entityTypes, err := t.svc.GetAllEntityTypes(t.su)
	t.nil(err)

	entityTypeMap := make(map[string]int64)
	for _, entityType := range entityTypes {
		entityTypeMap[entityType.Name] = entityType.Id
	}
	return entityTypeMap
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/justincampbell/tmux-pomodoro/tmux"
)

const timeFormat = time.RFC3339

var duration, _ = time.ParseDuration("25m")
var noTime time.Time

const usage = `
github.com/justincampbell/tmux-pomodoro

  pomodoro start   Start a timer for 25 minutes
  pomodoro status  Show the remaining time, or an exclamation point if done
  pomodoro clear   Clear the timer
`
const version = "v1.2.0"

// State is the state of the world passed through the functions to determine
// side-effects.
type State struct {
	endTime time.Time
	now     time.Time
}

// Output has fields for functions to set/append when they intend to output to
// the user.
type Output struct {
	text       string
	returnCode int
}

func init() {
	flag.Usage = func() {
		fmt.Printf("tmux-pomodoro %s\n", version)
		fmt.Printf("%s\n", strings.TrimSpace(usage))
	}

	flag.Parse()
}

func main() {
	state := State{
		endTime: readExistingTime(),
		now:     time.Now(),
	}

	args := flag.Args()
	var command string
	if len(args) == 0 {
		command = ""
	} else {
		command = args[0]
	}

	newState, output := parseCommand(state, command)

	if newState.endTime != state.endTime {
		writeTime(newState.endTime)
	}

	if output.text != "" {
		fmt.Println(output.text)
	}

	if output.returnCode != 0 {
		os.Exit(output.returnCode)
	}
}

func refreshTmux() {
	_ = tmux.RefreshClient("-S")
}

func parseCommand(state State, command string) (newState State, output Output) {
	newState = state

	switch command {
	case "start":
		newState.endTime = state.now.Add(duration)
		output.text = "Timer started, 25 minutes remaining"
		killRunningBeepers()
		_ = startBeeper()
		refreshTmux()
	case "status":
		if state.endTime == noTime {
			return
		}
		output.text = formatRemainingTime(state.endTime, state.now) + " 🍅 "
	case "clear":
		newState.endTime = noTime
		output.text = "Pomodoro cleared!"
		killRunningBeepers()
		refreshTmux()
	case "beep":
		<-time.NewTicker(duration).C
		_ = tmux.DisplayMessage("Pomodoro done, take a break!")
		refreshTmux()
	case "":
		flag.Usage()
	default:
		flag.Usage()
		output.returnCode = 1
	}

	return
}

func startBeeper() (err error) {
	command := exec.Command("pomodoro", "beep")
	err = command.Start()
	if err != nil {
		log.Println(err)
		return
	}
	bytes := []byte(strconv.Itoa(command.Process.Pid))
	err = ioutil.WriteFile(pidFilePath(), bytes, 0644)
	if err != nil {
		log.Println(err)
	}

	return
}

func killRunningBeepers() {
	bytes, err := ioutil.ReadFile(pidFilePath())
	if err != nil {
		return
	}
	pid, err := strconv.Atoi(string(bytes[:]))
	if err != nil {
		return
	}
	process, err := os.FindProcess(pid)
	if err != nil {
		return
	}
	_ = process.Kill()
}

func formatRemainingTime(existingTime time.Time, now time.Time) string {
	remaining := existingTime.Sub(now)
	remainingMinutes := remaining.Minutes()

	if remainingMinutes >= 0 {
		return strconv.FormatFloat(remainingMinutes, 'f', 0, 64)
	}

	return "❗"
}

func writeTime(t time.Time) {
	var bytes []byte
	if t != noTime {
		bytes = []byte(t.Format(timeFormat))
	}
	err := ioutil.WriteFile(filePath(), bytes, 0644)
	if err != nil {
		log.Fatal(err)
	}
}

func readExistingTime() time.Time {
	bytes, err := ioutil.ReadFile(filePath())
	if err != nil {
		return noTime
	}

	contents := string(bytes[:])
	contents = strings.TrimSpace(contents)

	result, err := time.Parse(timeFormat, contents)
	if err != nil {
		return noTime
	}

	return result
}

func filePath() string {
	return homeDir() + "/.pomodoro"
}

func pidFilePath() string {
	return homeDir() + "/.pomodoro.pid"
}

func homeDir() string {
	return os.Getenv("HOME")
}
<commit_msg>Bump version to v1.2.1<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/justincampbell/tmux-pomodoro/tmux"
)

const timeFormat = time.RFC3339

var duration, _ = time.ParseDuration("25m")
var noTime time.Time

const usage = `
github.com/justincampbell/tmux-pomodoro

  pomodoro start   Start a timer for 25 minutes
  pomodoro status  Show the remaining time, or an exclamation point if done
  pomodoro clear   Clear the timer
`
const version = "v1.2.1"

// State is the state of the world passed through the functions to determine
// side-effects.
type State struct {
	endTime time.Time
	now     time.Time
}

// Output has fields for functions to set/append when they intend to output to
// the user.
type Output struct {
	text       string
	returnCode int
}

func init() {
	flag.Usage = func() {
		fmt.Printf("tmux-pomodoro %s\n", version)
		fmt.Printf("%s\n", strings.TrimSpace(usage))
	}

	flag.Parse()
}

func main() {
	state := State{
		endTime: readExistingTime(),
		now:     time.Now(),
	}

	args := flag.Args()
	var command string
	if len(args) == 0 {
		command = ""
	} else {
		command = args[0]
	}

	newState, output := parseCommand(state, command)

	if newState.endTime != state.endTime {
		writeTime(newState.endTime)
	}

	if output.text != "" {
		fmt.Println(output.text)
	}

	if output.returnCode != 0 {
		os.Exit(output.returnCode)
	}
}

func refreshTmux() {
	_ = tmux.RefreshClient("-S")
}

func parseCommand(state State, command string) (newState State, output Output) {
	newState = state

	switch command {
	case "start":
		newState.endTime = state.now.Add(duration)
		output.text = "Timer started, 25 minutes remaining"
		killRunningBeepers()
		_ = startBeeper()
		refreshTmux()
	case "status":
		if state.endTime == noTime {
			return
		}
		output.text = formatRemainingTime(state.endTime, state.now) + " 🍅 "
	case "clear":
		newState.endTime = noTime
		output.text = "Pomodoro cleared!"
		killRunningBeepers()
		refreshTmux()
	case "beep":
		<-time.NewTicker(duration).C
		_ = tmux.DisplayMessage("Pomodoro done, take a break!")
		refreshTmux()
	case "":
		flag.Usage()
	default:
		flag.Usage()
		output.returnCode = 1
	}

	return
}

func startBeeper() (err error) {
	command := exec.Command("pomodoro", "beep")
	err = command.Start()
	if err != nil {
		log.Println(err)
		return
	}
	bytes := []byte(strconv.Itoa(command.Process.Pid))
	err = ioutil.WriteFile(pidFilePath(), bytes, 0644)
	if err != nil {
		log.Println(err)
	}

	return
}

func killRunningBeepers() {
	bytes, err := ioutil.ReadFile(pidFilePath())
	if err != nil {
		return
	}
	pid, err := strconv.Atoi(string(bytes[:]))
	if err != nil {
		return
	}
	process, err := os.FindProcess(pid)
	if err != nil {
		return
	}
	_ = process.Kill()
}

func formatRemainingTime(existingTime time.Time, now time.Time) string {
	remaining := existingTime.Sub(now)
	remainingMinutes := remaining.Minutes()

	if remainingMinutes >= 0 {
		return strconv.FormatFloat(remainingMinutes, 'f', 0, 64)
	}

	return "❗"
}

func writeTime(t time.Time) {
	var bytes []byte
	if t != noTime {
		bytes = []byte(t.Format(timeFormat))
	}
	err := ioutil.WriteFile(filePath(), bytes, 0644)
	if err != nil {
		log.Fatal(err)
	}
}

func readExistingTime() time.Time {
	bytes, err := ioutil.ReadFile(filePath())
	if err != nil {
		return noTime
	}

	contents := string(bytes[:])
	contents = strings.TrimSpace(contents)

	result, err := time.Parse(timeFormat, contents)
	if err != nil {
		return noTime
	}

	return result
}

func filePath() string {
	return homeDir() + "/.pomodoro"
}

func pidFilePath() string {
	return homeDir() + "/.pomodoro.pid"
}

func homeDir() string {
	return os.Getenv("HOME")
}
<|endoftext|>


<commit_before>// pormodoro timer
package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strconv"
	"time"
)

const defaultDuration = 25 * time.Minute

var silence = flag.Bool("silence", false, "")

func init() {
	const usage = `Usage of pomodoro:

	pomodoro [-silence] [duration]

Duration defaults to %d minutes. Durations may be expressed as integer minutes
(e.g. "15") or time with units (e.g. "1m30s" or "90s").

Chimes system bell at the end of the timer, unless -silence is set.
`
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, int(defaultDuration/time.Minute))
	}
	flag.Parse()
}

func waitDuration() (time.Duration, error) {
	if len(flag.Args()) > 1 {
		return 0, errors.New("Too many args...")
	}

	if flag.Arg(0) == "" {
		return defaultDuration, nil
	}

	if n, err := strconv.Atoi(flag.Arg(0)); err == nil {
		return time.Duration(n) * time.Minute, nil
	}

	if d, err := time.ParseDuration(flag.Arg(0)); err == nil {
		return d, nil
	}

	return 0, errors.New("Couldn't parse duration...")
}

func main() {
	wait, err := waitDuration()
	if err != nil {
		flag.Usage()
		os.Exit(2)
	}

	fmt.Printf("Start timer for %s.\n\n", wait)
	doneCh := time.After(wait)
	doneT := time.Now().Add(wait)

	go func() {
		for range time.Tick(200 * time.Millisecond) {
			duration := -time.Since(doneT)
			fmt.Printf("Countdown: %02d:%02d    \r",
				int(duration.Minutes()), int(duration.Seconds())%60)
		}
	}()
	<-doneCh

	if !*silence {
		fmt.Println("\a") // \a is the bell literal.
	} else {
		fmt.Println()
	}
}
<commit_msg>Different formats for different initial timers<commit_after>// pormodoro timer
package main

import (
	"errors"
	"flag"
	"fmt"
	"math"
	"os"
	"strconv"
	"time"
)

const defaultDuration = 25 * time.Minute

var silence = flag.Bool("silence", false, "")

func init() {
	const usage = `Usage of pomodoro:

	pomodoro [-silence] [duration]

Duration defaults to %d minutes. Durations may be expressed as integer minutes
(e.g. "15") or time with units (e.g. "1m30s" or "90s").

Chimes system bell at the end of the timer, unless -silence is set.
`
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, int(defaultDuration/time.Minute))
	}
	flag.Parse()
}

func waitDuration() (time.Duration, error) {
	if len(flag.Args()) > 1 {
		return 0, errors.New("Too many args...")
	}

	if flag.Arg(0) == "" {
		return defaultDuration, nil
	}

	if n, err := strconv.Atoi(flag.Arg(0)); err == nil {
		return time.Duration(n) * time.Minute, nil
	}

	if d, err := time.ParseDuration(flag.Arg(0)); err == nil {
		return d, nil
	}

	return 0, errors.New("Couldn't parse duration...")
}

func countdown(target time.Time, formatDuration func(time.Duration) string) {
	for range time.Tick(100 * time.Millisecond) {
		fmt.Print(formatDuration(-time.Since(target)))
	}
}

func formatDays(timeLeft time.Duration) string {
	days := int(timeLeft.Hours() / 24)
	hours := int(timeLeft.Hours()) % 24
	minutes := int(timeLeft.Minutes()) % 60
	seconds := int(timeLeft.Seconds()) % 60
	return fmt.Sprintf("Countdown: %d:%02d:%02d:%02d \r",
		days, hours, minutes, seconds)
}

func formatHours(timeLeft time.Duration) string {
	hours := int(timeLeft.Hours())
	minutes := int(timeLeft.Minutes()) % 60
	seconds := int(timeLeft.Seconds()) % 60
	return fmt.Sprintf("Countdown: %d:%02d:%02d\r",
		hours, minutes, seconds)
}

func formatMinutes(timeLeft time.Duration) string {
	minutes := int(timeLeft.Minutes())
	seconds := int(timeLeft.Seconds()) % 60
	return fmt.Sprintf("Countdown: %d:%02d\r", minutes, seconds)
}

func formatSeconds(timeLeft time.Duration) string {
	return fmt.Sprintf("Countdown: %02.1f \r", math.Abs(timeLeft.Seconds()))
}

func main() {
	wait, err := waitDuration()
	if err != nil {
		flag.Usage()
		os.Exit(2)
	}

	fmt.Printf("Start timer for %s.\n\n", wait)
	doneCh := time.After(wait)
	doneT := time.Now().Add(wait)

	switch {
	case wait >= 24*time.Hour:
		go countdown(doneT, formatDays)
	case wait >= time.Hour:
		go countdown(doneT, formatHours)
	case wait >= time.Minute:
		go countdown(doneT, formatMinutes)
	default:
		go countdown(doneT, formatSeconds)
	}

	<-doneCh

	if !*silence {
		fmt.Println("\a") // \a is the bell literal.
	} else {
		fmt.Println()
	}
}
<|endoftext|>


<commit_before>package segmenter

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strconv"
	"time"

	"github.com/ethereum/go-ethereum/logger/glog"
	"github.com/kz26/m3u8"
	"github.com/livepeer/lpms/stream"
	"github.com/nareix/joy4/av"
	"github.com/nareix/joy4/format/rtmp"
)

type SegmenterOptions struct {
	EnforceKeyframe bool //Enforce each segment starts with a keyframe
	SegLength       time.Duration
}

type VideoSegment struct {
	Codec  av.CodecType
	Format stream.VideoFormat
	Length time.Duration
	Data   []byte
	Name   string
}

type VideoPlaylist struct {
	Format stream.VideoFormat
	// Data   []byte
	Data *m3u8.MediaPlaylist
}

type VideoSegmenter interface{}

//FFMpegVideoSegmenter segments a RTMP stream by invoking FFMpeg and monitoring the file system.
type FFMpegVideoSegmenter struct {
	WorkDir      string
	LocalRtmpUrl string
	StrmID       string
	curSegment   int
	curPlaylist  *m3u8.MediaPlaylist
	curWaitTime  time.Duration
	SegLen       time.Duration
}

func NewFFMpegVideoSegmenter(workDir string, strmID string, localRtmpUrl string, segLen time.Duration) *FFMpegVideoSegmenter {
	return &FFMpegVideoSegmenter{WorkDir: workDir, StrmID: strmID, LocalRtmpUrl: localRtmpUrl, SegLen: segLen}
}

//RTMPToHLS invokes the FFMpeg command to do the segmenting.  This method blocks unless killed.
func (s *FFMpegVideoSegmenter) RTMPToHLS(ctx context.Context, opt SegmenterOptions) error {
	//Set up local workdir
	if _, err := os.Stat(s.WorkDir); os.IsNotExist(err) {
		err := os.Mkdir(s.WorkDir, 0700)
		if err != nil {
			return err
		}
	}

	//Test to make sure local RTMP is running.
	rtmpMux, err := rtmp.Dial(s.LocalRtmpUrl)
	if err != nil {
		glog.Errorf("Video Segmenter Error: %v.  Make sure local RTMP stream is available for segmenter.", err)
		rtmpMux.Close()
		return err
	}
	rtmpMux.Close()

	//Invoke the FFMpeg command
	// fmt.Println("ffmpeg", "-i", fmt.Sprintf("rtmp://localhost:%v/stream/%v", "1935", "test"), "-vcodec", "copy", "-acodec", "copy", "-bsf:v", "h264_mp4toannexb", "-f", "segment", "-muxdelay", "0", "-segment_list", "./tmp/stream.m3u8", "./tmp/stream_%d.ts")
	plfn := fmt.Sprintf("%s/%s.m3u8", s.WorkDir, s.StrmID)
	tsfn := s.WorkDir + "/" + s.StrmID + "_%d.ts"

	//This command needs to be manually killed, because ffmpeg doesn't seem to quit after getting a rtmp EOF
	cmd := exec.Command("ffmpeg", "-i", s.LocalRtmpUrl, "-vcodec", "copy", "-acodec", "copy", "-bsf:v", "h264_mp4toannexb", "-f", "segment", "-muxdelay", "0", "-segment_list", plfn, tsfn)
	err = cmd.Start()
	if err != nil {
		glog.Errorf("Cannot start ffmpeg command.")
		return err
	}

	ec := make(chan error, 1)
	go func() { ec <- cmd.Wait() }()

	select {
	case ffmpege := <-ec:
		glog.Errorf("Error from ffmpeg: %v", ffmpege)
		return ffmpege
	case <-ctx.Done():
		//Can't close RTMP server, joy4 doesn't support it.
		//server.Stop()
		cmd.Process.Kill()
		return ctx.Err()
	}
}

//PollSegment monitors the filesystem and returns a new segment as it becomes available
func (s *FFMpegVideoSegmenter) PollSegment(ctx context.Context) (*VideoSegment, error) {
	var length time.Duration
	tsfn := s.WorkDir + "/" + s.StrmID + "_" + strconv.Itoa(s.curSegment) + ".ts"
	seg, err := pollSegment(ctx, tsfn, time.Millisecond*100, s.SegLen)
	if err != nil {
		return nil, err
	}

	name := s.StrmID + "_" + strconv.Itoa(s.curSegment) + ".ts"
	if s.curPlaylist != nil && s.curPlaylist.Segments[s.curSegment] != nil {
		//This is ridiculous - but it's how we can round floats in Go
		sec, _ := strconv.Atoi(fmt.Sprintf("%.0f", s.curPlaylist.Segments[s.curSegment].Duration))
		length = time.Duration(sec) * 1000 * time.Millisecond
	}

	s.curSegment = s.curSegment + 1
	glog.Infof("Segment: %v, len:%v", name, len(seg))
	return &VideoSegment{Codec: av.H264, Format: stream.HLS, Length: length, Data: seg, Name: name}, err
}

//PollPlaylist monitors the filesystem and returns a new playlist as it becomes available
func (s *FFMpegVideoSegmenter) PollPlaylist(ctx context.Context) (*VideoPlaylist, error) {
	plfn := fmt.Sprintf("%s/%s.m3u8", s.WorkDir, s.StrmID)
	var lastPl []byte
	if s.curPlaylist == nil {
		lastPl = nil
	} else {
		lastPl = s.curPlaylist.Encode().Bytes()
	}

	pl, err := pollPlaylist(ctx, plfn, time.Millisecond*100, lastPl)
	if err != nil {
		return nil, err
	}

	p, err := m3u8.NewMediaPlaylist(50000, 50000)
	err = p.DecodeFrom(bytes.NewReader(pl), true)
	if err != nil {
		return nil, err
	}

	s.curPlaylist = p
	return &VideoPlaylist{Format: stream.HLS, Data: p}, err
}

func pollPlaylist(ctx context.Context, fn string, sleepTime time.Duration, lastFile []byte) (f []byte, err error) {
	for {
		if _, err := os.Stat(fn); err == nil {
			if err != nil {
				return nil, err
			}

			content, err := ioutil.ReadFile(fn)
			if err != nil {
				return nil, err
			}

			//The m3u8 package has some bugs, so the translation isn't 100% correct...
			p, err := m3u8.NewMediaPlaylist(50000, 50000)
			err = p.DecodeFrom(bytes.NewReader(content), true)
			if err != nil {
				return nil, err
			}
			curFile := p.Encode().Bytes()

			// fmt.Printf("p.Segments: %v\n", p.Segments[0])
			// fmt.Printf("lf: %s \ncf: %s \ncomp:%v\n\n", lastFile, curFile, bytes.Compare(lastFile, curFile))
			if lastFile == nil || bytes.Compare(lastFile, curFile) != 0 {
				return content, nil
			}
		}

		select {
		case <-ctx.Done():
			fmt.Println("ctx.Done()!!!")
			return nil, ctx.Err()
		default:
		}

		time.Sleep(sleepTime)
	}

}

func pollSegment(ctx context.Context, fn string, sleepTime time.Duration, segLen time.Duration) (f []byte, err error) {
	for {
		if _, err := os.Stat(fn); err == nil {
			// fmt.Printf("FileName: %v, FileSize: %v \n\n", fn, info.Size())
			time.Sleep(segLen)
			// fmt.Printf("FileName: %v, FileSize: %v \n\n", fn, info.Size())
			content, err := ioutil.ReadFile(fn)
			if err != nil {
				return nil, err
			}
			return content, err
		}

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		time.Sleep(sleepTime)
	}
}
<commit_msg>Removing go-ethereum/logger/glog dependency<commit_after>package segmenter

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strconv"
	"time"

	"github.com/golang/glog"
	"github.com/kz26/m3u8"
	"github.com/livepeer/lpms/stream"
	"github.com/nareix/joy4/av"
	"github.com/nareix/joy4/format/rtmp"
)

type SegmenterOptions struct {
	EnforceKeyframe bool //Enforce each segment starts with a keyframe
	SegLength       time.Duration
}

type VideoSegment struct {
	Codec  av.CodecType
	Format stream.VideoFormat
	Length time.Duration
	Data   []byte
	Name   string
}

type VideoPlaylist struct {
	Format stream.VideoFormat
	// Data   []byte
	Data *m3u8.MediaPlaylist
}

type VideoSegmenter interface{}

//FFMpegVideoSegmenter segments a RTMP stream by invoking FFMpeg and monitoring the file system.
type FFMpegVideoSegmenter struct {
	WorkDir      string
	LocalRtmpUrl string
	StrmID       string
	curSegment   int
	curPlaylist  *m3u8.MediaPlaylist
	curWaitTime  time.Duration
	SegLen       time.Duration
}

func NewFFMpegVideoSegmenter(workDir string, strmID string, localRtmpUrl string, segLen time.Duration) *FFMpegVideoSegmenter {
	return &FFMpegVideoSegmenter{WorkDir: workDir, StrmID: strmID, LocalRtmpUrl: localRtmpUrl, SegLen: segLen}
}

//RTMPToHLS invokes the FFMpeg command to do the segmenting.  This method blocks unless killed.
func (s *FFMpegVideoSegmenter) RTMPToHLS(ctx context.Context, opt SegmenterOptions) error {
	//Set up local workdir
	if _, err := os.Stat(s.WorkDir); os.IsNotExist(err) {
		err := os.Mkdir(s.WorkDir, 0700)
		if err != nil {
			return err
		}
	}

	//Test to make sure local RTMP is running.
	rtmpMux, err := rtmp.Dial(s.LocalRtmpUrl)
	if err != nil {
		glog.Errorf("Video Segmenter Error: %v.  Make sure local RTMP stream is available for segmenter.", err)
		rtmpMux.Close()
		return err
	}
	rtmpMux.Close()

	//Invoke the FFMpeg command
	// fmt.Println("ffmpeg", "-i", fmt.Sprintf("rtmp://localhost:%v/stream/%v", "1935", "test"), "-vcodec", "copy", "-acodec", "copy", "-bsf:v", "h264_mp4toannexb", "-f", "segment", "-muxdelay", "0", "-segment_list", "./tmp/stream.m3u8", "./tmp/stream_%d.ts")
	plfn := fmt.Sprintf("%s/%s.m3u8", s.WorkDir, s.StrmID)
	tsfn := s.WorkDir + "/" + s.StrmID + "_%d.ts"

	//This command needs to be manually killed, because ffmpeg doesn't seem to quit after getting a rtmp EOF
	cmd := exec.Command("ffmpeg", "-i", s.LocalRtmpUrl, "-vcodec", "copy", "-acodec", "copy", "-bsf:v", "h264_mp4toannexb", "-f", "segment", "-muxdelay", "0", "-segment_list", plfn, tsfn)
	err = cmd.Start()
	if err != nil {
		glog.Errorf("Cannot start ffmpeg command.")
		return err
	}

	ec := make(chan error, 1)
	go func() { ec <- cmd.Wait() }()

	select {
	case ffmpege := <-ec:
		glog.Errorf("Error from ffmpeg: %v", ffmpege)
		return ffmpege
	case <-ctx.Done():
		//Can't close RTMP server, joy4 doesn't support it.
		//server.Stop()
		cmd.Process.Kill()
		return ctx.Err()
	}
}

//PollSegment monitors the filesystem and returns a new segment as it becomes available
func (s *FFMpegVideoSegmenter) PollSegment(ctx context.Context) (*VideoSegment, error) {
	var length time.Duration
	tsfn := s.WorkDir + "/" + s.StrmID + "_" + strconv.Itoa(s.curSegment) + ".ts"
	seg, err := pollSegment(ctx, tsfn, time.Millisecond*100, s.SegLen)
	if err != nil {
		return nil, err
	}

	name := s.StrmID + "_" + strconv.Itoa(s.curSegment) + ".ts"
	if s.curPlaylist != nil && s.curPlaylist.Segments[s.curSegment] != nil {
		//This is ridiculous - but it's how we can round floats in Go
		sec, _ := strconv.Atoi(fmt.Sprintf("%.0f", s.curPlaylist.Segments[s.curSegment].Duration))
		length = time.Duration(sec) * 1000 * time.Millisecond
	}

	s.curSegment = s.curSegment + 1
	glog.Infof("Segment: %v, len:%v", name, len(seg))
	return &VideoSegment{Codec: av.H264, Format: stream.HLS, Length: length, Data: seg, Name: name}, err
}

//PollPlaylist monitors the filesystem and returns a new playlist as it becomes available
func (s *FFMpegVideoSegmenter) PollPlaylist(ctx context.Context) (*VideoPlaylist, error) {
	plfn := fmt.Sprintf("%s/%s.m3u8", s.WorkDir, s.StrmID)
	var lastPl []byte
	if s.curPlaylist == nil {
		lastPl = nil
	} else {
		lastPl = s.curPlaylist.Encode().Bytes()
	}

	pl, err := pollPlaylist(ctx, plfn, time.Millisecond*100, lastPl)
	if err != nil {
		return nil, err
	}

	p, err := m3u8.NewMediaPlaylist(50000, 50000)
	err = p.DecodeFrom(bytes.NewReader(pl), true)
	if err != nil {
		return nil, err
	}

	s.curPlaylist = p
	return &VideoPlaylist{Format: stream.HLS, Data: p}, err
}

func pollPlaylist(ctx context.Context, fn string, sleepTime time.Duration, lastFile []byte) (f []byte, err error) {
	for {
		if _, err := os.Stat(fn); err == nil {
			if err != nil {
				return nil, err
			}

			content, err := ioutil.ReadFile(fn)
			if err != nil {
				return nil, err
			}

			//The m3u8 package has some bugs, so the translation isn't 100% correct...
			p, err := m3u8.NewMediaPlaylist(50000, 50000)
			err = p.DecodeFrom(bytes.NewReader(content), true)
			if err != nil {
				return nil, err
			}
			curFile := p.Encode().Bytes()

			// fmt.Printf("p.Segments: %v\n", p.Segments[0])
			// fmt.Printf("lf: %s \ncf: %s \ncomp:%v\n\n", lastFile, curFile, bytes.Compare(lastFile, curFile))
			if lastFile == nil || bytes.Compare(lastFile, curFile) != 0 {
				return content, nil
			}
		}

		select {
		case <-ctx.Done():
			fmt.Println("ctx.Done()!!!")
			return nil, ctx.Err()
		default:
		}

		time.Sleep(sleepTime)
	}

}

func pollSegment(ctx context.Context, fn string, sleepTime time.Duration, segLen time.Duration) (f []byte, err error) {
	for {
		if _, err := os.Stat(fn); err == nil {
			// fmt.Printf("FileName: %v, FileSize: %v \n\n", fn, info.Size())
			time.Sleep(segLen)
			// fmt.Printf("FileName: %v, FileSize: %v \n\n", fn, info.Size())
			content, err := ioutil.ReadFile(fn)
			if err != nil {
				return nil, err
			}
			return content, err
		}

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		time.Sleep(sleepTime)
	}
}
<|endoftext|>


<commit_before>package allyourbase

import (
	"fmt"
	"math"
)

func ConvertToBase(inputBase int, inputDigits []int, outputBase int) (outputDigits []int, e error) {
	if len(inputDigits) == 0 {
		return []int{0}, nil
	}
	base10 := getBase10Input(inputBase, inputDigits)
	for base10 > 0 {
		digit := base10 % outputBase
		outputDigits = append([]int{digit}, outputDigits...)
		base10 = base10 / outputBase
	}
	return outputDigits, nil
}

func getBase10Input(inputBase int, inputDigits []int) (base10Input int) {
	for i, digit := range reverse(inputDigits) {
		base10Input += powInt(inputBase, i) * digit
	}
	fmt.Printf("getBase10Input(%d, %v)=%d\n", inputBase, inputDigits, base10Input)
	return base10Input
}

func reverse(input []int) (reversed []int) {
	for i := len(input) - 1; i >= 0; i-- {
		reversed = append(reversed, input[i])
	}
	return reversed
}

func powInt(x, y int) int {
	return int(math.Pow(float64(x), float64(y)))
}
<commit_msg>Handle multiple zeroes<commit_after>package allyourbase

import (
	"fmt"
	"math"
)

func ConvertToBase(inputBase int, inputDigits []int, outputBase int) (outputDigits []int, e error) {
	base10 := getBase10Input(inputBase, inputDigits)
	if base10 == 0 {
		return []int{0}, nil
	}
	for base10 > 0 {
		digit := base10 % outputBase
		outputDigits = append([]int{digit}, outputDigits...)
		base10 = base10 / outputBase
	}
	return outputDigits, nil
}

func getBase10Input(inputBase int, inputDigits []int) (base10Input int) {
	for i, digit := range reverse(inputDigits) {
		base10Input += powInt(inputBase, i) * digit
	}
	fmt.Printf("getBase10Input(%d, %v)=%d\n", inputBase, inputDigits, base10Input)
	return base10Input
}

func reverse(input []int) (reversed []int) {
	for i := len(input) - 1; i >= 0; i-- {
		reversed = append(reversed, input[i])
	}
	return reversed
}

func powInt(x, y int) int {
	return int(math.Pow(float64(x), float64(y)))
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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 mysqlctl

import (
	"bytes"
	"fmt"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/log"
)

// Note that definitions of local_metadata and shard_metadata should be the same
// as in testing which is defined in config/init_db.sql.
const (
	sqlCreateLocalMetadataTable = `CREATE TABLE IF NOT EXISTS _vt.local_metadata (
  name VARCHAR(255) NOT NULL,
  value VARCHAR(255) NOT NULL,
  PRIMARY KEY (name)
  ) ENGINE=InnoDB`
	sqlCreateShardMetadataTable = `CREATE TABLE IF NOT EXISTS _vt.shard_metadata (
  name VARCHAR(255) NOT NULL,
  value MEDIUMBLOB NOT NULL,
  PRIMARY KEY (name)
  ) ENGINE=InnoDB`
	sqlUpdateLocalMetadataTable = "UPDATE _vt.local_metadata SET db_name='%s' WHERE db_name=''"
	sqlUpdateShardMetadataTable = "UPDATE _vt.shard_metadata SET db_name='%s' WHERE db_name=''"
)

var (
	sqlAlterLocalMetadataTable = []string{
		`ALTER TABLE _vt.local_metadata ADD COLUMN db_name VARBINARY(255) NOT NULL DEFAULT ''`,
		`ALTER TABLE _vt.local_metadata DROP PRIMARY KEY, ADD PRIMARY KEY(name, db_name)`,
	}
	sqlAlterShardMetadataTable = []string{
		`ALTER TABLE _vt.shard_metadata ADD COLUMN db_name VARBINARY(255) NOT NULL DEFAULT ''`,
		`ALTER TABLE _vt.shard_metadata DROP PRIMARY KEY, ADD PRIMARY KEY(name, db_name)`,
	}
)

// PopulateMetadataTables creates and fills the _vt.local_metadata table and
// creates _vt.shard_metadata table. _vt.local_metadata table is
// a per-tablet table that is never replicated. This allows queries
// against local_metadata to return different values on different tablets,
// which is used for communicating between Vitess and MySQL-level tools like
// Orchestrator (http://github.com/github/orchestrator).
// _vt.shard_metadata is a replicated table with per-shard information, but it's
// created here to make it easier to create it on databases that were running
// old version of Vitess, or databases that are getting converted to run under
// Vitess.
func PopulateMetadataTables(mysqld MysqlDaemon, localMetadata map[string]string, dbName string) error {
	log.Infof("Populating _vt.local_metadata table...")

	// Get a non-pooled DBA connection.
	conn, err := mysqld.GetDbaConnection()
	if err != nil {
		return err
	}
	defer conn.Close()

	// Disable replication on this session. We close the connection after using
	// it, so there's no need to re-enable replication when we're done.
	if _, err := conn.ExecuteFetch("SET @@session.sql_log_bin = 0", 0, false); err != nil {
		return err
	}

	// Create the database and table if necessary.
	if _, err := conn.ExecuteFetch("CREATE DATABASE IF NOT EXISTS _vt", 0, false); err != nil {
		return err
	}
	if _, err := conn.ExecuteFetch(sqlCreateLocalMetadataTable, 0, false); err != nil {
		return err
	}
	for _, sql := range sqlAlterLocalMetadataTable {
		if _, err := conn.ExecuteFetch(sql, 0, false); err != nil {
			if merr, ok := err.(*mysql.SQLError); ok && merr.Num == mysql.ERDupFieldName {
				log.Errorf("Expected error executing %v: %v", sql, err)
			} else {
				log.Errorf("Unexpected error executing %v: %v", sql, err)
				return err
			}
		}
	}
	if _, err := conn.ExecuteFetch(fmt.Sprintf(sqlUpdateLocalMetadataTable, dbName), 0, false); err != nil {
		return err
	}
	if _, err := conn.ExecuteFetch(sqlCreateShardMetadataTable, 0, false); err != nil {
		return err
	}
	for _, sql := range sqlAlterShardMetadataTable {
		if _, err := conn.ExecuteFetch(sql, 0, false); err != nil {
			if merr, ok := err.(*mysql.SQLError); ok && merr.Num == mysql.ERDupFieldName {
				log.Errorf("Expected error executing %v: %v", sql, err)
			} else {
				log.Errorf("Unexpected error executing %v: %v", sql, err)
				return err
			}
		}
	}
	if _, err := conn.ExecuteFetch(fmt.Sprintf(sqlUpdateShardMetadataTable, dbName), 0, false); err != nil {
		return err
	}

	// Populate local_metadata from the passed list of values.
	if _, err := conn.ExecuteFetch("BEGIN", 0, false); err != nil {
		return err
	}
	for name, val := range localMetadata {
		nameValue := sqltypes.NewVarChar(name)
		valValue := sqltypes.NewVarChar(val)
		dbNameValue := sqltypes.NewVarBinary(dbName)

		queryBuf := bytes.Buffer{}
		queryBuf.WriteString("INSERT INTO _vt.local_metadata (name,value, db_name) VALUES (")
		nameValue.EncodeSQL(&queryBuf)
		queryBuf.WriteByte(',')
		valValue.EncodeSQL(&queryBuf)
		queryBuf.WriteByte(',')
		dbNameValue.EncodeSQL(&queryBuf)
		queryBuf.WriteString(") ON DUPLICATE KEY UPDATE value = ")
		valValue.EncodeSQL(&queryBuf)

		if _, err := conn.ExecuteFetch(queryBuf.String(), 0, false); err != nil {
			return err
		}
	}
	_, err = conn.ExecuteFetch("COMMIT", 0, false)
	return err
}
<commit_msg>change fatal to error if we are unable to update local_metadata or shard_metadata<commit_after>/*
Copyright 2017 Google Inc.

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 mysqlctl

import (
	"bytes"
	"fmt"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/log"
)

// Note that definitions of local_metadata and shard_metadata should be the same
// as in testing which is defined in config/init_db.sql.
const (
	sqlCreateLocalMetadataTable = `CREATE TABLE IF NOT EXISTS _vt.local_metadata (
  name VARCHAR(255) NOT NULL,
  value VARCHAR(255) NOT NULL,
  PRIMARY KEY (name)
  ) ENGINE=InnoDB`
	sqlCreateShardMetadataTable = `CREATE TABLE IF NOT EXISTS _vt.shard_metadata (
  name VARCHAR(255) NOT NULL,
  value MEDIUMBLOB NOT NULL,
  PRIMARY KEY (name)
  ) ENGINE=InnoDB`
	sqlUpdateLocalMetadataTable = "UPDATE _vt.local_metadata SET db_name='%s' WHERE db_name=''"
	sqlUpdateShardMetadataTable = "UPDATE _vt.shard_metadata SET db_name='%s' WHERE db_name=''"
)

var (
	sqlAlterLocalMetadataTable = []string{
		`ALTER TABLE _vt.local_metadata ADD COLUMN db_name VARBINARY(255) NOT NULL DEFAULT ''`,
		`ALTER TABLE _vt.local_metadata DROP PRIMARY KEY, ADD PRIMARY KEY(name, db_name)`,
	}
	sqlAlterShardMetadataTable = []string{
		`ALTER TABLE _vt.shard_metadata ADD COLUMN db_name VARBINARY(255) NOT NULL DEFAULT ''`,
		`ALTER TABLE _vt.shard_metadata DROP PRIMARY KEY, ADD PRIMARY KEY(name, db_name)`,
	}
)

// PopulateMetadataTables creates and fills the _vt.local_metadata table and
// creates _vt.shard_metadata table. _vt.local_metadata table is
// a per-tablet table that is never replicated. This allows queries
// against local_metadata to return different values on different tablets,
// which is used for communicating between Vitess and MySQL-level tools like
// Orchestrator (http://github.com/github/orchestrator).
// _vt.shard_metadata is a replicated table with per-shard information, but it's
// created here to make it easier to create it on databases that were running
// old version of Vitess, or databases that are getting converted to run under
// Vitess.
func PopulateMetadataTables(mysqld MysqlDaemon, localMetadata map[string]string, dbName string) error {
	log.Infof("Populating _vt.local_metadata table...")

	// Get a non-pooled DBA connection.
	conn, err := mysqld.GetDbaConnection()
	if err != nil {
		return err
	}
	defer conn.Close()

	// Disable replication on this session. We close the connection after using
	// it, so there's no need to re-enable replication when we're done.
	if _, err := conn.ExecuteFetch("SET @@session.sql_log_bin = 0", 0, false); err != nil {
		return err
	}

	// Create the database and table if necessary.
	if _, err := conn.ExecuteFetch("CREATE DATABASE IF NOT EXISTS _vt", 0, false); err != nil {
		return err
	}
	if _, err := conn.ExecuteFetch(sqlCreateLocalMetadataTable, 0, false); err != nil {
		return err
	}
	for _, sql := range sqlAlterLocalMetadataTable {
		if _, err := conn.ExecuteFetch(sql, 0, false); err != nil {
			if merr, ok := err.(*mysql.SQLError); ok && merr.Num == mysql.ERDupFieldName {
				log.Errorf("Expected error executing %v: %v", sql, err)
			} else {
				log.Errorf("Unexpected error executing %v: %v", sql, err)
				return err
			}
		}
	}
	sql := fmt.Sprintf(sqlUpdateLocalMetadataTable, dbName)
	if _, err := conn.ExecuteFetch(sql, 0, false); err != nil {
		log.Errorf("unexpected error executing %v: %v, continuing. Please check the data in _vt.local_metadata and take corrective action", sql, err)
	}
	if _, err := conn.ExecuteFetch(sqlCreateShardMetadataTable, 0, false); err != nil {
		return err
	}
	for _, sql := range sqlAlterShardMetadataTable {
		if _, err := conn.ExecuteFetch(sql, 0, false); err != nil {
			if merr, ok := err.(*mysql.SQLError); ok && merr.Num == mysql.ERDupFieldName {
				log.Errorf("Expected error executing %v: %v", sql, err)
			} else {
				log.Errorf("Unexpected error executing %v: %v", sql, err)
				return err
			}
		}
	}
	sql = fmt.Sprintf(sqlUpdateShardMetadataTable, dbName)
	if _, err := conn.ExecuteFetch(sql, 0, false); err != nil {
		log.Errorf("unexpected error executing %v: %v, continuing. Please check the data in _vt.shard_metadata and take corrective action", sql, err)
	}

	// Populate local_metadata from the passed list of values.
	if _, err := conn.ExecuteFetch("BEGIN", 0, false); err != nil {
		return err
	}
	for name, val := range localMetadata {
		nameValue := sqltypes.NewVarChar(name)
		valValue := sqltypes.NewVarChar(val)
		dbNameValue := sqltypes.NewVarBinary(dbName)

		queryBuf := bytes.Buffer{}
		queryBuf.WriteString("INSERT INTO _vt.local_metadata (name,value, db_name) VALUES (")
		nameValue.EncodeSQL(&queryBuf)
		queryBuf.WriteByte(',')
		valValue.EncodeSQL(&queryBuf)
		queryBuf.WriteByte(',')
		dbNameValue.EncodeSQL(&queryBuf)
		queryBuf.WriteString(") ON DUPLICATE KEY UPDATE value = ")
		valValue.EncodeSQL(&queryBuf)

		if _, err := conn.ExecuteFetch(queryBuf.String(), 0, false); err != nil {
			return err
		}
	}
	_, err = conn.ExecuteFetch("COMMIT", 0, false)
	return err
}
<|endoftext|>