<commit_before>package data

// this package uses panic, as data validation is very important here.
// if something goes wrong, panic is safer than getting corrupted data.

import "unicode/utf8"

func runeToBytes(r rune) []byte {
	bs := make([]byte, utf8.RuneLen(r))
	utf8.EncodeRune(bs, r)
	return bs
}

type Clip struct {
	data     []byte
	newlines []int
}

func DataClip(data []byte) Clip {
	newlines := []int{}
	for i, b := range data {
		if b == '\n' {
			newlines = append(newlines, i)
		}
	}
	return Clip{
		data:     data,
		newlines: newlines,
	}
}

func Clips(datas ...[]byte) []Clip {
	clips := make([]Clip, 0)
	for _, data := range datas {
		c := DataClip(data)
		clips = append(clips, c)
	}
	return clips
}

func (c Clip) Len() int {
	return len(c.data)
}

func cut(c Clip, o int) (a, b Clip) {
	aNewlines := make([]int, 0)
	bNewlines := make([]int, 0)
	for _, n := range c.newlines {
		if o < n {
			aNewlines = append(aNewlines, n)
		} else {
			bNewlines = append(bNewlines, n-o)
		}
	}
	a = Clip{data: c.data[:o], newlines: aNewlines}
	b = Clip{data: c.data[o:], newlines: bNewlines}
	return a, b
}

func (c Clip) Cut(o int) (a, b Clip) {
	aNewlines := make([]int, 0)
	bNewlines := make([]int, 0)
	for _, n := range c.newlines {
		if o < n {
			aNewlines = append(aNewlines, n)
		} else {
			bNewlines = append(bNewlines, n-o)
		}
	}
	a = Clip{data: c.data[:o], newlines: aNewlines}
	b = Clip{data: c.data[o:], newlines: bNewlines}
	return a, b
}

func (c Clip) Append(r rune) Clip {
	if r == '\n' {
		c.newlines = append(c.newlines, len(c.data))
	}
	c.data = append(c.data, runeToBytes(r)...)
	return c
}

type Cursor struct {
	clips []Clip

	i int // clip index
	o int // byte offset on the clip

	appending bool
}

func NewCursor(clips []Clip) *Cursor {
	return &Cursor{clips: clips}
}

func nextOffset(data []byte, o int) int {
	remain := data[o:]
	r, n := utf8.DecodeRune(remain)
	remain = remain[n:]
	if r == '\r' {
		r, _ := utf8.DecodeRune(remain)
		if r == '\n' {
			n += 1
		}
	}
	o += n
	if o == len(data) {
		return -1
	}
	return o
}

func prevOffset(data []byte, o int) int {
	if o == 0 {
		return -1
	}
	remain := data[:o]
	r, n := utf8.DecodeLastRune(remain)
	remain = remain[:len(remain)-n]
	if r == '\n' {
		r, _ := utf8.DecodeLastRune(remain)
		if r == '\r' {
			n += 1
		}
	}
	return o - n
}

func (c *Cursor) MoveNext() {
	c.appending = false
	if c.i == len(c.clips) {
		if c.o != 0 {
			panic("c.o should 0 when c.i == len(c.clips)")
		}
		return
	}
	o := nextOffset(c.clips[c.i].data, c.o)
	if o == -1 {
		c.i++
		c.o = 0
		return
	}
	c.o = o
}

func (c *Cursor) MovePrev() {
	c.appending = false
	if c.i == 0 && c.o == 0 {
		return
	}
	if c.i == len(c.clips) {
		if c.o != 0 {
			panic("c.o should 0 when c.i == len(c.clips)")
		}
		c.i--
		c.o = len(c.clips[c.i].data)
	}
	if c.o == 0 {
		c.i--
		c.o = prevOffset(c.clips[c.i].data, len(c.clips[c.i].data))
		return
	}
	c.o = prevOffset(c.clips[c.i].data, c.o)
}

func (c *Cursor) Move(o int) {
	if o == 0 {
		return
	}
	if o > 0 {
		for i := 0; i < o; i++ {
			c.MoveNext()
		}
	} else {
		for i := 0; i > o; i-- {
			c.MovePrev()
		}
	}
}

func (c *Cursor) GotoStart() {
	c.appending = false
	c.i = 0
	c.o = 0
}

func (c *Cursor) GotoEnd() {
	c.appending = false
	c.i = len(c.clips)
	c.o = 0
}

func (c *Cursor) GotoNextLine() {
	c.appending = false
	if len(c.clips) == 0 {
		panic("length of clips should not be zero")
	}
	for {
		if c.i == len(c.clips) {
			if c.o != 0 {
				panic("c.o should 0 when c.i == len(c.clips)")
			}
			return
		}
		nls := c.clips[c.i].newlines
		for i := range nls {
			o := nls[i]
			if o <= c.o {
				continue
			}
			c.o = o
			return
		}
		c.i++
		c.o = 0
	}
}

func (c *Cursor) GotoPrevLine() {
	c.appending = false
	if len(c.clips) == 0 {
		panic("length of clips should not be zero")
	}
	if c.i == len(c.clips) {
		c.i--
		c.o = len(c.clips[c.i].data)
	}
	for {
		nls := c.clips[c.i].newlines
		for i := range nls {
			o := nls[len(nls)-1-i]
			if o >= c.o {
				continue
			}
			c.o = o
			return
		}
		if c.i == 0 {
			// no more previous clip
			c.o = 0
			return
		}
		c.i--
		c.o = len(c.clips[c.i].data)
	}
}

// Cut cuts the underlying clip it stands.
// If it is standing at edge of clip(s), it doesn't do anything.
func (c *Cursor) Cut() {
	if c.i == len(c.clips) {
		if c.o != 0 {
			panic("c.o should 0 when c.i == len(c.clips)")
		}
		return
	}
	if c.o == 0 {
		// edge
		return
	}
	clipA, clipB := c.clips[c.i].Cut(c.o)
	c.clips = append(c.clips[:c.i], append([]Clip{clipA, clipB}, c.clips[c.i+1:]...)...)
	c.i++
	c.o = 0
}

func (c *Cursor) Write(r rune) {
	if c.appending {
		if c.o != 0 {
			panic("c.o should 0 when appending")
		}
		c.clips[c.i-1] = c.clips[c.i-1].Append(r)
		return
	}
	c.appending = true
	c.Cut()
	clipInsert := DataClip(runeToBytes(r))
	c.clips = append(c.clips[:c.i], append([]Clip{clipInsert}, c.clips[c.i:]...)...)
	c.i++
	c.o = 0
}

func (c *Cursor) Delete() {
	c.appending = false
	if c.i == len(c.clips) {
		if c.o != 0 {
			panic("c.o should 0 when c.i == len(c.clips)")
		}
		return
	}
	c.Cut()
	p := nextOffset(c.clips[c.i].data, 0)
	if p == -1 {
		c.clips = append(c.clips[:c.i], c.clips[c.i+1:]...)
		return
	}
	_, c.clips[c.i] = c.clips[c.i].Cut(p)
}

func (c *Cursor) Backspace() {
	c.appending = false
	if c.i == 0 && c.o == 0 {
		return
	}
	c.Cut()
	p := prevOffset(c.clips[c.i-1].data, len(c.clips[c.i-1].data))
	if p == 0 {
		c.clips = append(c.clips[:c.i-1], c.clips[c.i:]...)
		c.i--
		return
	}
	c.clips[c.i-1], _ = c.clips[c.i-1].Cut(p)
}
<commit_msg>data: remove cut which is duplication of Clip.Cut<commit_after>package data

// this package uses panic, as data validation is very important here.
// if something goes wrong, panic is safer than getting corrupted data.

import "unicode/utf8"

func runeToBytes(r rune) []byte {
	bs := make([]byte, utf8.RuneLen(r))
	utf8.EncodeRune(bs, r)
	return bs
}

type Clip struct {
	data     []byte
	newlines []int
}

func DataClip(data []byte) Clip {
	newlines := []int{}
	for i, b := range data {
		if b == '\n' {
			newlines = append(newlines, i)
		}
	}
	return Clip{
		data:     data,
		newlines: newlines,
	}
}

func Clips(datas ...[]byte) []Clip {
	clips := make([]Clip, 0)
	for _, data := range datas {
		c := DataClip(data)
		clips = append(clips, c)
	}
	return clips
}

func (c Clip) Len() int {
	return len(c.data)
}

func (c Clip) Cut(o int) (a, b Clip) {
	aNewlines := make([]int, 0)
	bNewlines := make([]int, 0)
	for _, n := range c.newlines {
		if o < n {
			aNewlines = append(aNewlines, n)
		} else {
			bNewlines = append(bNewlines, n-o)
		}
	}
	a = Clip{data: c.data[:o], newlines: aNewlines}
	b = Clip{data: c.data[o:], newlines: bNewlines}
	return a, b
}

func (c Clip) Append(r rune) Clip {
	if r == '\n' {
		c.newlines = append(c.newlines, len(c.data))
	}
	c.data = append(c.data, runeToBytes(r)...)
	return c
}

type Cursor struct {
	clips []Clip

	i int // clip index
	o int // byte offset on the clip

	appending bool
}

func NewCursor(clips []Clip) *Cursor {
	return &Cursor{clips: clips}
}

func nextOffset(data []byte, o int) int {
	remain := data[o:]
	r, n := utf8.DecodeRune(remain)
	remain = remain[n:]
	if r == '\r' {
		r, _ := utf8.DecodeRune(remain)
		if r == '\n' {
			n += 1
		}
	}
	o += n
	if o == len(data) {
		return -1
	}
	return o
}

func prevOffset(data []byte, o int) int {
	if o == 0 {
		return -1
	}
	remain := data[:o]
	r, n := utf8.DecodeLastRune(remain)
	remain = remain[:len(remain)-n]
	if r == '\n' {
		r, _ := utf8.DecodeLastRune(remain)
		if r == '\r' {
			n += 1
		}
	}
	return o - n
}

func (c *Cursor) MoveNext() {
	c.appending = false
	if c.i == len(c.clips) {
		if c.o != 0 {
			panic("c.o should 0 when c.i == len(c.clips)")
		}
		return
	}
	o := nextOffset(c.clips[c.i].data, c.o)
	if o == -1 {
		c.i++
		c.o = 0
		return
	}
	c.o = o
}

func (c *Cursor) MovePrev() {
	c.appending = false
	if c.i == 0 && c.o == 0 {
		return
	}
	if c.i == len(c.clips) {
		if c.o != 0 {
			panic("c.o should 0 when c.i == len(c.clips)")
		}
		c.i--
		c.o = len(c.clips[c.i].data)
	}
	if c.o == 0 {
		c.i--
		c.o = prevOffset(c.clips[c.i].data, len(c.clips[c.i].data))
		return
	}
	c.o = prevOffset(c.clips[c.i].data, c.o)
}

func (c *Cursor) Move(o int) {
	if o == 0 {
		return
	}
	if o > 0 {
		for i := 0; i < o; i++ {
			c.MoveNext()
		}
	} else {
		for i := 0; i > o; i-- {
			c.MovePrev()
		}
	}
}

func (c *Cursor) GotoStart() {
	c.appending = false
	c.i = 0
	c.o = 0
}

func (c *Cursor) GotoEnd() {
	c.appending = false
	c.i = len(c.clips)
	c.o = 0
}

func (c *Cursor) GotoNextLine() {
	c.appending = false
	if len(c.clips) == 0 {
		panic("length of clips should not be zero")
	}
	for {
		if c.i == len(c.clips) {
			if c.o != 0 {
				panic("c.o should 0 when c.i == len(c.clips)")
			}
			return
		}
		nls := c.clips[c.i].newlines
		for i := range nls {
			o := nls[i]
			if o <= c.o {
				continue
			}
			c.o = o
			return
		}
		c.i++
		c.o = 0
	}
}

func (c *Cursor) GotoPrevLine() {
	c.appending = false
	if len(c.clips) == 0 {
		panic("length of clips should not be zero")
	}
	if c.i == len(c.clips) {
		c.i--
		c.o = len(c.clips[c.i].data)
	}
	for {
		nls := c.clips[c.i].newlines
		for i := range nls {
			o := nls[len(nls)-1-i]
			if o >= c.o {
				continue
			}
			c.o = o
			return
		}
		if c.i == 0 {
			// no more previous clip
			c.o = 0
			return
		}
		c.i--
		c.o = len(c.clips[c.i].data)
	}
}

// Cut cuts the underlying clip it stands.
// If it is standing at edge of clip(s), it doesn't do anything.
func (c *Cursor) Cut() {
	if c.i == len(c.clips) {
		if c.o != 0 {
			panic("c.o should 0 when c.i == len(c.clips)")
		}
		return
	}
	if c.o == 0 {
		// edge
		return
	}
	clipA, clipB := c.clips[c.i].Cut(c.o)
	c.clips = append(c.clips[:c.i], append([]Clip{clipA, clipB}, c.clips[c.i+1:]...)...)
	c.i++
	c.o = 0
}

func (c *Cursor) Write(r rune) {
	if c.appending {
		if c.o != 0 {
			panic("c.o should 0 when appending")
		}
		c.clips[c.i-1] = c.clips[c.i-1].Append(r)
		return
	}
	c.appending = true
	c.Cut()
	clipInsert := DataClip(runeToBytes(r))
	c.clips = append(c.clips[:c.i], append([]Clip{clipInsert}, c.clips[c.i:]...)...)
	c.i++
	c.o = 0
}

func (c *Cursor) Delete() {
	c.appending = false
	if c.i == len(c.clips) {
		if c.o != 0 {
			panic("c.o should 0 when c.i == len(c.clips)")
		}
		return
	}
	c.Cut()
	p := nextOffset(c.clips[c.i].data, 0)
	if p == -1 {
		c.clips = append(c.clips[:c.i], c.clips[c.i+1:]...)
		return
	}
	_, c.clips[c.i] = c.clips[c.i].Cut(p)
}

func (c *Cursor) Backspace() {
	c.appending = false
	if c.i == 0 && c.o == 0 {
		return
	}
	c.Cut()
	p := prevOffset(c.clips[c.i-1].data, len(c.clips[c.i-1].data))
	if p == 0 {
		c.clips = append(c.clips[:c.i-1], c.clips[c.i:]...)
		c.i--
		return
	}
	c.clips[c.i-1], _ = c.clips[c.i-1].Cut(p)
}
<|endoftext|>


<commit_before>package mock

import (
	"errors"
	"net"
	"time"
)

type MockConn struct {
	t            int
	err          error
	block        bool
	readIndex    int
	first        bool
	readTimeout  time.Time
	writeTimeout time.Time
	readings     [][]byte
	Closed       bool
	Written      [][]byte
}

var TimeoutErr = &TimeoutError{errors.New("i/o timeout")}

// Creates a new mock object which satisfies the net.Conn interface
func Conn() *MockConn {
	return &MockConn{
		first:    true,
		Written:  make([][]byte, 0, 4),
		readings: make([][]byte, 0, 4),
	}
}

func (c *MockConn) Read(out []byte) (int, error) {
	if c.first {
		c.first = false
		if c.err != nil {
			return 0, c.err
		}
	}

	if c.readIndex == len(c.readings) {
		if c.block == false {
			return 0, nil
		}
		time.Sleep(time.Minute)
	}
	bytes := c.readings[c.readIndex]
	end := len(bytes)
	if end > len(out) {
		end = len(out)
	}
	copy(out, bytes[:end])
	length := len(bytes[:end])
	c.readings[c.readIndex] = bytes[end:]
	if len(bytes[end:]) == 0 {
		c.readIndex++
	} else {
		c.readings[c.readIndex] = bytes[end:]
	}
	return length, nil
}

func (c *MockConn) Write(b []byte) (int, error) {
	if c.err != nil {
		return 0, c.err
	}
	if c.writeTimeout.IsZero() == false {
		return 0, TimeoutErr
	}
	c.Written = append(c.Written, b)
	return len(b), nil
}

func (c *MockConn) Close() error {
	c.Closed = true
	return nil
}

func (c *MockConn) LocalAddr() net.Addr {
	return nil
}

func (c *MockConn) RemoteAddr() net.Addr {
	return nil
}

func (c *MockConn) SetDeadline(t time.Time) error {
	c.SetReadDeadline(t)
	c.SetWriteDeadline(t)
	return nil
}

func (c *MockConn) SetReadDeadline(t time.Time) error {
	c.readTimeout = t
	return nil
}

// noop
func (c *MockConn) SetWriteDeadline(t time.Time) error {
	c.writeTimeout = t
	return nil
}

func (c *MockConn) Error(err error) *MockConn {
	c.err = err
	return c
}

// Define what each call to Read() will read
func (c *MockConn) Reading(data ...[]byte) *MockConn {
	for _, d := range data {
		c.readings = append(c.readings, d)
	}
	return c
}

func (c *MockConn) Block() *MockConn {
	c.block = true
	return c
}

// Returns whatever data hasn't been read yet
func (c *MockConn) Drain() []byte {
	left := 0
	for i := c.readIndex; i < len(c.readings); i++ {
		left += len(c.readings[i])
	}
	data := make([]byte, 0, left)
	for i := c.readIndex; i < len(c.readings); i++ {
		data = append(data, c.readings[i]...)
	}
	return data
}

type TimeoutError struct {
	error
}

func (t *TimeoutError) Timeout() bool {
	return true
}

func (t *TimeoutError) Temporary() bool {
	return true
}
<commit_msg>thought I undid this yesterday. Timeout attempt failed miserably<commit_after>package mock

import (
	"net"
	"time"
)

type MockConn struct {
	t            int
	err          error
	block        bool
	readIndex    int
	first        bool
	readings     [][]byte
	Closed       bool
	Written      [][]byte
}

// Creates a new mock object which satisfies the net.Conn interface
func Conn() *MockConn {
	return &MockConn{
		first:    true,
		Written:  make([][]byte, 0, 4),
		readings: make([][]byte, 0, 4),
	}
}

func (c *MockConn) Read(out []byte) (int, error) {
	if c.first {
		c.first = false
		if c.err != nil {
			return 0, c.err
		}
	}

	if c.readIndex == len(c.readings) {
		if c.block == false {
			return 0, nil
		}
		time.Sleep(time.Minute)
	}
	bytes := c.readings[c.readIndex]
	end := len(bytes)
	if end > len(out) {
		end = len(out)
	}
	copy(out, bytes[:end])
	length := len(bytes[:end])
	c.readings[c.readIndex] = bytes[end:]
	if len(bytes[end:]) == 0 {
		c.readIndex++
	} else {
		c.readings[c.readIndex] = bytes[end:]
	}
	return length, nil
}

func (c *MockConn) Write(b []byte) (int, error) {
	if c.err != nil {
		return 0, c.err
	}
	c.Written = append(c.Written, b)
	return len(b), nil
}

func (c *MockConn) Close() error {
	c.Closed = true
	return nil
}

func (c *MockConn) LocalAddr() net.Addr {
	return nil
}

func (c *MockConn) RemoteAddr() net.Addr {
	return nil
}

func (c *MockConn) SetDeadline(t time.Time) error {
	c.SetReadDeadline(t)
	c.SetWriteDeadline(t)
	return nil
}

func (c *MockConn) SetReadDeadline(t time.Time) error {
	return nil
}

// noop
func (c *MockConn) SetWriteDeadline(t time.Time) error {
	return nil
}

func (c *MockConn) Error(err error) *MockConn {
	c.err = err
	return c
}

// Define what each call to Read() will read
func (c *MockConn) Reading(data ...[]byte) *MockConn {
	for _, d := range data {
		c.readings = append(c.readings, d)
	}
	return c
}

func (c *MockConn) Block() *MockConn {
	c.block = true
	return c
}

// Returns whatever data hasn't been read yet
func (c *MockConn) Drain() []byte {
	left := 0
	for i := c.readIndex; i < len(c.readings); i++ {
		left += len(c.readings[i])
	}
	data := make([]byte, 0, left)
	for i := c.readIndex; i < len(c.readings); i++ {
		data = append(data, c.readings[i]...)
	}
	return data
}

type TimeoutError struct {
	error
}

func (t *TimeoutError) Timeout() bool {
	return true
}

func (t *TimeoutError) Temporary() bool {
	return true
}
<|endoftext|>


<commit_before>// Package bq includes all code related to BigQuery.
//
// NB: NOTES ON MEMORY USE AND HTTP SIZE
// The bigquery library uses JSON encoding of data, which appears to be the only
// option at this time.  Furthermore, it uses intermediate data representations,
// eventually creating a map[string]Value (unless you pass that in to begin with).
// In general, when we start pumping large volumes of data, both the map and the
// JSON will cause some memory pressure, and likely pretty severe limits on the size
// of the insert we can send, likely on the order of a couple MB of actual row footprint
// in BQ.
//
// Passing in slice of structs makes memory pressure a bit worse, but probably isn't
// worth worrying about.

package bq

import (
	"context"
	"encoding/json"
	"log"
	"math/rand"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"cloud.google.com/go/bigquery"

	"github.com/m-lab/etl/etl"
	"github.com/m-lab/etl/metrics"
)

// insertsBeforeRowJSONCount controls how often we perform a wasted JSON marshal
// on a row before inserting it into BQ. The smaller the number the more wasted
// CPU time. -- 1% of inserts.
const insertsBeforeRowJSONCount = 100

// tableInsertCounts tracks the number of inserts to each table. Keys are table
// names and values are *uint32.
var tableInsertCounts = sync.Map{}

func init() {
}

// NewInserter creates a new BQInserter with appropriate characteristics.
func NewInserter(dt etl.DataType, partition time.Time) (etl.Inserter, error) {
	suffix := ""
	if etl.IsBatchService() || time.Since(partition) > 30*24*time.Hour {
		// If batch, or too far in the past, we use a templated table, and must merge it later.
		suffix = "_" + partition.Format("20060102")
	} else {
		// Otherwise, we can stream directly to correct partition.
		suffix = "$" + partition.Format("20060102")
	}

	bqProject := dt.BigqueryProject()
	dataset := dt.Dataset()
	table := dt.Table()

	return NewBQInserter(
		etl.InserterParams{Project: bqProject, Dataset: dataset, Table: table, Suffix: suffix,
			PutTimeout: 60 * time.Second, BufferSize: dt.BQBufferSize(), RetryDelay: 30 * time.Second},
		nil)
}

// NewBQInserter initializes a new BQInserter
// Pass in nil uploader for normal use, custom uploader for custom behavior
// TODO - improve the naming between here and NewInserter.
func NewBQInserter(params etl.InserterParams, uploader etl.Uploader) (etl.Inserter, error) {
	if uploader == nil {
		client, err := GetClient(params.Project)
		if err != nil {
			return nil, err
		}
		table := params.Table
		if params.Suffix[0] == '$' {
			// Suffix starting with $ is just a partition spec.
			table += params.Suffix
		}
		u := client.Dataset(params.Dataset).Table(table).Uploader()
		if params.Suffix[0] == '_' {
			// Suffix starting with _ is a template suffix.
			u.TableTemplateSuffix = params.Suffix
		}
		// This avoids problems when a single row of the insert has invalid
		// data.  We then have to carefully parse the returned error object.
		u.SkipInvalidRows = true
		uploader = u
	}
	in := BQInserter{params: params, uploader: uploader, putTimeout: params.PutTimeout}
	in.rows = make([]interface{}, 0, in.params.BufferSize)
	return &in, nil
}

//===============================================================================

// GetClient returns an appropriate bigquery client.
// This incurs network delay, so it should not be inside fast loops.
func GetClient(project string) (*bigquery.Client, error) {
	// We do this here, instead of in init(), because we only want to do it
	// when we actually want to access the bigquery backend.

	// Network request
	// ctx is used only for the request to create the client.  It is not used by
	// the client.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	return bigquery.NewClient(ctx, project)
}

//===============================================================================

// MapSaver is a generic implementation of bq.ValueSaver, based on maps.  This avoids extra
// conversion steps in the bigquery library (except for the JSON conversion).
// IMPLEMENTS: bigquery.ValueSaver
type MapSaver struct {
	Values map[string]bigquery.Value
}

func (s *MapSaver) Save() (row map[string]bigquery.Value, insertID string, err error) {
	return s.Values, "", nil
}

//----------------------------------------------------------------------------

type BQInserter struct {
	etl.Inserter
	params     etl.InserterParams
	uploader   etl.Uploader  // May be a BQ Uploader, or a test Uploader
	putTimeout time.Duration // Timeout used for BQ put operations.
	rows       []interface{}
	inserted   int // Number of rows successfully inserted.
	badRows    int // Number of row failures, including rows in full failures.
	failures   int // Number of complete insert failures.
}

// Caller should check error, and take appropriate action before calling again.
func (in *BQInserter) InsertRow(data interface{}) error {
	return in.InsertRows([]interface{}{data})
}

// maybeCountRowSize periodically converts a row to JSON to measure its size.
func (in *BQInserter) maybeCountRowSize(data []interface{}) {
	if len(data) == 0 {
		return
	}
	counter, ok := tableInsertCounts.Load(in.TableBase())
	if !ok {
		tableInsertCounts.Store(in.TableBase(), new(uint32))
		counter, ok = tableInsertCounts.Load(in.TableBase())
		if !ok {
			log.Println("No counter for", in.TableBase())
			return
		}
	}
	count := atomic.AddUint32(counter.(*uint32), 1)
	// Do this just once in a while, so it doesn't take much resource.
	if count%insertsBeforeRowJSONCount != 0 {
		return
	}
	// Note: this estimate works very well for map[]bigquery.Value types. And, we
	// believe it is an okay estimate for struct types.
	jsonRow, _ := json.Marshal(data[0])
	metrics.RowSizeHistogram.WithLabelValues(
		in.TableBase()).Observe(float64(len(jsonRow)))
}

// Caller should check error, and take appropriate action before calling again.
// TODO - should this return a specific error to indicate that a flush is needed
// instead of flushing internally?  The "handle errors in the middle" would
// be easier, though other complications would ensue.
func (in *BQInserter) InsertRows(data []interface{}) error {
	metrics.WorkerState.WithLabelValues(in.TableBase(), "insert").Inc()
	defer metrics.WorkerState.WithLabelValues(in.TableBase(), "insert").Dec()

	in.maybeCountRowSize(data)

	for len(data)+len(in.rows) >= in.params.BufferSize {
		// space >= len(data)
		space := cap(in.rows) - len(in.rows)
		var add []interface{}
		add, data = data[:space], data[space:] // does this break?
		in.rows = append(in.rows, add...)
		err := in.Flush()
		if err != nil {
			// TODO - handle errors in middle better?
			return err
		}
	}
	in.rows = append(in.rows, data...)
	return nil
}

func (in *BQInserter) HandleInsertErrors(err error) error {
	switch typedErr := err.(type) {
	case bigquery.PutMultiError:
		if len(typedErr) == len(in.rows) {
			log.Printf("%v\n", err)
			metrics.BackendFailureCount.WithLabelValues(
				in.TableBase(), "failed insert").Inc()
			in.failures += 1
		}
		// If ALL rows failed, and number of rows is large, just report single failure.
		if len(typedErr) > 10 && len(typedErr) == len(in.rows) {
			log.Printf("Insert error: %v\n", err)
			metrics.ErrorCount.WithLabelValues(
				in.TableBase(), "unknown", "insert row error").
				Add(float64(len(typedErr)))
		} else {
			// Handle each error individually.
			// TODO Should we try to handle large numbers of row errors?
			for _, rowError := range typedErr {
				// These are rowInsertionErrors
				log.Printf("%v\n", rowError)
				// rowError.Errors is a MultiError
				for _, oneErr := range rowError.Errors {
					log.Printf("Insert error: %v\n", oneErr)
					metrics.ErrorCount.WithLabelValues(
						in.TableBase(), "unknown", "insert row error").Inc()
				}
			}
		}
		in.inserted += len(in.rows) - len(typedErr)
		in.badRows += len(typedErr)
		err = nil
	default:
		log.Printf("Unhandled insert error %v\n", typedErr)
		metrics.BackendFailureCount.WithLabelValues(
			in.TableBase(), "failed insert").Inc()
		metrics.ErrorCount.WithLabelValues(
			in.TableBase(), "unknown", "UNHANDLED insert error").Inc()
		// TODO - Conservative, but possibly not correct.
		// This at least preserves the count invariance.
		in.badRows += len(in.rows)
		err = nil
	}
	return err
}

// TODO(dev) Should have a recovery mechanism for failed inserts.
func (in *BQInserter) Flush() error {
	metrics.WorkerState.WithLabelValues(in.TableBase(), "flush").Inc()
	defer metrics.WorkerState.WithLabelValues(in.TableBase(), "flush").Dec()

	if len(in.rows) == 0 {
		return nil
	}

	// If we exceed the quota, this basically backs off and tries again.  When
	// operating near quota, this will fire enough times to slow down each task
	// enough to stay within quota.  It may result in AppEngine reducing the
	// number of workers, but that is fine - it will also result in staying
	// under the quota.
	// Analysis:
	//  We can handle a minimum of 10 inserts per second, because
	//   the default quota is 100MB/sec, and the limit on requests
	//   is 10MB per request.  Since generally inserts are smaller,
	//   the typical number is more like 20 inserts/sec.
	//  The net effect we need to see is that, if the pipeline capacity
	//   exceeds the quota by 10%, then the pipeline needs to slow down
	//   by roughly 10% to fit within the quota.  The incoming request
	//   rate is dictated by the task queue, and ultimately the handler
	//   must reject 10% of the incoming requests.  This only happens
	//   when 10% of the instances have hit MAX_WORKERS.
	//  If the capacity of the pipeline is, e.g., 2X the task queue rate,
	//   then each task will need to be slowed down to the point that it
	//   takes roughly 2.2X longer than it would without any Quota exceeded
	//   errors.  For NDT, the 100MB tasks require about 35 concurrent tasks
	//   to process 60 tasks/min, indicating that they require about 35
	//   seconds per task.  There are about 70 tests/task, so this is about
	//   7 buffer flushes per second (of 10 tests each), or on average, about
	//   one buffer flush every 5 seconds for each task.
	//  The batch job might have 50 instances, and process 900 tasks
	//   concurrently.  If this had to be scaled back to 50%, the tasks
	//   would have to spend 50% of their time sleeping between Put requests.
	//   Since each task typically takes about 35 seconds, each task would
	//   on average experience just over one 'Quota exceeded' error in order
	//   to slow the pipeline down by 50%.
	//  Note that a secondary effect is to reduce the CPU utilization, which
	//   will likely trigger a reduction in the number of instances running.
	//   Under these conditions, AppEngine would reduce the number of instances
	//   until the target utilization is reaches, reducing the number of
	//   concurrent tasks, and thus the frequency at which the tasks would
	//   experience 'Quota error' events.

	var err error
	for i := 0; i < 10; i++ {
		// This is heavyweight, and may run forever without a context deadline.
		ctx, _ := context.WithTimeout(context.Background(), in.putTimeout)
		err = in.uploader.Put(ctx, in.rows)
		if err == nil || !strings.Contains(err.Error(), "Quota exceeded:") {
			break
		}
		metrics.WarningCount.WithLabelValues(in.TableBase(), "", "Quota Exceeded").Inc()
		// Use some randomness to reduce risk of synchronization across tasks.
		t := in.params.RetryDelay.Seconds() * (0.5 + rand.Float64()) // between 0.5 and 1.5 * RetryDelay
		time.Sleep(time.Duration(1000000*t) * time.Microsecond)
	}

	// If there is still an error, then handle it.
	if err == nil {
		in.inserted += len(in.rows)
	} else {
		// This adjusts the inserted count, failure count, and updates in.rows.
		err = in.HandleInsertErrors(err)
	}
	// Allocate new slice of rows.  Any failed rows are lost.
	in.rows = make([]interface{}, 0, in.params.BufferSize)
	return err
}

func (in *BQInserter) FullTableName() string {
	return in.TableBase() + in.TableSuffix()
}
func (in *BQInserter) TableBase() string {
	return in.params.Table
}

// The $ or _ suffix.
func (in *BQInserter) TableSuffix() string {
	return in.params.Suffix
}
func (in *BQInserter) Project() string {
	return in.params.Project
}
func (in *BQInserter) Dataset() string {
	return in.params.Dataset
}
func (in *BQInserter) RowsInBuffer() int {
	return len(in.rows)
}
func (in *BQInserter) Accepted() int {
	return in.inserted + in.badRows + len(in.rows)
}
func (in *BQInserter) Committed() int {
	return in.inserted
}
func (in *BQInserter) Failed() int {
	return in.badRows
}

//----------------------------------------------------------------------------

// Inserter wrapper that handles flush metrics.
// TODO - add prometheus counters for attempts, number of rows.
type DurationWrapper struct {
	etl.Inserter
}

func (dw DurationWrapper) Flush() error {
	t := time.Now()
	status := "succeed"
	err := dw.Inserter.Flush()
	if err != nil {
		status = "fail"
	}
	metrics.InsertionHistogram.WithLabelValues(
		dw.TableBase(), status).Observe(time.Since(t).Seconds())
	return err
}
<commit_msg>fix PutTimeout, clarify comments<commit_after>// Package bq includes all code related to BigQuery.
//
// NB: NOTES ON MEMORY USE AND HTTP SIZE
// The bigquery library uses JSON encoding of data, which appears to be the only
// option at this time.  Furthermore, it uses intermediate data representations,
// eventually creating a map[string]Value (unless you pass that in to begin with).
// In general, when we start pumping large volumes of data, both the map and the
// JSON will cause some memory pressure, and likely pretty severe limits on the size
// of the insert we can send, likely on the order of a couple MB of actual row footprint
// in BQ.
//
// Passing in slice of structs makes memory pressure a bit worse, but probably isn't
// worth worrying about.

package bq

import (
	"context"
	"encoding/json"
	"log"
	"math/rand"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"cloud.google.com/go/bigquery"

	"github.com/m-lab/etl/etl"
	"github.com/m-lab/etl/metrics"
)

// insertsBeforeRowJSONCount controls how often we perform a wasted JSON marshal
// on a row before inserting it into BQ. The smaller the number the more wasted
// CPU time. -- 1% of inserts.
const insertsBeforeRowJSONCount = 100

// tableInsertCounts tracks the number of inserts to each table. Keys are table
// names and values are *uint32.
var tableInsertCounts = sync.Map{}

func init() {
}

// NewInserter creates a new BQInserter with appropriate characteristics.
func NewInserter(dt etl.DataType, partition time.Time) (etl.Inserter, error) {
	suffix := ""
	if etl.IsBatchService() || time.Since(partition) > 30*24*time.Hour {
		// If batch, or too far in the past, we use a templated table, and must merge it later.
		suffix = "_" + partition.Format("20060102")
	} else {
		// Otherwise, we can stream directly to correct partition.
		suffix = "$" + partition.Format("20060102")
	}

	bqProject := dt.BigqueryProject()
	dataset := dt.Dataset()
	table := dt.Table()

	return NewBQInserter(
		etl.InserterParams{Project: bqProject, Dataset: dataset, Table: table, Suffix: suffix,
			PutTimeout: 60 * time.Second, BufferSize: dt.BQBufferSize(), RetryDelay: 30 * time.Second},
		nil)
}

// NewBQInserter initializes a new BQInserter
// Pass in nil uploader for normal use, custom uploader for custom behavior
// TODO - improve the naming between here and NewInserter.
func NewBQInserter(params etl.InserterParams, uploader etl.Uploader) (etl.Inserter, error) {
	if uploader == nil {
		client, err := GetClient(params.Project)
		if err != nil {
			return nil, err
		}
		table := params.Table
		if params.Suffix[0] == '$' {
			// Suffix starting with $ is just a partition spec.
			table += params.Suffix
		}
		u := client.Dataset(params.Dataset).Table(table).Uploader()
		if params.Suffix[0] == '_' {
			// Suffix starting with _ is a template suffix.
			u.TableTemplateSuffix = params.Suffix
		}
		// This avoids problems when a single row of the insert has invalid
		// data.  We then have to carefully parse the returned error object.
		u.SkipInvalidRows = true
		uploader = u
	}
	in := BQInserter{params: params, uploader: uploader, putTimeout: params.PutTimeout}
	in.rows = make([]interface{}, 0, in.params.BufferSize)
	return &in, nil
}

//===============================================================================

// GetClient returns an appropriate bigquery client.
// This incurs network delay, so it should not be inside fast loops.
func GetClient(project string) (*bigquery.Client, error) {
	// We do this here, instead of in init(), because we only want to do it
	// when we actually want to access the bigquery backend.

	// Network request
	// ctx is used only for the request to create the client.  It is not used by
	// the client.
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	return bigquery.NewClient(ctx, project)
}

//===============================================================================

// MapSaver is a generic implementation of bq.ValueSaver, based on maps.  This avoids extra
// conversion steps in the bigquery library (except for the JSON conversion).
// IMPLEMENTS: bigquery.ValueSaver
type MapSaver struct {
	Values map[string]bigquery.Value
}

func (s *MapSaver) Save() (row map[string]bigquery.Value, insertID string, err error) {
	return s.Values, "", nil
}

//----------------------------------------------------------------------------

type BQInserter struct {
	etl.Inserter
	params     etl.InserterParams
	uploader   etl.Uploader  // May be a BQ Uploader, or a test Uploader
	putTimeout time.Duration // Timeout used for BQ put operations.
	rows       []interface{}
	inserted   int // Number of rows successfully inserted.
	badRows    int // Number of row failures, including rows in full failures.
	failures   int // Number of complete insert failures.
}

// Caller should check error, and take appropriate action before calling again.
func (in *BQInserter) InsertRow(data interface{}) error {
	return in.InsertRows([]interface{}{data})
}

// maybeCountRowSize periodically converts a row to JSON to measure its size.
func (in *BQInserter) maybeCountRowSize(data []interface{}) {
	if len(data) == 0 {
		return
	}
	counter, ok := tableInsertCounts.Load(in.TableBase())
	if !ok {
		tableInsertCounts.Store(in.TableBase(), new(uint32))
		counter, ok = tableInsertCounts.Load(in.TableBase())
		if !ok {
			log.Println("No counter for", in.TableBase())
			return
		}
	}
	count := atomic.AddUint32(counter.(*uint32), 1)
	// Do this just once in a while, so it doesn't take much resource.
	if count%insertsBeforeRowJSONCount != 0 {
		return
	}
	// Note: this estimate works very well for map[]bigquery.Value types. And, we
	// believe it is an okay estimate for struct types.
	jsonRow, _ := json.Marshal(data[0])
	metrics.RowSizeHistogram.WithLabelValues(
		in.TableBase()).Observe(float64(len(jsonRow)))
}

// Caller should check error, and take appropriate action before calling again.
// TODO - should this return a specific error to indicate that a flush is needed
// instead of flushing internally?  The "handle errors in the middle" would
// be easier, though other complications would ensue.
func (in *BQInserter) InsertRows(data []interface{}) error {
	metrics.WorkerState.WithLabelValues(in.TableBase(), "insert").Inc()
	defer metrics.WorkerState.WithLabelValues(in.TableBase(), "insert").Dec()

	in.maybeCountRowSize(data)

	for len(data)+len(in.rows) >= in.params.BufferSize {
		// space >= len(data)
		space := cap(in.rows) - len(in.rows)
		var add []interface{}
		add, data = data[:space], data[space:] // does this break?
		in.rows = append(in.rows, add...)
		err := in.Flush()
		if err != nil {
			// TODO - handle errors in middle better?
			return err
		}
	}
	in.rows = append(in.rows, data...)
	return nil
}

func (in *BQInserter) HandleInsertErrors(err error) error {
	switch typedErr := err.(type) {
	case bigquery.PutMultiError:
		if len(typedErr) == len(in.rows) {
			log.Printf("%v\n", err)
			metrics.BackendFailureCount.WithLabelValues(
				in.TableBase(), "failed insert").Inc()
			in.failures += 1
		}
		// If ALL rows failed, and number of rows is large, just report single failure.
		if len(typedErr) > 10 && len(typedErr) == len(in.rows) {
			log.Printf("Insert error: %v\n", err)
			metrics.ErrorCount.WithLabelValues(
				in.TableBase(), "unknown", "insert row error").
				Add(float64(len(typedErr)))
		} else {
			// Handle each error individually.
			// TODO Should we try to handle large numbers of row errors?
			for _, rowError := range typedErr {
				// These are rowInsertionErrors
				log.Printf("%v\n", rowError)
				// rowError.Errors is a MultiError
				for _, oneErr := range rowError.Errors {
					log.Printf("Insert error: %v\n", oneErr)
					metrics.ErrorCount.WithLabelValues(
						in.TableBase(), "unknown", "insert row error").Inc()
				}
			}
		}
		in.inserted += len(in.rows) - len(typedErr)
		in.badRows += len(typedErr)
		err = nil
	default:
		log.Printf("Unhandled insert error %v\n", typedErr)
		metrics.BackendFailureCount.WithLabelValues(
			in.TableBase(), "failed insert").Inc()
		metrics.ErrorCount.WithLabelValues(
			in.TableBase(), "unknown", "UNHANDLED insert error").Inc()
		// TODO - Conservative, but possibly not correct.
		// This at least preserves the count invariance.
		in.badRows += len(in.rows)
		err = nil
	}
	return err
}

// TODO(dev) Should have a recovery mechanism for failed inserts.
func (in *BQInserter) Flush() error {
	metrics.WorkerState.WithLabelValues(in.TableBase(), "flush").Inc()
	defer metrics.WorkerState.WithLabelValues(in.TableBase(), "flush").Dec()

	if len(in.rows) == 0 {
		return nil
	}

	// If we exceed the quota, this basically backs off and tries again.  When
	// operating near quota, this will fire enough times to slow down each task
	// enough to stay within quota.  It may result in AppEngine reducing the
	// number of workers, but that is fine - it will also result in staying
	// under the quota.
	// Analysis:
	//  We can handle a minimum of 10 inserts per second, because
	//   the default quota is 100MB/sec, and the limit on requests
	//   is 10MB per request.  Since generally inserts are smaller,
	//   the typical number is more like 20 inserts/sec.
	//  The net effect we need to see is that, if the pipeline capacity
	//   exceeds the quota by 10%, then the pipeline needs to slow down
	//   by roughly 10% to fit within the quota.  The incoming request
	//   rate is dictated by the task queue, and ultimately the handler
	//   must reject 10% of the incoming requests.  This only happens
	//   when 10% of the instances have hit MAX_WORKERS.
	//  If the capacity of the pipeline is, e.g., 2X the task queue rate,
	//   then each task will need to be slowed down to the point that it
	//   takes roughly 2.2X longer than it would without any Quota exceeded
	//   errors.  For NDT, the 100MB tasks require about 35 concurrent tasks
	//   to process 60 tasks/min, indicating that they require about 35
	//   seconds per task.  There are about 70 tests/task, so this is about
	//   7 buffer flushes per second (of 10 tests each), or on average, about
	//   one buffer flush every 5 seconds for each task.
	//  The batch job might have 50 instances, and process 900 tasks
	//   concurrently.  If this had to be scaled back to 50%, the tasks
	//   would have to spend 50% of their time sleeping between Put requests.
	//   Since each task typically takes about 35 seconds, each task would
	//   on average experience just over one 'Quota exceeded' error in order
	//   to slow the pipeline down by 50%.
	//  Note that a secondary effect is to reduce the CPU utilization, which
	//   will likely trigger a reduction in the number of instances running.
	//   Under these conditions, AppEngine would reduce the number of instances
	//   until the target utilization is reaches, reducing the number of
	//   concurrent tasks, and thus the frequency at which the tasks would
	//   experience 'Quota error' events.

	var err error
	for i := 0; i < 10; i++ {
		// This is heavyweight, and may run forever without a context deadline.
		ctx, _ := context.WithTimeout(context.Background(), in.putTimeout)
		err = in.uploader.Put(ctx, in.rows)
		if err == nil || !strings.Contains(err.Error(), "Quota exceeded:") {
			break
		}
		metrics.WarningCount.WithLabelValues(in.TableBase(), "", "Quota Exceeded").Inc()
		// Use some randomness to reduce risk of synchronization across tasks.
		t := in.params.RetryDelay.Seconds() * (0.5 + rand.Float64()) // between 0.5 and 1.5 * RetryDelay
		time.Sleep(time.Duration(1000000*t) * time.Microsecond)
	}

	// If there is still an error, then handle it.
	if err == nil {
		in.inserted += len(in.rows)
	} else {
		// This adjusts the inserted count, failure count, and updates in.rows.
		err = in.HandleInsertErrors(err)
	}
	// Allocate new slice of rows.  Any failed rows are lost.
	in.rows = make([]interface{}, 0, in.params.BufferSize)
	return err
}

func (in *BQInserter) FullTableName() string {
	return in.TableBase() + in.TableSuffix()
}
func (in *BQInserter) TableBase() string {
	return in.params.Table
}

// The $ or _ suffix.
func (in *BQInserter) TableSuffix() string {
	return in.params.Suffix
}
func (in *BQInserter) Project() string {
	return in.params.Project
}
func (in *BQInserter) Dataset() string {
	return in.params.Dataset
}
func (in *BQInserter) RowsInBuffer() int {
	return len(in.rows)
}
func (in *BQInserter) Accepted() int {
	return in.inserted + in.badRows + len(in.rows)
}
func (in *BQInserter) Committed() int {
	return in.inserted
}
func (in *BQInserter) Failed() int {
	return in.badRows
}

//----------------------------------------------------------------------------

// Inserter wrapper that handles flush metrics.
// TODO - add prometheus counters for attempts, number of rows.
type DurationWrapper struct {
	etl.Inserter
}

func (dw DurationWrapper) Flush() error {
	t := time.Now()
	status := "succeed"
	err := dw.Inserter.Flush()
	if err != nil {
		status = "fail"
	}
	metrics.InsertionHistogram.WithLabelValues(
		dw.TableBase(), status).Observe(time.Since(t).Seconds())
	return err
}
<|endoftext|>


<commit_before>package commons

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
)

const (
	CF_PUSH_TIMEOUT                       = 4 * time.Minute
	LONG_CURL_TIMEOUT                     = 4 * time.Minute
	DOCKER_IMAGE_DOWNLOAD_DEFAULT_TIMEOUT = 10 * time.Minute
	NOT_FOUND                             = "404 Not Found"
	OK_RESPONSE                           = "0"
)

func GuidForAppName(appName string) string {
	cfApp := cf.Cf("app", appName, "--guid")
	Expect(cfApp.Wait()).To(Exit(0))

	appGuid := strings.TrimSpace(string(cfApp.Out.Contents()))
	Expect(appGuid).NotTo(Equal(""))
	return appGuid
}

func GuidForSpaceName(spaceName string) string {
	cfSpace := cf.Cf("space", spaceName, "--guid")
	Expect(cfSpace.Wait()).To(Exit(0))

	spaceGuid := strings.TrimSpace(string(cfSpace.Out.Contents()))
	Expect(spaceGuid).NotTo(Equal(""))
	return spaceGuid
}

func AssertImageAvailable(registryAddress string, imageName string) {
	client := http.Client{}
	resp, err := client.Get(fmt.Sprintf("http://%s/v1/search?q=%s", registryAddress, imageName))
	Expect(err).NotTo(HaveOccurred())
	Expect(resp.StatusCode).To(Equal(http.StatusOK))
	bytes, err := ioutil.ReadAll(resp.Body)
	Expect(err).NotTo(HaveOccurred())
	Expect(string(bytes)).To(ContainSubstring("library/" + imageName))
}

func EnableDockerFeatureFlag(context helpers.SuiteContext) {
	cf.AsUser(context.AdminUserContext(), time.Minute, func() {
		Eventually(cf.Cf("enable-feature-flag", "diego_docker")).Should(Exit(0))
	})
}

func DisableDockerFeatureFlag(context helpers.SuiteContext) {
	cf.AsUser(context.AdminUserContext(), time.Minute, func() {
		Eventually(cf.Cf("disable-feature-flag", "diego_docker")).Should(Exit(0))
	})
}

func GetAppLogs(appName string) string {
	cfLogs := cf.Cf("logs", appName, "--recent")
	Expect(cfLogs.Wait()).To(Exit(0))
	return string(cfLogs.Out.Contents())
}

func AssertDockerEnabled() {
	session := cf.Cf("feature-flag", "diego_docker")
	Eventually(session.Wait()).Should(Exit(0))

	Expect(session.Out).To(Say("enabled"))
}
<commit_msg>Replace catalog with image tag listing<commit_after>package commons

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
)

const (
	CF_PUSH_TIMEOUT                       = 4 * time.Minute
	LONG_CURL_TIMEOUT                     = 4 * time.Minute
	DOCKER_IMAGE_DOWNLOAD_DEFAULT_TIMEOUT = 10 * time.Minute
	NOT_FOUND                             = "404 Not Found"
	OK_RESPONSE                           = "0"
)

func GuidForAppName(appName string) string {
	cfApp := cf.Cf("app", appName, "--guid")
	Expect(cfApp.Wait()).To(Exit(0))

	appGuid := strings.TrimSpace(string(cfApp.Out.Contents()))
	Expect(appGuid).NotTo(Equal(""))
	return appGuid
}

func GuidForSpaceName(spaceName string) string {
	cfSpace := cf.Cf("space", spaceName, "--guid")
	Expect(cfSpace.Wait()).To(Exit(0))

	spaceGuid := strings.TrimSpace(string(cfSpace.Out.Contents()))
	Expect(spaceGuid).NotTo(Equal(""))
	return spaceGuid
}

func AssertImageAvailable(registryAddress string, imageName string) {
	client := http.Client{}
	resp, err := client.Get(fmt.Sprintf("http://%s/v2/%s/tags/list", registryAddress, imageName))
	Expect(err).NotTo(HaveOccurred())
	Expect(resp.StatusCode).To(Equal(http.StatusOK))
	bytes, err := ioutil.ReadAll(resp.Body)
	Expect(err).NotTo(HaveOccurred())
	Expect(string(bytes)).To(ContainSubstring(fmt.Sprintf(`"name":"%s"`, imageName)))
}

func EnableDockerFeatureFlag(context helpers.SuiteContext) {
	cf.AsUser(context.AdminUserContext(), time.Minute, func() {
		Eventually(cf.Cf("enable-feature-flag", "diego_docker")).Should(Exit(0))
	})
}

func DisableDockerFeatureFlag(context helpers.SuiteContext) {
	cf.AsUser(context.AdminUserContext(), time.Minute, func() {
		Eventually(cf.Cf("disable-feature-flag", "diego_docker")).Should(Exit(0))
	})
}

func GetAppLogs(appName string) string {
	cfLogs := cf.Cf("logs", appName, "--recent")
	Expect(cfLogs.Wait()).To(Exit(0))
	return string(cfLogs.Out.Contents())
}

func AssertDockerEnabled() {
	session := cf.Cf("feature-flag", "diego_docker")
	Eventually(session.Wait()).Should(Exit(0))

	Expect(session.Out).To(Say("enabled"))
}
<|endoftext|>


<commit_before>package github

import (
	"strings"
	"sync"

	"github.com/estafette/estafette-ci-api/estafette"
	"github.com/rs/zerolog/log"
)

// EventWorker processes events pushed to channels
type EventWorker interface {
	ListenToEventChannels()
	Stop()
	CreateJobForGithubPush(PushEvent)
}

type eventWorkerImpl struct {
	waitGroup       *sync.WaitGroup
	quitChannel     chan bool
	eventsChannel   chan PushEvent
	apiClient       APIClient
	ciBuilderClient estafette.CiBuilderClient
}

// NewGithubEventWorker returns a new github.EventWorker to handle events channeled by github.EventHandler
func NewGithubEventWorker(waitGroup *sync.WaitGroup, apiClient APIClient, ciBuilderClient estafette.CiBuilderClient, eventsChannel chan PushEvent) EventWorker {
	return &eventWorkerImpl{
		waitGroup:       waitGroup,
		quitChannel:     make(chan bool),
		eventsChannel:   eventsChannel,
		apiClient:       apiClient,
		ciBuilderClient: ciBuilderClient,
	}
}

func (w *eventWorkerImpl) ListenToEventChannels() {
	go func() {
		// handle github events via channels
		log.Debug().Msg("Listening to Github events channels...")
		for {
			select {
			case pushEvent := <-w.eventsChannel:
				go func() {
					w.waitGroup.Add(1)
					w.CreateJobForGithubPush(pushEvent)
					w.waitGroup.Done()
				}()
			case <-w.quitChannel:
				log.Debug().Msg("Stopping Github event worker...")
				return
			}
		}
	}()
}

func (w *eventWorkerImpl) Stop() {
	go func() {
		w.quitChannel <- true
	}()
}

func (w *eventWorkerImpl) CreateJobForGithubPush(pushEvent PushEvent) {

	// check to see that it's a cloneable event
	if !strings.HasPrefix(pushEvent.Ref, "refs/heads/") {
		return
	}

	// get access token
	accessToken, err := w.apiClient.GetInstallationToken(pushEvent.Installation.ID)
	if err != nil {
		log.Error().Err(err).
			Msg("Retrieving access token failed")
		return
	}

	// get manifest file
	manifestExists, manifest, err := w.apiClient.GetEstafetteManifest(accessToken, pushEvent)
	if err != nil {
		log.Error().Err(err).
			Msg("Retrieving Estafettte manifest failed")
		return
	}

	if !manifestExists {
		log.Info().Interface("pushEvent", pushEvent).Msgf("No Estaffette manifest for repo %v and revision %v, not creating a job", pushEvent.Repository.FullName, pushEvent.After)
		return
	}
	log.Debug().Interface("pushEvent", pushEvent).Str("manifest", manifest).Msgf("Estaffette manifest for repo %v and revision %v exists creating a builder job...", pushEvent.Repository.FullName, pushEvent.After)

	// get authenticated url for the repository
	authenticatedRepositoryURL, err := w.apiClient.GetAuthenticatedRepositoryURL(accessToken, pushEvent.Repository.HTMLURL)
	if err != nil {
		log.Error().Err(err).
			Msg("Retrieving authenticated repository failed")
		return
	}

	// define ci builder params
	ciBuilderParams := estafette.CiBuilderParams{
		RepoFullName:         pushEvent.Repository.FullName,
		RepoURL:              authenticatedRepositoryURL,
		RepoBranch:           strings.Replace(pushEvent.Ref, "refs/heads/", "", 1),
		RepoRevision:         pushEvent.After,
		EnvironmentVariables: map[string]string{"ESTAFETTE_GITHUB_API_TOKEN": accessToken.Token},
	}

	// create ci builder job
	_, err = w.ciBuilderClient.CreateCiBuilderJob(ciBuilderParams)
	if err != nil {
		log.Error().Err(err).
			Str("fullname", ciBuilderParams.RepoFullName).
			Str("url", ciBuilderParams.RepoURL).
			Str("branch", ciBuilderParams.RepoBranch).
			Str("revision", ciBuilderParams.RepoRevision).
			Msgf("Creating estafette-ci-builder job for Github repository %v revision %v failed", ciBuilderParams.RepoFullName, ciBuilderParams.RepoRevision)

		return
	}

	log.Info().
		Str("fullname", ciBuilderParams.RepoFullName).
		Str("url", ciBuilderParams.RepoURL).
		Str("branch", ciBuilderParams.RepoBranch).
		Str("revision", ciBuilderParams.RepoRevision).
		Msgf("Created estafette-ci-builder job for Github repository %v revision %v", ciBuilderParams.RepoFullName, ciBuilderParams.RepoRevision)
}
<commit_msg>Fixed spelling mistakes<commit_after>package github

import (
	"strings"
	"sync"

	"github.com/estafette/estafette-ci-api/estafette"
	"github.com/rs/zerolog/log"
)

// EventWorker processes events pushed to channels
type EventWorker interface {
	ListenToEventChannels()
	Stop()
	CreateJobForGithubPush(PushEvent)
}

type eventWorkerImpl struct {
	waitGroup       *sync.WaitGroup
	quitChannel     chan bool
	eventsChannel   chan PushEvent
	apiClient       APIClient
	ciBuilderClient estafette.CiBuilderClient
}

// NewGithubEventWorker returns a new github.EventWorker to handle events channeled by github.EventHandler
func NewGithubEventWorker(waitGroup *sync.WaitGroup, apiClient APIClient, ciBuilderClient estafette.CiBuilderClient, eventsChannel chan PushEvent) EventWorker {
	return &eventWorkerImpl{
		waitGroup:       waitGroup,
		quitChannel:     make(chan bool),
		eventsChannel:   eventsChannel,
		apiClient:       apiClient,
		ciBuilderClient: ciBuilderClient,
	}
}

func (w *eventWorkerImpl) ListenToEventChannels() {
	go func() {
		// handle github events via channels
		log.Debug().Msg("Listening to Github events channels...")
		for {
			select {
			case pushEvent := <-w.eventsChannel:
				go func() {
					w.waitGroup.Add(1)
					w.CreateJobForGithubPush(pushEvent)
					w.waitGroup.Done()
				}()
			case <-w.quitChannel:
				log.Debug().Msg("Stopping Github event worker...")
				return
			}
		}
	}()
}

func (w *eventWorkerImpl) Stop() {
	go func() {
		w.quitChannel <- true
	}()
}

func (w *eventWorkerImpl) CreateJobForGithubPush(pushEvent PushEvent) {

	// check to see that it's a cloneable event
	if !strings.HasPrefix(pushEvent.Ref, "refs/heads/") {
		return
	}

	// get access token
	accessToken, err := w.apiClient.GetInstallationToken(pushEvent.Installation.ID)
	if err != nil {
		log.Error().Err(err).
			Msg("Retrieving access token failed")
		return
	}

	// get manifest file
	manifestExists, manifest, err := w.apiClient.GetEstafetteManifest(accessToken, pushEvent)
	if err != nil {
		log.Error().Err(err).
			Msg("Retrieving Estafettte manifest failed")
		return
	}

	if !manifestExists {
		log.Info().Interface("pushEvent", pushEvent).Msgf("No Estafette manifest for repo %v and revision %v, not creating a job", pushEvent.Repository.FullName, pushEvent.After)
		return
	}
	log.Debug().Interface("pushEvent", pushEvent).Str("manifest", manifest).Msgf("Estafette manifest for repo %v and revision %v exists creating a builder job...", pushEvent.Repository.FullName, pushEvent.After)

	// get authenticated url for the repository
	authenticatedRepositoryURL, err := w.apiClient.GetAuthenticatedRepositoryURL(accessToken, pushEvent.Repository.HTMLURL)
	if err != nil {
		log.Error().Err(err).
			Msg("Retrieving authenticated repository failed")
		return
	}

	// define ci builder params
	ciBuilderParams := estafette.CiBuilderParams{
		RepoFullName:         pushEvent.Repository.FullName,
		RepoURL:              authenticatedRepositoryURL,
		RepoBranch:           strings.Replace(pushEvent.Ref, "refs/heads/", "", 1),
		RepoRevision:         pushEvent.After,
		EnvironmentVariables: map[string]string{"ESTAFETTE_GITHUB_API_TOKEN": accessToken.Token},
	}

	// create ci builder job
	_, err = w.ciBuilderClient.CreateCiBuilderJob(ciBuilderParams)
	if err != nil {
		log.Error().Err(err).
			Str("fullname", ciBuilderParams.RepoFullName).
			Str("url", ciBuilderParams.RepoURL).
			Str("branch", ciBuilderParams.RepoBranch).
			Str("revision", ciBuilderParams.RepoRevision).
			Msgf("Creating estafette-ci-builder job for Github repository %v revision %v failed", ciBuilderParams.RepoFullName, ciBuilderParams.RepoRevision)

		return
	}

	log.Info().
		Str("fullname", ciBuilderParams.RepoFullName).
		Str("url", ciBuilderParams.RepoURL).
		Str("branch", ciBuilderParams.RepoBranch).
		Str("revision", ciBuilderParams.RepoRevision).
		Msgf("Created estafette-ci-builder job for Github repository %v revision %v", ciBuilderParams.RepoFullName, ciBuilderParams.RepoRevision)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"sync"
	"time"

	"github.com/glycerine/rbuf"
)

type BroadcastString struct {
	writers map[chan string]bool
	mu      sync.Mutex
}

func NewBroadcastString() *BroadcastString {
	return &BroadcastString{
		writers: make(map[chan string]bool, 0),
	}
}

func (b *BroadcastString) WriteMessage(message string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for c := range b.writers {
		select {
		case c <- message:
		case <-time.After(500 * time.Millisecond):
			log.Println("channel closed, remove from queue")
			delete(b.writers, c)
		}
	}
}

func (b *BroadcastString) AddListener(c chan string) chan string {
	b.mu.Lock()
	defer b.mu.Unlock()
	if c == nil {
		c = make(chan string, 0)
	}
	b.writers[c] = true
	return c
}

func (b *BroadcastString) RemoveListener(c chan string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	delete(b.writers, c)
}

type BufferBroadcast struct {
	bs *BroadcastString

	maxSize int
	buf     *rbuf.FixedSizeRingBuf // *bytes.Buffer
	mu      sync.Mutex
}

func NewBufferBroadcast(size int) *BufferBroadcast {
	if size <= 0 {
		size = 4 * 1024 // 4K
	}
	return &BufferBroadcast{
		maxSize: size,
		bs:      NewBroadcastString(),
		buf:     rbuf.NewFixedSizeRingBuf(size), //  bytes.NewBuffer(nil), // buffer.NewRing(buffer.New(size)),
	}
}

func (b *BufferBroadcast) Write(data []byte) (n int, err error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	// if b.buf.Len() >= b.maxSize*2 {
	// 	b.buf = bytes.NewBuffer(b.buf.Bytes()[b.buf.Len()-b.maxSize : b.buf.Len()])
	// }
	b.bs.WriteMessage(string(data))
	return b.buf.Write(data)
}

func (b *BufferBroadcast) Reset() {
	b.buf.Reset()
}

func (b *BufferBroadcast) AddHookFunc(wf func(string) error) chan error {
	b.mu.Lock()
	defer b.mu.Unlock()
	c := b.bs.AddListener(nil)
	errC := make(chan error, 1)
	go func() {
		data := b.buf.Bytes()
		// data, _ := ioutil.ReadAll(b.buf)
		if err := wf(string(data)); err != nil {
			errC <- err
			return
		}
		for msg := range c {
			err := wf(msg)
			if err != nil {
				errC <- err
				break
			}
		}
	}()
	return errC
}
<commit_msg>fix pipe error<commit_after>package main

import (
	"sync"
	"time"

	"github.com/glycerine/rbuf"
	"github.com/qiniu/log"
)

type BroadcastString struct {
	msgC    chan string
	writers map[chan string]bool
	mu      sync.Mutex
}

func NewBroadcastString() *BroadcastString {
	b := &BroadcastString{
		msgC:    make(chan string, 20), // in case of cmd pipe error
		writers: make(map[chan string]bool, 0),
	}
	go func() {
		for message := range b.msgC {
			b.writeToAll(message)
		}
	}()
	return b
}

func (b *BroadcastString) writeToAll(message string) {
	for c := range b.writers {
		select {
		case c <- message:
		case <-time.After(500 * time.Millisecond):
			log.Println("channel closed, remove from queue")
			delete(b.writers, c)
		}
	}
}

func (b *BroadcastString) WriteMessage(message string) {
	select {
	case b.msgC <- message:
	default:
	}
}

func (b *BroadcastString) Reset() {
	b.msgC = make(chan string, 20)
}

func (b *BroadcastString) AddListener(c chan string) chan string {
	b.mu.Lock()
	defer b.mu.Unlock()
	if c == nil {
		c = make(chan string, 0)
	}
	b.writers[c] = true
	return c
}

func (b *BroadcastString) RemoveListener(c chan string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	delete(b.writers, c)
}

func (b *BroadcastString) Close() {
	close(b.msgC)
}

type BufferBroadcast struct {
	bs *BroadcastString

	maxSize int
	buf     *rbuf.FixedSizeRingBuf // *bytes.Buffer
	mu      sync.Mutex
}

func NewBufferBroadcast(size int) *BufferBroadcast {
	if size <= 0 {
		size = 4 * 1024 // 4K
	}
	bufb := &BufferBroadcast{
		maxSize: size,
		bs:      NewBroadcastString(),
		buf:     rbuf.NewFixedSizeRingBuf(size), //  bytes.NewBuffer(nil), // buffer.NewRing(buffer.New(size)),
	}
	bufC := bufb.bs.AddListener(nil)
	go func() {
		for msg := range bufC {
			bufb.buf.Write([]byte(msg))
		}
	}()
	return bufb
}

func (b *BufferBroadcast) Write(data []byte) (n int, err error) {
	b.bs.WriteMessage(string(data)) // should return immediatiely, in case of pipe error
	return len(data), nil
}

func (b *BufferBroadcast) Reset() {
	b.buf.Reset()
	b.bs.Reset()
}

func (b *BufferBroadcast) AddHookFunc(wf func(string) error) chan error {
	b.mu.Lock()
	defer b.mu.Unlock()
	c := b.bs.AddListener(nil)
	errC := make(chan error, 1)
	go func() {
		data := b.buf.Bytes()
		// data, _ := ioutil.ReadAll(b.buf)
		if err := wf(string(data)); err != nil {
			errC <- err
			return
		}
		for msg := range c {
			err := wf(msg)
			if err != nil {
				errC <- err
				break
			}
		}
	}()
	return errC
}
<|endoftext|>


<commit_before>package models

import (
	"errors"

	// "fmt"
	"time"

	"github.com/koding/bongo"
)

type NotificationContent struct {
	// unique identifier of NotificationContent
	Id int64 `json:"id"`

	// target of the activity (replied messageId, followed accountId etc.)
	TargetId int64 `json:"targetId,string"   sql:"NOT NULL"`

	// Type of the NotificationContent
	TypeConstant string `json:"typeConstant" sql:"NOT NULL;TYPE:VARCHAR(100);"`

	// Creation date of the NotificationContent
	CreatedAt time.Time `json:"createdAt"`
}

const (
	// NotificationContent Types
	NotificationContent_TYPE_LIKE    = "like"
	NotificationContent_TYPE_COMMENT = "comment"
	NotificationContent_TYPE_MENTION = "mention"
	NotificationContent_TYPE_PM      = "pm"
)

func (n *NotificationContent) FindByTarget() error {
	s := map[string]interface{}{
		"type_constant": n.TypeConstant,
		"target_id":     n.TargetId,
	}
	q := bongo.NewQS(s)

	return n.One(q)
}

// CreateNotification validates notifiable instance and creates a new notification
// with actor activity.
func CreateNotificationContent(i Notifiable) (*NotificationContent, error) {
	// first check for type constant and target id
	if i.GetType() == "" {
		return nil, errors.New("Type must be set")
	}

	if i.GetTargetId() == 0 {
		return nil, errors.New("TargetId must be set")
	}

	if i.GetActorId() == 0 {
		return nil, errors.New("ActorId must be set")
	}

	// check for previous NotificationContent create if it does not exist (type:comment targetId:messageId)
	nc := NewNotificationContent()
	nc.TypeConstant = i.GetType()
	nc.TargetId = i.GetTargetId()

	if err := nc.Create(); err != nil {
		return nil, err
	}
	a := NewNotificationActivity()
	a.NotificationContentId = nc.Id
	a.ActorId = i.GetActorId()
	a.MessageId = i.GetMessageId()

	if err := a.Create(); err != nil {
		return nil, err
	}

	return nc, nil
}

// FetchByIds fetches notification contents with given ids
func (n *NotificationContent) FetchByIds(ids []int64) ([]NotificationContent, error) {
	notificationContents := make([]NotificationContent, 0)
	if err := bongo.B.FetchByIds(n, &notificationContents, ids); err != nil {
		return nil, err
	}
	return notificationContents, nil
}

// FetchMapByIds returns NotificationContent map with given ids
func (n *NotificationContent) FetchMapByIds(ids []int64) (map[int64]NotificationContent, error) {
	ncList, err := n.FetchByIds(ids)
	if err != nil {
		return nil, err
	}

	ncMap := make(map[int64]NotificationContent, 0)
	for _, nc := range ncList {
		ncMap[nc.Id] = nc
	}

	return ncMap, nil
}

// CreateNotificationType creates an instance of notifiable subclasses
func CreateNotificationContentType(notificationType string) (Notifiable, error) {
	switch notificationType {
	case NotificationContent_TYPE_LIKE:
		return NewInteractionNotification(notificationType), nil
	case NotificationContent_TYPE_COMMENT:
		return NewReplyNotification(), nil
	case NotificationContent_TYPE_MENTION:
		return NewMentionNotification(), nil
	case NotificationContent_TYPE_PM:
		return NewPMNotification(), nil
	default:
		return nil, errors.New("undefined notification type")
	}

}

func (n *NotificationContent) GetContentType() (Notifiable, error) {
	return CreateNotificationContentType(n.TypeConstant)
}

func (n *NotificationContent) GetDefinition() string {
	nt, err := CreateNotificationContentType(n.TypeConstant)
	if err != nil {
		return ""
	}

	return nt.GetDefinition()
}
<commit_msg>Notification: Add FetchIdsByTarget method to notification content<commit_after>package models

import (
	"errors"

	// "fmt"
	"time"

	"github.com/koding/bongo"
)

type NotificationContent struct {
	// unique identifier of NotificationContent
	Id int64 `json:"id"`

	// target of the activity (replied messageId, followed accountId etc.)
	TargetId int64 `json:"targetId,string"   sql:"NOT NULL"`

	// Type of the NotificationContent
	TypeConstant string `json:"typeConstant" sql:"NOT NULL;TYPE:VARCHAR(100);"`

	// Creation date of the NotificationContent
	CreatedAt time.Time `json:"createdAt"`
}

const (
	// NotificationContent Types
	NotificationContent_TYPE_LIKE    = "like"
	NotificationContent_TYPE_COMMENT = "comment"
	NotificationContent_TYPE_MENTION = "mention"
	NotificationContent_TYPE_PM      = "pm"
)

func (n *NotificationContent) FindByTarget() error {
	s := map[string]interface{}{
		"type_constant": n.TypeConstant,
		"target_id":     n.TargetId,
	}
	q := bongo.NewQS(s)

	return n.One(q)
}

// CreateNotification validates notifiable instance and creates a new notification
// with actor activity.
func CreateNotificationContent(i Notifiable) (*NotificationContent, error) {
	// first check for type constant and target id
	if i.GetType() == "" {
		return nil, errors.New("Type must be set")
	}

	if i.GetTargetId() == 0 {
		return nil, errors.New("TargetId must be set")
	}

	if i.GetActorId() == 0 {
		return nil, errors.New("ActorId must be set")
	}

	// check for previous NotificationContent create if it does not exist (type:comment targetId:messageId)
	nc := NewNotificationContent()
	nc.TypeConstant = i.GetType()
	nc.TargetId = i.GetTargetId()

	if err := nc.Create(); err != nil {
		return nil, err
	}
	a := NewNotificationActivity()
	a.NotificationContentId = nc.Id
	a.ActorId = i.GetActorId()
	a.MessageId = i.GetMessageId()

	if err := a.Create(); err != nil {
		return nil, err
	}

	return nc, nil
}

// FetchByIds fetches notification contents with given ids
func (n *NotificationContent) FetchByIds(ids []int64) ([]NotificationContent, error) {
	notificationContents := make([]NotificationContent, 0)
	if err := bongo.B.FetchByIds(n, &notificationContents, ids); err != nil {
		return nil, err
	}
	return notificationContents, nil
}

// FetchMapByIds returns NotificationContent map with given ids
func (n *NotificationContent) FetchMapByIds(ids []int64) (map[int64]NotificationContent, error) {
	ncList, err := n.FetchByIds(ids)
	if err != nil {
		return nil, err
	}

	ncMap := make(map[int64]NotificationContent, 0)
	for _, nc := range ncList {
		ncMap[nc.Id] = nc
	}

	return ncMap, nil
}

func (n *NotificationContent) FetchIdsByTargetId(targetId int64) ([]int64, error) {
	var ids []int64

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"target_id": targetId,
		},
		Pluck: "id",
	}

	return ids, n.Some(&ids, query)
}

// CreateNotificationType creates an instance of notifiable subclasses
func CreateNotificationContentType(notificationType string) (Notifiable, error) {
	switch notificationType {
	case NotificationContent_TYPE_LIKE:
		return NewInteractionNotification(notificationType), nil
	case NotificationContent_TYPE_COMMENT:
		return NewReplyNotification(), nil
	case NotificationContent_TYPE_MENTION:
		return NewMentionNotification(), nil
	case NotificationContent_TYPE_PM:
		return NewPMNotification(), nil
	default:
		return nil, errors.New("undefined notification type")
	}

}

func (n *NotificationContent) GetContentType() (Notifiable, error) {
	return CreateNotificationContentType(n.TypeConstant)
}

func (n *NotificationContent) GetDefinition() string {
	nt, err := CreateNotificationContentType(n.TypeConstant)
	if err != nil {
		return ""
	}

	return nt.GetDefinition()
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package master

import (
	"context"
	"flag"
	"os"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/mysql"

	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance *cluster.LocalProcessCluster
	vtParams        mysql.ConnParams
	hostname        = "localhost"
	keyspaceName    = "ks"
	cell            = "zone1"
	sqlSchema       = `
	create table t1(
		id bigint,
		value varchar(16),
		primary key(id)
	) Engine=InnoDB;
`

	vSchema = `
	{
    "sharded": true,
    "vindexes": {
      "hash": {
        "type": "hash"
      }
    },
    "tables": {
      "t1": {
        "column_vindexes": [
          {
            "column": "id",
            "name": "hash"
          }
        ]
      }
    }
  }`
)

const dbName = "myDbName"

func TestMain(m *testing.M) {
	flag.Parse()

	exitCode := func() int {
		clusterInstance = cluster.NewCluster(cell, hostname)
		defer clusterInstance.Teardown()

		// Start topo server
		err := clusterInstance.StartTopo()
		if err != nil {
			return 1
		}

		// Set extra tablet args for lock timeout
		clusterInstance.VtTabletExtraArgs = []string{
			"-init_db_name_override", dbName,
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name:      keyspaceName,
			SchemaSQL: sqlSchema,
			VSchema:   vSchema,
		}

		if err = clusterInstance.StartUnshardedKeyspace(*keyspace, 1, false); err != nil {
			return 1
		}

		if err = clusterInstance.StartVtgate(); err != nil {
			return 1
		}
		vtParams = mysql.ConnParams{
			Host: clusterInstance.Hostname,
			Port: clusterInstance.VtgateMySQLPort,
		}

		return m.Run()
	}()
	os.Exit(exitCode)
}

func TestDbNameOverride(t *testing.T) {
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()
	qr, err := conn.ExecuteFetch("SELECT database() FROM information_schema.tables WHERE table_schema = database()", 1000, true)

	require.NoError(t, err)
	require.Equal(t, 1, len(qr.Rows), "did not get enough rows back")
	require.Equal(t, strings.ToLower(dbName), qr.Rows[0][0].ToString())
}
<commit_msg>remove ToLower from expected value, since the test uses camel case for db name<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package master

import (
	"context"
	"flag"
	"os"
	"testing"

	"github.com/stretchr/testify/require"
	"vitess.io/vitess/go/mysql"

	"vitess.io/vitess/go/test/endtoend/cluster"
)

var (
	clusterInstance *cluster.LocalProcessCluster
	vtParams        mysql.ConnParams
	hostname        = "localhost"
	keyspaceName    = "ks"
	cell            = "zone1"
	sqlSchema       = `
	create table t1(
		id bigint,
		value varchar(16),
		primary key(id)
	) Engine=InnoDB;
`

	vSchema = `
	{
    "sharded": true,
    "vindexes": {
      "hash": {
        "type": "hash"
      }
    },
    "tables": {
      "t1": {
        "column_vindexes": [
          {
            "column": "id",
            "name": "hash"
          }
        ]
      }
    }
  }`
)

const dbName = "myDbName"

func TestMain(m *testing.M) {
	flag.Parse()

	exitCode := func() int {
		clusterInstance = cluster.NewCluster(cell, hostname)
		defer clusterInstance.Teardown()

		// Start topo server
		err := clusterInstance.StartTopo()
		if err != nil {
			return 1
		}

		// Set extra tablet args for lock timeout
		clusterInstance.VtTabletExtraArgs = []string{
			"-init_db_name_override", dbName,
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name:      keyspaceName,
			SchemaSQL: sqlSchema,
			VSchema:   vSchema,
		}

		if err = clusterInstance.StartUnshardedKeyspace(*keyspace, 1, false); err != nil {
			return 1
		}

		if err = clusterInstance.StartVtgate(); err != nil {
			return 1
		}
		vtParams = mysql.ConnParams{
			Host: clusterInstance.Hostname,
			Port: clusterInstance.VtgateMySQLPort,
		}

		return m.Run()
	}()
	os.Exit(exitCode)
}

func TestDbNameOverride(t *testing.T) {
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()
	qr, err := conn.ExecuteFetch("SELECT database() FROM information_schema.tables WHERE table_schema = database()", 1000, true)

	require.NoError(t, err)
	require.Equal(t, 1, len(qr.Rows), "did not get enough rows back")
	require.Equal(t, dbName, qr.Rows[0][0].ToString())
}
<|endoftext|>


<commit_before>/* Copyright (c) 2016 PaddlePaddle Authors All Rights Reserve.

   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 autoscaler

import (
	"fmt"
	"sort"
	"time"

	// TODO(typhoonzero): this package still depends on k8s API, try to remove this.
	"k8s.io/apimachinery/pkg/api/resource"
	batchv1 "k8s.io/client-go/pkg/apis/batch/v1"

	paddlejob "github.com/PaddlePaddle/cloud/go/api"
	log "github.com/inconshreveable/log15"
)

const (
	defaultLoopDur = time.Second * 5
)

// ClusterResource is the resource of a cluster
type ClusterResource struct {
	NodeCount int

	GPURequest int
	GPULimit   int
	GPUTotal   int

	CPURequestMilli int64
	CPULimitMilli   int64
	CPUTotalMilli   int64

	MemoryRequestMega int64
	MemoryLimitMega   int64
	MemoryTotalMega   int64
}

// Cluster represents the cluster managment system such as Kubernetes.
type Cluster interface {
	// SyncResource will sync resource values with the cluster.
	// should call this function in every tick.
	SyncResource() (ClusterResource, error)

	// GetTrainerJob gets the trainer job spec.
	GetTrainerJob(job *paddlejob.TrainingJob) (*batchv1.Job, error)

	// UpdateTrainerJob updates the trainer job spec.
	UpdateTrainerJob(job *batchv1.Job) error

	// JobPods returns the number total desired pods and the
	// number of running pods of a job.
	JobPods(job *paddlejob.TrainingJob) (int, int, error)
}

type job struct {
	Config     *paddlejob.TrainingJob
	TrainerJob *batchv1.Job
}

func (j job) TrainerGPULimit() int {
	q := j.Config.Spec.Trainer.Resources.Limits.NvidiaGPU()
	return int(q.Value())
}

func (j job) TrainerCPURequestMilli() int64 {
	q := j.Config.Spec.Trainer.Resources.Requests.Cpu()
	return q.ScaledValue(resource.Milli)
}

func (j job) TrainerMemRequestMega() int64 {
	q := j.Config.Spec.Trainer.Resources.Requests.Memory()
	return q.ScaledValue(resource.Mega)
}

func (j job) Fulfillment() float64 {
	minInstance := j.Config.Spec.Trainer.MinInstance
	maxInstance := j.Config.Spec.Trainer.MaxInstance

	if minInstance == maxInstance {
		return 1
	}

	curInstance := int(*j.TrainerJob.Spec.Parallelism)
	return float64(curInstance-minInstance) / float64(maxInstance-minInstance)
}

// Autoscaler launches and scales the training jobs.
type Autoscaler struct {
	ticker  *time.Ticker
	cluster Cluster
	jobs    map[string]job
	eventCh chan event
}

// New creates a new Autoscaler.
func New(cluster Cluster, options ...func(*Autoscaler)) *Autoscaler {
	c := &Autoscaler{
		cluster: cluster,
		ticker:  time.NewTicker(defaultLoopDur),
		jobs:    make(map[string]job),
		eventCh: make(chan event),
	}
	for _, option := range options {
		option(c)
	}
	return c
}

type jobs []job

func (j jobs) Len() int {
	return len(j)
}

func (j jobs) Less(a int, b int) bool {
	scoreA := j[a].Fulfillment()
	scoreB := j[b].Fulfillment()

	if scoreA == scoreB {
		resA := j[a].Config.Spec.Trainer.Resources
		resB := j[b].Config.Spec.Trainer.Resources
		resALimitsGPU := *resA.Limits.NvidiaGPU()
		resBLimitsGPU := *resB.Limits.NvidiaGPU()
		if resALimitsGPU.Cmp(resBLimitsGPU) == 0 {
			resARequestsCPU := *resA.Requests.Cpu()
			resBRequestsCPU := *resB.Requests.Cpu()
			if resARequestsCPU.Cmp(resBRequestsCPU) == 0 {
				resARequestsMem := *resA.Requests.Memory()
				resBRequestsMem := *resB.Requests.Memory()
				return resARequestsMem.Cmp(resBRequestsMem) == -1
			}
			return resARequestsCPU.Cmp(resBRequestsCPU) == -1
		}
		return resALimitsGPU.Cmp(resBLimitsGPU) == -1
	}
	return scoreA < scoreB
}

func (j jobs) Swap(a int, b int) {
	j[a], j[b] = j[b], j[a]
}

// elastic job filter.
func elastic(j job) bool {
	return j.Config.Elastic()
}

// gpu job filter.
func gpu(j job) bool {
	return j.Config.NeedGPU()
}

type eventType int

const (
	add eventType = iota
	del
	update
)

type event struct {
	Type eventType
	Job  *paddlejob.TrainingJob
}

func (e event) GoString() string {
	return fmt.Sprintf("event{type: %d, job name: %s}", e.Type, e.Job.Name)
}

// OnAdd notifies the autoscaler that a job has been added.
func (a *Autoscaler) OnAdd(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: add, Job: trainingjob}
}

// OnDel notifies the autoscaler that a job has been deleted.
func (a *Autoscaler) OnDel(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: del, Job: trainingjob}
}

// OnUpdate notifies the autoscaler that a job has been deleted.
func (a *Autoscaler) OnUpdate(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: update, Job: trainingjob}
}

// sortedJobs return the names of sorted jobs by fulfillment and
// tiebreakers in ascending order.
func sortedJobs(j []job, filters ...func(job) bool) []job {
	var js jobs
nextJob:
	for _, v := range j {
		for _, f := range filters {
			if !f(v) {
				continue nextJob
			}
		}
		js = append(js, v)
	}

	sort.Sort(js)
	return js
}

func scaleDryRun(r *ClusterResource, j job, curDiff int, scaleDown bool) (additional int) {
	additionalGPUInstance := 0
	additionalCPUInstance := 0
	cpuRequestMilli := j.TrainerCPURequestMilli()
	memRequestMega := j.TrainerMemRequestMega()
	gpuLimit := j.TrainerGPULimit()

	// Adjust resource upon return.
	defer func() {
		r.GPULimit += gpuLimit * additional
		r.CPURequestMilli += cpuRequestMilli * int64(additional)
		r.MemoryRequestMega += memRequestMega * int64(additional)
	}()

	if r.GPULimit > r.GPUTotal || r.CPURequestMilli > r.CPUTotalMilli {
		return -1
	}

	// TODO(helin): j.TrainerJob.Spec.Parallelism may not reflect
	// the actual pod running for the trainer job. We need to
	// count the pod manually. And calculate the additional value
	// based on the running pod count,
	// j.TrainerJob.Spec.Parallelism, and curDiff.
	plannedInstance := int(*j.TrainerJob.Spec.Parallelism) + curDiff
	instanceMax := j.Config.Spec.Trainer.MaxInstance

	if plannedInstance >= instanceMax {
		// Do not scale or scale down, don't need to check if
		// there are available free resources.
		additional = instanceMax - plannedInstance
		return
	}

	if r.MemoryTotalMega-r.MemoryRequestMega <= memRequestMega {
		// insufficient memory, do not scale
		additional = 0
		return
	}

	if r.CPUTotalMilli-r.CPURequestMilli >= cpuRequestMilli {
		additionalCPUInstance = 1
	}

	needGPU := gpuLimit > 0
	if needGPU && r.GPUTotal-r.GPULimit >= gpuLimit {
		additionalGPUInstance = 1
	}

	if needGPU {
		if additionalGPUInstance < additionalCPUInstance {
			additional = additionalGPUInstance
		} else {
			additional = additionalCPUInstance
		}
	} else {
		additional = additionalCPUInstance
	}

	return
}

func scaleAllDryRun(jobs []job, r ClusterResource) map[string]int {
	// Iteratively calculate scaling diff until nothing changes.
	diff := make(map[string]int)
	for {
		noChange := true
		sorted := sortedJobs(jobs, elastic)
		dryRun := func(j job, scaleDirection bool) {
			name := j.Config.Name
			additional := scaleDryRun(&r, j, diff[name], scaleDirection)
			log.Debug(
				"dry run scale job",
				"name", name, "current scale difference", diff[name],
				"scale up number of instances", additional, "cluster resource", r,
			)
			diff[name] += additional

			if additional != 0 {
				noChange = false
			}
		}

		// TODO(typhoonzero): implement GPU priority CFS scheduler from here.

		// scale up from the ones that need scaling up the
		// most.
		for _, j := range sorted {
			dryRun(j, false)
		}

		// scale down from the ones that need scaling up the
		// least.
		for i := len(sorted) - 1; i >= 0; i-- {
			dryRun(sorted[i], true)
		}

		if noChange {
			break
		}
	}

	return diff
}

func (a *Autoscaler) scaleAll(diff map[string]int) {
	for name := range diff {
		if diff[name] != 0 {
			log.Info("scaling job",
				"name", name, "number of instances", diff[name])
			target := *a.jobs[name].TrainerJob.Spec.Parallelism + int32(diff[name])

			var err error
			for retry := 0; retry < 3; retry++ {
				var tj *batchv1.Job
				// don't shadow err
				tj, err = a.cluster.GetTrainerJob(a.jobs[name].Config)
				if err != nil {
					log.Warn("sync trainer job error",
						"error", err, "remaining retry", retry)
					continue
				}
				j := a.jobs[name]
				// NOTE: only update batchv1.job from k8s api-server before updating
				// for efficiency. Update the job resource to get latest k8s
				// resource reversion.
				j.TrainerJob = tj
				a.jobs[name] = j
				*a.jobs[name].TrainerJob.Spec.Parallelism = target
				err = a.cluster.UpdateTrainerJob(a.jobs[name].TrainerJob)
				if err != nil {
					log.Error("error updating trainer job",
						"error", err, "remaining retry", retry)
					continue
				}

				break
			}

			if err != nil {
				log.Warn("Error updating trainer job", "error", err)
			}
		}
	}
}

// Monitor monitors the cluster resources and training jobs in a loop,
// scales the training jobs according to the cluster resource.
func (a *Autoscaler) Monitor() {
	for {
		select {
		case <-a.ticker.C:
		case e := <-a.eventCh:
			log.Debug("monitor received event", "event", e)
			switch e.Type {
			case add:
				fallthrough
			case update:
				// TODO(helin): schedule the training
				// k8s Job. Currently we don't
				// schedule the trainer job, but only
				// scale it.
				var tj *batchv1.Job
				var err error
				for {
					tj, err = a.cluster.GetTrainerJob(e.Job)
					if err == nil {
						break
					}

					log.Error(
						"Error getting the trainer k8s job, retrying soon.",
						"name", e.Job.ObjectMeta.Name,
					)
					time.Sleep(3 * time.Second)
				}

				j := job{
					Config:     e.Job,
					TrainerJob: tj,
				}
				a.jobs[e.Job.ObjectMeta.Name] = j
			case del:
				// TODO(helin): delete all created
				// resources (e.g., trainer Job,
				// pserver Replica Set) when we
				// schedules the resources.
				delete(a.jobs, e.Job.ObjectMeta.Name)
			default:
				log.Error("unrecognized event", "event", e)
			}
		}

		r, err := a.cluster.SyncResource()
		if err != nil {
			log.Error("error sync resource", "error", err)
			continue
		}

		log.Info("latest cluster resource", "resource", r)
		var js []job
		for _, j := range a.jobs {
			// Scale jobs only when it's running (defined
			// by all pods are in the "running"
			// status). Pods are
			// pending/starting/terminating if the job is
			// just submited or just scaled up/down.
			total, running, err := a.cluster.JobPods(j.Config)
			if err != nil {
				log.Error("check if job is running failed", "error", err)
				continue
			}

			if total == running {
				js = append(js, j)
			}
		}
		diff := scaleAllDryRun(js, r)
		log.Info("calculated scaling plan", "plan", diff)
		a.scaleAll(diff)
	}
}
<commit_msg>Add retry limit when get trainer job errors, improve error logging.<commit_after>/* Copyright (c) 2016 PaddlePaddle Authors All Rights Reserve.

   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 autoscaler

import (
	"fmt"
	"sort"
	"time"

	// TODO(typhoonzero): this package still depends on k8s API, try to remove this.
	"k8s.io/apimachinery/pkg/api/resource"
	batchv1 "k8s.io/client-go/pkg/apis/batch/v1"

	paddlejob "github.com/PaddlePaddle/cloud/go/api"
	log "github.com/inconshreveable/log15"
)

const (
	defaultLoopDur = time.Second * 5
)

// ClusterResource is the resource of a cluster
type ClusterResource struct {
	NodeCount int

	GPURequest int
	GPULimit   int
	GPUTotal   int

	CPURequestMilli int64
	CPULimitMilli   int64
	CPUTotalMilli   int64

	MemoryRequestMega int64
	MemoryLimitMega   int64
	MemoryTotalMega   int64
}

// Cluster represents the cluster managment system such as Kubernetes.
type Cluster interface {
	// SyncResource will sync resource values with the cluster.
	// should call this function in every tick.
	SyncResource() (ClusterResource, error)

	// GetTrainerJob gets the trainer job spec.
	GetTrainerJob(job *paddlejob.TrainingJob) (*batchv1.Job, error)

	// UpdateTrainerJob updates the trainer job spec.
	UpdateTrainerJob(job *batchv1.Job) error

	// JobPods returns the number total desired pods and the
	// number of running pods of a job.
	JobPods(job *paddlejob.TrainingJob) (int, int, error)
}

type job struct {
	Config     *paddlejob.TrainingJob
	TrainerJob *batchv1.Job
}

func (j job) TrainerGPULimit() int {
	q := j.Config.Spec.Trainer.Resources.Limits.NvidiaGPU()
	return int(q.Value())
}

func (j job) TrainerCPURequestMilli() int64 {
	q := j.Config.Spec.Trainer.Resources.Requests.Cpu()
	return q.ScaledValue(resource.Milli)
}

func (j job) TrainerMemRequestMega() int64 {
	q := j.Config.Spec.Trainer.Resources.Requests.Memory()
	return q.ScaledValue(resource.Mega)
}

func (j job) Fulfillment() float64 {
	minInstance := j.Config.Spec.Trainer.MinInstance
	maxInstance := j.Config.Spec.Trainer.MaxInstance

	if minInstance == maxInstance {
		return 1
	}

	curInstance := int(*j.TrainerJob.Spec.Parallelism)
	return float64(curInstance-minInstance) / float64(maxInstance-minInstance)
}

// Autoscaler launches and scales the training jobs.
type Autoscaler struct {
	ticker  *time.Ticker
	cluster Cluster
	jobs    map[string]job
	eventCh chan event
}

// New creates a new Autoscaler.
func New(cluster Cluster, options ...func(*Autoscaler)) *Autoscaler {
	c := &Autoscaler{
		cluster: cluster,
		ticker:  time.NewTicker(defaultLoopDur),
		jobs:    make(map[string]job),
		eventCh: make(chan event),
	}
	for _, option := range options {
		option(c)
	}
	return c
}

type jobs []job

func (j jobs) Len() int {
	return len(j)
}

func (j jobs) Less(a int, b int) bool {
	scoreA := j[a].Fulfillment()
	scoreB := j[b].Fulfillment()

	if scoreA == scoreB {
		resA := j[a].Config.Spec.Trainer.Resources
		resB := j[b].Config.Spec.Trainer.Resources
		resALimitsGPU := *resA.Limits.NvidiaGPU()
		resBLimitsGPU := *resB.Limits.NvidiaGPU()
		if resALimitsGPU.Cmp(resBLimitsGPU) == 0 {
			resARequestsCPU := *resA.Requests.Cpu()
			resBRequestsCPU := *resB.Requests.Cpu()
			if resARequestsCPU.Cmp(resBRequestsCPU) == 0 {
				resARequestsMem := *resA.Requests.Memory()
				resBRequestsMem := *resB.Requests.Memory()
				return resARequestsMem.Cmp(resBRequestsMem) == -1
			}
			return resARequestsCPU.Cmp(resBRequestsCPU) == -1
		}
		return resALimitsGPU.Cmp(resBLimitsGPU) == -1
	}
	return scoreA < scoreB
}

func (j jobs) Swap(a int, b int) {
	j[a], j[b] = j[b], j[a]
}

// elastic job filter.
func elastic(j job) bool {
	return j.Config.Elastic()
}

// gpu job filter.
func gpu(j job) bool {
	return j.Config.NeedGPU()
}

type eventType int

const (
	add eventType = iota
	del
	update
)

type event struct {
	Type eventType
	Job  *paddlejob.TrainingJob
}

func (e event) GoString() string {
	return fmt.Sprintf("event{type: %d, job name: %s}", e.Type, e.Job.Name)
}

// OnAdd notifies the autoscaler that a job has been added.
func (a *Autoscaler) OnAdd(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: add, Job: trainingjob}
}

// OnDel notifies the autoscaler that a job has been deleted.
func (a *Autoscaler) OnDel(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: del, Job: trainingjob}
}

// OnUpdate notifies the autoscaler that a job has been deleted.
func (a *Autoscaler) OnUpdate(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: update, Job: trainingjob}
}

// sortedJobs return the names of sorted jobs by fulfillment and
// tiebreakers in ascending order.
func sortedJobs(j []job, filters ...func(job) bool) []job {
	var js jobs
nextJob:
	for _, v := range j {
		for _, f := range filters {
			if !f(v) {
				continue nextJob
			}
		}
		js = append(js, v)
	}

	sort.Sort(js)
	return js
}

func scaleDryRun(r *ClusterResource, j job, curDiff int, scaleDown bool) (additional int) {
	additionalGPUInstance := 0
	additionalCPUInstance := 0
	cpuRequestMilli := j.TrainerCPURequestMilli()
	memRequestMega := j.TrainerMemRequestMega()
	gpuLimit := j.TrainerGPULimit()

	// Adjust resource upon return.
	defer func() {
		r.GPULimit += gpuLimit * additional
		r.CPURequestMilli += cpuRequestMilli * int64(additional)
		r.MemoryRequestMega += memRequestMega * int64(additional)
	}()

	if r.GPULimit > r.GPUTotal || r.CPURequestMilli > r.CPUTotalMilli {
		return -1
	}

	// TODO(helin): j.TrainerJob.Spec.Parallelism may not reflect
	// the actual pod running for the trainer job. We need to
	// count the pod manually. And calculate the additional value
	// based on the running pod count,
	// j.TrainerJob.Spec.Parallelism, and curDiff.
	plannedInstance := int(*j.TrainerJob.Spec.Parallelism) + curDiff
	instanceMax := j.Config.Spec.Trainer.MaxInstance

	if plannedInstance >= instanceMax {
		// Do not scale or scale down, don't need to check if
		// there are available free resources.
		additional = instanceMax - plannedInstance
		return
	}

	if r.MemoryTotalMega-r.MemoryRequestMega <= memRequestMega {
		// insufficient memory, do not scale
		additional = 0
		return
	}

	if r.CPUTotalMilli-r.CPURequestMilli >= cpuRequestMilli {
		additionalCPUInstance = 1
	}

	needGPU := gpuLimit > 0
	if needGPU && r.GPUTotal-r.GPULimit >= gpuLimit {
		additionalGPUInstance = 1
	}

	if needGPU {
		if additionalGPUInstance < additionalCPUInstance {
			additional = additionalGPUInstance
		} else {
			additional = additionalCPUInstance
		}
	} else {
		additional = additionalCPUInstance
	}

	return
}

func scaleAllDryRun(jobs []job, r ClusterResource) map[string]int {
	// Iteratively calculate scaling diff until nothing changes.
	diff := make(map[string]int)
	for {
		noChange := true
		sorted := sortedJobs(jobs, elastic)
		dryRun := func(j job, scaleDirection bool) {
			name := j.Config.Name
			additional := scaleDryRun(&r, j, diff[name], scaleDirection)
			log.Debug(
				"dry run scale job",
				"name", name, "current scale difference", diff[name],
				"scale up number of instances", additional, "cluster resource", r,
			)
			diff[name] += additional

			if additional != 0 {
				noChange = false
			}
		}

		// TODO(typhoonzero): implement GPU priority CFS scheduler from here.

		// scale up from the ones that need scaling up the
		// most.
		for _, j := range sorted {
			dryRun(j, false)
		}

		// scale down from the ones that need scaling up the
		// least.
		for i := len(sorted) - 1; i >= 0; i-- {
			dryRun(sorted[i], true)
		}

		if noChange {
			break
		}
	}

	return diff
}

func (a *Autoscaler) scaleAll(diff map[string]int) {
	for name := range diff {
		if diff[name] != 0 {
			log.Info("scaling job",
				"name", name, "number of instances", diff[name])
			target := *a.jobs[name].TrainerJob.Spec.Parallelism + int32(diff[name])

			var err error
			for retry := 5; retry > 0; retry-- {
				var tj *batchv1.Job
				// don't shadow err
				tj, err = a.cluster.GetTrainerJob(a.jobs[name].Config)
				if err != nil {
					log.Warn("sync trainer job error",
						"error", err, "remaining retry", retry)
					continue
				}
				j := a.jobs[name]
				// NOTE: only update batchv1.job from k8s api-server before updating
				// for efficiency. Update the job resource to get latest k8s
				// resource reversion.
				j.TrainerJob = tj
				a.jobs[name] = j
				*a.jobs[name].TrainerJob.Spec.Parallelism = target
				err = a.cluster.UpdateTrainerJob(a.jobs[name].TrainerJob)
				if err != nil {
					log.Error("error updating trainer job",
						"error", err, "remaining retry", retry)
					continue
				}

				break
			}

			if err != nil {
				log.Warn("Error updating trainer job", "error", err)
			}
		}
	}
}

// Monitor monitors the cluster resources and training jobs in a loop,
// scales the training jobs according to the cluster resource.
func (a *Autoscaler) Monitor() {
	for {
		select {
		case <-a.ticker.C:
		case e := <-a.eventCh:
			log.Debug("monitor received event", "event", e)
			switch e.Type {
			case add:
				fallthrough
			case update:
				// TODO(helin): schedule the training
				// k8s Job. Currently we don't
				// schedule the trainer job, but only
				// scale it.
				var tj *batchv1.Job
				var err error
				for retry := 5; retry > 0; retry-- {
					tj, err = a.cluster.GetTrainerJob(e.Job)
					if err == nil {
						break
					}

					log.Error(
						"Error getting the trainer k8s job, retry if have retry remaining",
						"name", e.Job.ObjectMeta.Name,
						"retry remaining", retry,
						"error", err,
					)
					time.Sleep(3 * time.Second)
				}

				if err != nil {
					log.Error(
						"Error getting the trainer k8s job, skip event.",
						"name", e.Job.ObjectMeta.Name,
						"error", err,
					)
					continue
				}

				j := job{
					Config:     e.Job,
					TrainerJob: tj,
				}
				a.jobs[e.Job.ObjectMeta.Name] = j
			case del:
				// TODO(helin): delete all created
				// resources (e.g., trainer Job,
				// pserver Replica Set) when we
				// schedules the resources.
				delete(a.jobs, e.Job.ObjectMeta.Name)
			default:
				log.Error("unrecognized event", "event", e)
			}
		}

		r, err := a.cluster.SyncResource()
		if err != nil {
			log.Error("error sync resource", "error", err)
			continue
		}

		log.Info("latest cluster resource", "resource", r)
		var js []job
		for _, j := range a.jobs {
			// Scale jobs only when it's running (defined
			// by all pods are in the "running"
			// status). Pods are
			// pending/starting/terminating if the job is
			// just submited or just scaled up/down.
			total, running, err := a.cluster.JobPods(j.Config)
			if err != nil {
				log.Error("check if job is running failed", "error", err)
				continue
			}

			if total == running {
				js = append(js, j)
			}
		}
		diff := scaleAllDryRun(js, r)
		log.Info("calculated scaling plan", "plan", diff)
		a.scaleAll(diff)
	}
}
<|endoftext|>