<commit_before>// Package sentences provides a scanner for Unicode text segmentation sentence boundaries: https://unicode.org/reports/tr29/#Sentence_Boundaries
package sentences

import (
	"bufio"
	"bytes"
	"io"
	"unicode"
)

// NewScanner tokenizes a reader into a stream of sentence tokens according to Unicode Text Segmentation sentence boundaries https://unicode.org/reports/tr29/#Sentence_Boundaries
// Iterate through the stream by calling Scan() until false.
//	text := "This is an example. And another!"
//	reader := strings.NewReader(text)
//
//	scanner := sentences.NewScanner(reader)
//	for scanner.Scan() {
//		fmt.Printf("%s\n", scanner.Text())
//	}
//	if err := scanner.Err(); err != nil {
//		log.Fatal(err)
//	}
func NewScanner(r io.Reader) *Scanner {
	return &Scanner{
		incoming: bufio.NewReaderSize(r, 64*1024),
	}
}

// Scanner is the structure for scanning an input Reader. Use NewScanner to instantiate.
type Scanner struct {
	incoming *bufio.Reader

	// a buffer of runes to evaluate
	buffer []rune
	// a cursor for runes in the buffer
	pos int

	bb bytes.Buffer

	// outputs
	bytes []byte
	err   error
}

// reset creates a new bytes.Buffer on the Scanner, and clears previous values
func (sc *Scanner) reset() {
	// Drop the emitted runes (optimization to avoid growing array)
	copy(sc.buffer, sc.buffer[sc.pos:])
	sc.buffer = sc.buffer[:len(sc.buffer)-sc.pos]

	sc.pos = 0

	var bb bytes.Buffer
	sc.bb = bb

	sc.bytes = nil
	sc.err = nil
}

// Scan advances to the next token, returning true if successful. Returns false on error or EOF.
func (sc *Scanner) Scan() bool {
	sc.reset()

	for {
		// Fill the buffer with enough runes for lookahead
		for len(sc.buffer) < sc.pos+8 {
			r, eof, err := sc.readRune()
			if err != nil {
				sc.err = err
				return false
			}
			if eof {
				break
			}
			sc.buffer = append(sc.buffer, r)
		}

		// SB1
		sot := sc.pos == 0 // "start of text"
		eof := len(sc.buffer) == sc.pos
		if sot && !eof {
			sc.accept()
			continue
		}

		// SB2
		if eof {
			break
		}

		current := sc.buffer[sc.pos]
		previous := sc.buffer[sc.pos-1]

		// SB3
		if is(LF, current) && is(CR, previous) {
			sc.accept()
			continue
		}

		// SB4
		if is(_mergedParaSep, previous) {
			break
		}

		// SB5
		if is(_mergedExtendFormat, current) {
			sc.accept()
			continue
		}

		// SB6
		if is(Numeric, current) && sc.seekPrevious(sc.pos, ATerm) {
			sc.accept()
			continue
		}

		// SB7
		if is(Upper, current) {
			previousIndex := sc.seekPreviousIndex(sc.pos, ATerm)
			if previousIndex >= 0 && sc.seekPrevious(previousIndex, _mergedUpperLower) {
				sc.accept()
				continue
			}
		}

		// SB8
		{
			// This loop is the 'regex':
			// ( ¬(OLetter | Upper | Lower | ParaSep | SATerm) )*
			pos := sc.pos
			for pos < len(sc.buffer) {
				current := sc.buffer[pos]
				if is(_mergedOLetterUpperLowerParaSepSATerm, current) {
					break
				}
				pos++
			}

			if sc.seekForward(pos-1, Lower) {
				pos := sc.pos

				sp := pos
				for {
					sp = sc.seekPreviousIndex(sp, Sp)
					if sp < 0 {
						break
					}
					pos = sp
				}

				close := pos
				for {
					close = sc.seekPreviousIndex(close, Close)
					if close < 0 {
						break
					}
					pos = close
				}

				if sc.seekPrevious(pos, ATerm) {
					sc.accept()
					continue
				}
			}
		}

		// SB8a
		if is(_mergedSContinueSATerm, current) {
			pos := sc.pos

			sp := pos
			for {
				sp = sc.seekPreviousIndex(sp, Sp)
				if sp < 0 {
					break
				}
				pos = sp
			}

			close := pos
			for {
				close = sc.seekPreviousIndex(close, Close)
				if close < 0 {
					break
				}
				pos = close
			}

			if sc.seekPrevious(pos, _mergedSATerm) {
				sc.accept()
				continue
			}
		}

		// SB9
		if is(_mergedCloseSpParaSep, current) {
			pos := sc.pos

			close := pos
			for {
				close = sc.seekPreviousIndex(close, Close)
				if close < 0 {
					break
				}
				pos = close
			}

			if sc.seekPrevious(pos, _mergedSATerm) {
				sc.accept()
				continue
			}
		}

		// SB10
		if is(_mergedSpParaSep, current) {
			pos := sc.pos

			sp := pos
			for {
				sp = sc.seekPreviousIndex(sp, Sp)
				if sp < 0 {
					break
				}
				pos = sp
			}

			close := pos
			for {
				close = sc.seekPreviousIndex(close, Close)
				if close < 0 {
					break
				}
				pos = close
			}

			if sc.seekPrevious(pos, _mergedSATerm) {
				sc.accept()
				continue
			}
		}

		// SB11
		{
			pos := sc.pos

			ps := sc.seekPreviousIndex(pos, _mergedSpParaSep)
			if ps >= 0 {
				pos = ps
			}

			sp := pos
			for {
				sp = sc.seekPreviousIndex(sp, Sp)
				if sp < 0 {
					break
				}
				pos = sp
			}

			close := pos
			for {
				close = sc.seekPreviousIndex(close, Close)
				if close < 0 {
					break
				}
				pos = close
			}

			if sc.seekPrevious(pos, _mergedSATerm) {
				break
			}
		}

		// SB998
		if sc.pos > 0 {
			sc.accept()
			continue
		}

		// If we fall through all the above rules, it's a sentence break
		break
	}

	return sc.token()
}

// Bytes returns the current token as a byte slice, after a successful call to Scan
func (sc *Scanner) Bytes() []byte {
	return sc.bytes
}

// Text returns the current token, after a successful call to Scan
func (sc *Scanner) Text() string {
	return string(sc.bytes)
}

// Err returns the current error, after an unsuccessful call to Scan
func (sc *Scanner) Err() error {
	return sc.err
}

// Sentence boundary rules: https://unicode.org/reports/tr29/#Sentence_Boundaries
// In most cases, returning true means 'keep going'; check the name of the return var for clarity

var is = unicode.Is

// seekForward looks ahead until it hits a rune satisfying one of the range tables,
// ignoring Extend|Format
// See: https://unicode.org/reports/tr29/#Grapheme_Cluster_and_Format_Rules (driven by SB5)
func (sc *Scanner) seekForward(pos int, rts ...*unicode.RangeTable) bool {
	for i := pos + 1; i < len(sc.buffer); i++ {
		r := sc.buffer[i]

		// Ignore Extend|Format
		if is(_mergedExtendFormat, r) {
			continue
		}

		// See if any of the range tables apply
		for _, rt := range rts {
			if is(rt, r) {
				return true
			}
		}

		// If we get this far, it's not there
		break
	}

	return false
}

// seekPreviousIndex works backward until it hits a rune satisfying one of the range tables,
// ignoring Extend|Format, and returns the index of the rune in the buffer
// See: https://unicode.org/reports/tr29/#Grapheme_Cluster_and_Format_Rules (driven by SB5)
func (sc *Scanner) seekPreviousIndex(pos int, rts ...*unicode.RangeTable) int {
	// Start at the end of the buffer and move backwards
	for i := pos - 1; i >= 0; i-- {
		r := sc.buffer[i]

		// Ignore Extend|Format
		if is(_mergedExtendFormat, r) {
			continue
		}

		// See if any of the range tables apply
		for _, rt := range rts {
			if is(rt, r) {
				return i
			}
		}

		// If we get this far, it's not there
		break
	}

	return -1
}

// seekPreviousIndex works backward ahead until it hits a rune satisfying one of the range tables,
// ignoring Extend|Format, reporting success
// Logic is here: https://unicode.org/reports/tr29/#Grapheme_Cluster_and_Format_Rules (driven by SB5)
func (sc *Scanner) seekPrevious(pos int, rts ...*unicode.RangeTable) bool {
	return sc.seekPreviousIndex(pos, rts...) >= 0
}

func (sc *Scanner) token() bool {
	sc.bytes = sc.bb.Bytes()
	return len(sc.bytes) > 0
}

// accept forwards the buffer cursor (pos) by 1
func (sc *Scanner) accept() {
	sc.bb.WriteRune(sc.buffer[sc.pos])
	sc.pos++
}

// readRune gets the next rune, advancing the reader
func (sc *Scanner) readRune() (r rune, eof bool, err error) {
	r, _, err = sc.incoming.ReadRune()

	if err != nil {
		if err == io.EOF {
			return r, true, nil
		}
		return r, false, err
	}

	return r, false, nil
}
<commit_msg>bytes Buffer need not be member of Scanner<commit_after>// Package sentences provides a scanner for Unicode text segmentation sentence boundaries: https://unicode.org/reports/tr29/#Sentence_Boundaries
package sentences

import (
	"bufio"
	"bytes"
	"io"
	"unicode"
)

// NewScanner tokenizes a reader into a stream of sentence tokens according to Unicode Text Segmentation sentence boundaries https://unicode.org/reports/tr29/#Sentence_Boundaries
// Iterate through the stream by calling Scan() until false.
//	text := "This is an example. And another!"
//	reader := strings.NewReader(text)
//
//	scanner := sentences.NewScanner(reader)
//	for scanner.Scan() {
//		fmt.Printf("%s\n", scanner.Text())
//	}
//	if err := scanner.Err(); err != nil {
//		log.Fatal(err)
//	}
func NewScanner(r io.Reader) *Scanner {
	return &Scanner{
		incoming: bufio.NewReaderSize(r, 64*1024),
	}
}

// Scanner is the structure for scanning an input Reader. Use NewScanner to instantiate.
type Scanner struct {
	incoming *bufio.Reader

	// a buffer of runes to evaluate
	buffer []rune
	// a cursor for runes in the buffer
	pos int

	// outputs
	bytes []byte
	err   error
}

// reset creates a new bytes.Buffer on the Scanner, and clears previous values
func (sc *Scanner) reset() {
	// Drop the emitted runes (optimization to avoid growing array)
	copy(sc.buffer, sc.buffer[sc.pos:])
	sc.buffer = sc.buffer[:len(sc.buffer)-sc.pos]

	sc.pos = 0

	sc.bytes = nil
	sc.err = nil
}

// Scan advances to the next token, returning true if successful. Returns false on error or EOF.
func (sc *Scanner) Scan() bool {
	sc.reset()

	for {
		// Fill the buffer with enough runes for lookahead
		for len(sc.buffer) < sc.pos+8 {
			r, eof, err := sc.readRune()
			if err != nil {
				sc.err = err
				return false
			}
			if eof {
				break
			}
			sc.buffer = append(sc.buffer, r)
		}

		// SB1
		sot := sc.pos == 0 // "start of text"
		eof := len(sc.buffer) == sc.pos
		if sot && !eof {
			sc.accept()
			continue
		}

		// SB2
		if eof {
			break
		}

		current := sc.buffer[sc.pos]
		previous := sc.buffer[sc.pos-1]

		// SB3
		if is(LF, current) && is(CR, previous) {
			sc.accept()
			continue
		}

		// SB4
		if is(_mergedParaSep, previous) {
			break
		}

		// SB5
		if is(_mergedExtendFormat, current) {
			sc.accept()
			continue
		}

		// SB6
		if is(Numeric, current) && sc.seekPrevious(sc.pos, ATerm) {
			sc.accept()
			continue
		}

		// SB7
		if is(Upper, current) {
			previousIndex := sc.seekPreviousIndex(sc.pos, ATerm)
			if previousIndex >= 0 && sc.seekPrevious(previousIndex, _mergedUpperLower) {
				sc.accept()
				continue
			}
		}

		// SB8
		{
			// This loop is the 'regex':
			// ( ¬(OLetter | Upper | Lower | ParaSep | SATerm) )*
			pos := sc.pos
			for pos < len(sc.buffer) {
				current := sc.buffer[pos]
				if is(_mergedOLetterUpperLowerParaSepSATerm, current) {
					break
				}
				pos++
			}

			if sc.seekForward(pos-1, Lower) {
				pos := sc.pos

				sp := pos
				for {
					sp = sc.seekPreviousIndex(sp, Sp)
					if sp < 0 {
						break
					}
					pos = sp
				}

				close := pos
				for {
					close = sc.seekPreviousIndex(close, Close)
					if close < 0 {
						break
					}
					pos = close
				}

				if sc.seekPrevious(pos, ATerm) {
					sc.accept()
					continue
				}
			}
		}

		// SB8a
		if is(_mergedSContinueSATerm, current) {
			pos := sc.pos

			sp := pos
			for {
				sp = sc.seekPreviousIndex(sp, Sp)
				if sp < 0 {
					break
				}
				pos = sp
			}

			close := pos
			for {
				close = sc.seekPreviousIndex(close, Close)
				if close < 0 {
					break
				}
				pos = close
			}

			if sc.seekPrevious(pos, _mergedSATerm) {
				sc.accept()
				continue
			}
		}

		// SB9
		if is(_mergedCloseSpParaSep, current) {
			pos := sc.pos

			close := pos
			for {
				close = sc.seekPreviousIndex(close, Close)
				if close < 0 {
					break
				}
				pos = close
			}

			if sc.seekPrevious(pos, _mergedSATerm) {
				sc.accept()
				continue
			}
		}

		// SB10
		if is(_mergedSpParaSep, current) {
			pos := sc.pos

			sp := pos
			for {
				sp = sc.seekPreviousIndex(sp, Sp)
				if sp < 0 {
					break
				}
				pos = sp
			}

			close := pos
			for {
				close = sc.seekPreviousIndex(close, Close)
				if close < 0 {
					break
				}
				pos = close
			}

			if sc.seekPrevious(pos, _mergedSATerm) {
				sc.accept()
				continue
			}
		}

		// SB11
		{
			pos := sc.pos

			ps := sc.seekPreviousIndex(pos, _mergedSpParaSep)
			if ps >= 0 {
				pos = ps
			}

			sp := pos
			for {
				sp = sc.seekPreviousIndex(sp, Sp)
				if sp < 0 {
					break
				}
				pos = sp
			}

			close := pos
			for {
				close = sc.seekPreviousIndex(close, Close)
				if close < 0 {
					break
				}
				pos = close
			}

			if sc.seekPrevious(pos, _mergedSATerm) {
				break
			}
		}

		// SB998
		if sc.pos > 0 {
			sc.accept()
			continue
		}

		// If we fall through all the above rules, it's a sentence break
		break
	}

	return sc.token()
}

// Bytes returns the current token as a byte slice, after a successful call to Scan
func (sc *Scanner) Bytes() []byte {
	return sc.bytes
}

// Text returns the current token, after a successful call to Scan
func (sc *Scanner) Text() string {
	return string(sc.bytes)
}

// Err returns the current error, after an unsuccessful call to Scan
func (sc *Scanner) Err() error {
	return sc.err
}

// Sentence boundary rules: https://unicode.org/reports/tr29/#Sentence_Boundaries
// In most cases, returning true means 'keep going'; check the name of the return var for clarity

var is = unicode.Is

// seekForward looks ahead until it hits a rune satisfying one of the range tables,
// ignoring Extend|Format
// See: https://unicode.org/reports/tr29/#Grapheme_Cluster_and_Format_Rules (driven by SB5)
func (sc *Scanner) seekForward(pos int, rts ...*unicode.RangeTable) bool {
	for i := pos + 1; i < len(sc.buffer); i++ {
		r := sc.buffer[i]

		// Ignore Extend|Format
		if is(_mergedExtendFormat, r) {
			continue
		}

		// See if any of the range tables apply
		for _, rt := range rts {
			if is(rt, r) {
				return true
			}
		}

		// If we get this far, it's not there
		break
	}

	return false
}

// seekPreviousIndex works backward until it hits a rune satisfying one of the range tables,
// ignoring Extend|Format, and returns the index of the rune in the buffer
// See: https://unicode.org/reports/tr29/#Grapheme_Cluster_and_Format_Rules (driven by SB5)
func (sc *Scanner) seekPreviousIndex(pos int, rts ...*unicode.RangeTable) int {
	// Start at the end of the buffer and move backwards
	for i := pos - 1; i >= 0; i-- {
		r := sc.buffer[i]

		// Ignore Extend|Format
		if is(_mergedExtendFormat, r) {
			continue
		}

		// See if any of the range tables apply
		for _, rt := range rts {
			if is(rt, r) {
				return i
			}
		}

		// If we get this far, it's not there
		break
	}

	return -1
}

// seekPreviousIndex works backward ahead until it hits a rune satisfying one of the range tables,
// ignoring Extend|Format, reporting success
// Logic is here: https://unicode.org/reports/tr29/#Grapheme_Cluster_and_Format_Rules (driven by SB5)
func (sc *Scanner) seekPrevious(pos int, rts ...*unicode.RangeTable) bool {
	return sc.seekPreviousIndex(pos, rts...) >= 0
}

func (sc *Scanner) token() bool {
	var bb bytes.Buffer
	for _, r := range sc.buffer[:sc.pos] {
		bb.WriteRune(r)
	}
	sc.bytes = bb.Bytes()
	return len(sc.bytes) > 0
}

// accept forwards the buffer cursor (pos) by 1
func (sc *Scanner) accept() {
	sc.pos++
}

// readRune gets the next rune, advancing the reader
func (sc *Scanner) readRune() (r rune, eof bool, err error) {
	r, _, err = sc.incoming.ReadRune()

	if err != nil {
		if err == io.EOF {
			return r, true, nil
		}
		return r, false, err
	}

	return r, false, nil
}
<|endoftext|>


<commit_before>// Copyright 2013-3014 Adam Presley. All rights reserved
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package server

import (
	"fmt"
	"log"
	"net"
	"time"

	"github.com/mailslurper/libmailslurper/model/mailitem"
	"github.com/mailslurper/libmailslurper/sanitization"
	"github.com/mailslurper/libmailslurper/smtpio"
)

/*
ServerPool represents a pool of SMTP workers. This will
manage how many workers may respond to SMTP client requests
and allocation of those workers.
*/
type ServerPool chan *SmtpWorker

/*
JoinQueue adds a worker to the queue.
*/
func (pool ServerPool) JoinQueue(worker *SmtpWorker) {
	pool <- worker
}

/*
Create a new server pool with a maximum number of SMTP
workers. An array of workers is initialized with an ID
and an initial state of SMTP_WORKER_IDLE.
*/
func NewServerPool(maxWorkers int) ServerPool {
	xssService := sanitization.NewXSSService()
	emailValidationService := sanitization.NewEmailValidationService()

	pool := make(ServerPool, maxWorkers)

	for index := 0; index < maxWorkers; index++ {
		pool.JoinQueue(NewSmtpWorker(
			index+1,
			pool,
			emailValidationService,
			xssService,
		))
	}

	log.Println("libmailslurper: INFO - Worker pool configured for", maxWorkers, "worker(s)")
	return pool
}

/*
NextWorker retrieves the next available worker from
the queue.
*/
func (pool ServerPool) NextWorker(connection net.Conn, receiver chan mailitem.MailItem) (*SmtpWorker, error) {
	/*
	 * TODO: This blocks until a worker is available. Perhaps implement a timeout?
	 */
	select {
	case worker := <-pool:
		worker.Prepare(
			connection,
			receiver,
			smtpio.SmtpReader{Connection: connection},
			smtpio.SmtpWriter{Connection: connection},
		)

		log.Println("libmailslurper: INFO - Worker", worker.WorkerId, "queued to handle connection from", connection.RemoteAddr().String())
		return worker, nil

	case <-time.After(time.Second * 2):
		return &SmtpWorker{}, fmt.Errorf("No worker available. Timeout has been exceeded")
	}
}
<commit_msg>Making use of custom error<commit_after>// Copyright 2013-3014 Adam Presley. All rights reserved
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package server

import (
	"log"
	"net"
	"time"

	"github.com/mailslurper/libmailslurper/customerror"
	"github.com/mailslurper/libmailslurper/model/mailitem"
	"github.com/mailslurper/libmailslurper/sanitization"
	"github.com/mailslurper/libmailslurper/smtpio"
)

/*
ServerPool represents a pool of SMTP workers. This will
manage how many workers may respond to SMTP client requests
and allocation of those workers.
*/
type ServerPool chan *SmtpWorker

/*
JoinQueue adds a worker to the queue.
*/
func (pool ServerPool) JoinQueue(worker *SmtpWorker) {
	pool <- worker
}

/*
Create a new server pool with a maximum number of SMTP
workers. An array of workers is initialized with an ID
and an initial state of SMTP_WORKER_IDLE.
*/
func NewServerPool(maxWorkers int) ServerPool {
	xssService := sanitization.NewXSSService()
	emailValidationService := sanitization.NewEmailValidationService()

	pool := make(ServerPool, maxWorkers)

	for index := 0; index < maxWorkers; index++ {
		pool.JoinQueue(NewSmtpWorker(
			index+1,
			pool,
			emailValidationService,
			xssService,
		))
	}

	log.Println("libmailslurper: INFO - Worker pool configured for", maxWorkers, "worker(s)")
	return pool
}

/*
NextWorker retrieves the next available worker from
the queue.
*/
func (pool ServerPool) NextWorker(connection net.Conn, receiver chan mailitem.MailItem) (*SmtpWorker, error) {
	/*
	 * TODO: This blocks until a worker is available. Perhaps implement a timeout?
	 */
	select {
	case worker := <-pool:
		worker.Prepare(
			connection,
			receiver,
			smtpio.SmtpReader{Connection: connection},
			smtpio.SmtpWriter{Connection: connection},
		)

		log.Println("libmailslurper: INFO - Worker", worker.WorkerId, "queued to handle connection from", connection.RemoteAddr().String())
		return worker, nil

	case <-time.After(time.Second * 2):
		return &SmtpWorker{}, customerror.NoWorkerAvailable()
	}
}
<|endoftext|>


<commit_before>package service

import (
	"fmt"
	"reflect"
	"time"

	log "github.com/Cepave/open-falcon-backend/common/logruslog"
	nqmModel "github.com/Cepave/open-falcon-backend/common/model/nqm"
	commonQueue "github.com/Cepave/open-falcon-backend/common/queue"
	"github.com/Cepave/open-falcon-backend/common/utils"
	"github.com/Cepave/open-falcon-backend/modules/mysqlapi/model"
	"github.com/Cepave/open-falcon-backend/modules/mysqlapi/rdb"
)

var logger = log.NewDefaultLogger("INFO")

type mode byte

const (
	_DRAIN mode = 1
	_FLUSH mode = 2
)

var NqmQueue *nqmAgentUpdateService

func InitNqmHeartbeat(c *commonQueue.Config) {
	NqmQueue = newNqmAgentUpdateService(c)
	NqmQueue.updateToDatabase = updateNqmAgentHeartbeatImpl
	NqmQueue.Start()
}

func CloseNqmHeartbeat() {
	logger.Info("Closing NQM heartbeat queue service...")
	NqmQueue.Stop()
	logger.Info("Finish.")
}

var typeOfNqmAgentHeartbeat = reflect.TypeOf(new(nqmModel.HeartbeatRequest))

type nqmAgentUpdateService struct {
	q                *commonQueue.Queue
	c                *commonQueue.Config
	cnt              uint64 // counter for the dequeued elements
	running          bool
	flush            chan struct{}
	done             chan struct{}
	updateToDatabase func([]*nqmModel.HeartbeatRequest)
}

func newNqmAgentUpdateService(c *commonQueue.Config) *nqmAgentUpdateService {
	return &nqmAgentUpdateService{
		q:     commonQueue.New(),
		c:     c,
		done:  make(chan struct{}),
		flush: make(chan struct{}),
	}
}

// Gets the number of consumed updating requests(not guarantee on database)
func (q *nqmAgentUpdateService) ConsumedCount() uint64 {
	return q.cnt
}

// Gets the number of pending request of heartbeats
func (q *nqmAgentUpdateService) PendingLen() int {
	return q.q.Len()
}

func (q *nqmAgentUpdateService) Start() {
	if q.running {
		return
	}
	q.running = true
	go q.draining()
}

func (q *nqmAgentUpdateService) Stop() {
	if !q.running {
		return
	}
	q.running = false

	time.Sleep(q.c.Dur) // for all `q.q.Enqueue()`s to be done

	close(q.flush)
	<-q.done
}

func (q *nqmAgentUpdateService) Put(req *nqmModel.HeartbeatRequest) {
	if !q.running {
		return
	}
	q.q.Enqueue(req)
}

func (q *nqmAgentUpdateService) draining() {
	for {
		select {
		default:
			q.syncToDatabase(_DRAIN)
		case <-q.flush:
			q.syncToDatabase(_FLUSH)
			close(q.done)
			return
		}
	}
}

func (q *nqmAgentUpdateService) syncToDatabase(m mode) {
	var config commonQueue.Config = *q.c

	var reqs []*nqmModel.HeartbeatRequest

	switch m {
	case _FLUSH:
		config.Dur = 0
		reqs = q.drainFromQueue(&config)
		if len(reqs) > 0 {
			logger.Infof("Flushing [%d] heartbeats of NQM agent from queue", len(reqs))
		}
	default:
		reqs = q.drainFromQueue(&config)
	}

	q.updateToDatabase(reqs)
	q.cnt += uint64(len(reqs))

	if len(reqs) > 0 {
		logger.Debugf("[%d] heartbeats of NQM agent from queue", len(reqs))

		if m == _FLUSH {
			q.syncToDatabase(m)
		}
	}
}

func (q *nqmAgentUpdateService) drainFromQueue(config *commonQueue.Config) []*nqmModel.HeartbeatRequest {
	return q.q.DrainNWithDurationByType(
		config, typeOfNqmAgentHeartbeat,
	).([]*nqmModel.HeartbeatRequest)
}

func updateNqmAgentHeartbeatImpl(reqs []*nqmModel.HeartbeatRequest) {
	utils.BuildPanicCapture(
		func() {
			rdb.UpdateNqmAgentHeartbeat(reqs)
		},
		func(p interface{}) {
			logger.Errorf("[PANIC] Update heartbeats of NQM agent[#%d]: %v", len(reqs), p)
		},
	)()
}

var NqmCachedTargetList *nqmCachedTargetListService

func InitCachedTargetList(c *NqmCachedTargetListConfig) {
	NqmCachedTargetList = newNqmCachedTargetListService(c)
	logger.Infof("Target list service for agent. Timeout: %v. Queue Size: %v", c.Dur, c.Size)
	NqmCachedTargetList.Start()
}

func CloseCachedTargetList() {
	logger.Info("Closing NQM target list service...")
	NqmCachedTargetList.Stop()
	logger.Info("Finish.")
}

type NqmCachedTargetListConfig struct {
	Size int
	Dur  time.Duration
}

func newNqmCachedTargetListService(c *NqmCachedTargetListConfig) *nqmCachedTargetListService {
	return &nqmCachedTargetListService{
		agentIDQueueForRefreshCache: make(chan int32, c.Size),
		cacheTimeout:                c.Dur,
	}
}

type nqmCachedTargetListService struct {
	cacheTimeout                time.Duration
	agentIDQueueForRefreshCache chan int32
}

// Load get the current cached target list for an agent
func (s *nqmCachedTargetListService) Load(agentID int32) []*nqmModel.HeartbeatTarget {
	result, cacheLog := rdb.GetPingListFromCache(agentID, time.Now())

	go utils.BuildPanicCapture(
		func() {
			s.addRefreshCache(agentID, cacheLog)
		},
		func(p interface{}) {
			logger.Errorf("Cannot add agent id [%d] to queue: %v", agentID, p)
		},
	)()

	return result
}

// Start the service for refreshing cache of ping list
func (s *nqmCachedTargetListService) Start() {
	go func() {
		for agentID := range s.agentIDQueueForRefreshCache {
			logger.Debugf("Refresh cache for agent: [%d]", agentID)

			err := s.buildCacheOfPingList(agentID)
			if err != nil {
				logger.Errorf("Agent[%d]. Refresh has error: %v", agentID, err)
			}
		}
	}()
}

// Release resources of this service
func (s *nqmCachedTargetListService) Stop() {
	queueSize := len(s.agentIDQueueForRefreshCache)
	logger.Infof("Stopping nqmCachedTargetListService. Size of queue(refreshing cache): [%d]", queueSize)

	close(s.agentIDQueueForRefreshCache)

	/**
	 * Waiting for queue to be processed
	 */
	maxTimes := 30
	for queueSize > 0 && maxTimes > 0 {
		logger.Infof("Sleep for 2 seconds to wait queue to be processed... Current size: [%d]", queueSize)
		time.Sleep(2 * time.Second)
		maxTimes--
		queueSize = len(s.agentIDQueueForRefreshCache)
	}
	// :~)
}

func (s *nqmCachedTargetListService) addRefreshCache(agentID int32, cacheLog *model.PingListLog) {
	now := time.Now()

	/**
	 * If the timeout has reached, adds the id of agent into queue for refreshing cache
	 */
	diffDuration := now.Sub(cacheLog.RefreshTime)

	logger.Debugf(
		"Queue Size(refreshing cache of ping list): [%d]. Minutes: [%d]",
		len(s.agentIDQueueForRefreshCache),
		diffDuration/time.Minute,
	)

	if now.Sub(cacheLog.RefreshTime) >= s.cacheTimeout {
		s.agentIDQueueForRefreshCache <- agentID
	}
	// :~)
}

func (s *nqmCachedTargetListService) buildCacheOfPingList(agentID int32) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			err = fmt.Errorf("%v", p)
		}
	}()

	rdb.BuildCacheOfPingList(agentID, time.Now())

	return nil
}
<commit_msg>[OWL-1977] Add free-of-empty data for calling to database<commit_after>package service

import (
	"fmt"
	"reflect"
	"time"

	log "github.com/Cepave/open-falcon-backend/common/logruslog"
	nqmModel "github.com/Cepave/open-falcon-backend/common/model/nqm"
	commonQueue "github.com/Cepave/open-falcon-backend/common/queue"
	"github.com/Cepave/open-falcon-backend/common/utils"
	"github.com/Cepave/open-falcon-backend/modules/mysqlapi/model"
	"github.com/Cepave/open-falcon-backend/modules/mysqlapi/rdb"
)

var logger = log.NewDefaultLogger("INFO")

type mode byte

const (
	_DRAIN mode = 1
	_FLUSH mode = 2
)

var NqmQueue *nqmAgentUpdateService

func InitNqmHeartbeat(c *commonQueue.Config) {
	NqmQueue = newNqmAgentUpdateService(c)
	NqmQueue.updateToDatabase = updateNqmAgentHeartbeatImpl
	NqmQueue.Start()
}

func CloseNqmHeartbeat() {
	logger.Info("Closing NQM heartbeat queue service...")
	NqmQueue.Stop()
	logger.Info("Finish.")
}

var typeOfNqmAgentHeartbeat = reflect.TypeOf(new(nqmModel.HeartbeatRequest))

type nqmAgentUpdateService struct {
	q                *commonQueue.Queue
	c                *commonQueue.Config
	cnt              uint64 // counter for the dequeued elements
	running          bool
	flush            chan struct{}
	done             chan struct{}
	updateToDatabase func([]*nqmModel.HeartbeatRequest)
}

func newNqmAgentUpdateService(c *commonQueue.Config) *nqmAgentUpdateService {
	return &nqmAgentUpdateService{
		q:     commonQueue.New(),
		c:     c,
		done:  make(chan struct{}),
		flush: make(chan struct{}),
	}
}

// Gets the number of consumed updating requests(not guarantee on database)
func (q *nqmAgentUpdateService) ConsumedCount() uint64 {
	return q.cnt
}

// Gets the number of pending request of heartbeats
func (q *nqmAgentUpdateService) PendingLen() int {
	return q.q.Len()
}

func (q *nqmAgentUpdateService) Start() {
	if q.running {
		return
	}
	q.running = true
	go q.draining()
}

func (q *nqmAgentUpdateService) Stop() {
	if !q.running {
		return
	}
	q.running = false

	time.Sleep(q.c.Dur) // for all `q.q.Enqueue()`s to be done

	close(q.flush)
	<-q.done
}

func (q *nqmAgentUpdateService) Put(req *nqmModel.HeartbeatRequest) {
	if !q.running {
		return
	}
	q.q.Enqueue(req)
}

func (q *nqmAgentUpdateService) draining() {
	for {
		select {
		default:
			q.syncToDatabase(_DRAIN)
		case <-q.flush:
			q.syncToDatabase(_FLUSH)
			close(q.done)
			return
		}
	}
}

func (q *nqmAgentUpdateService) syncToDatabase(m mode) {
	var config commonQueue.Config = *q.c

	var reqs []*nqmModel.HeartbeatRequest

	switch m {
	case _FLUSH:
		config.Dur = 0
		reqs = q.drainFromQueue(&config)
		if len(reqs) > 0 {
			logger.Infof("Flushing [%d] heartbeats of NQM agent from queue", len(reqs))
		}
	default:
		reqs = q.drainFromQueue(&config)
	}

	q.updateToDatabase(reqs)
	q.cnt += uint64(len(reqs))

	if len(reqs) > 0 {
		logger.Debugf("[%d] heartbeats of NQM agent from queue", len(reqs))

		if m == _FLUSH {
			q.syncToDatabase(m)
		}
	}
}

func (q *nqmAgentUpdateService) drainFromQueue(config *commonQueue.Config) []*nqmModel.HeartbeatRequest {
	return q.q.DrainNWithDurationByType(
		config, typeOfNqmAgentHeartbeat,
	).([]*nqmModel.HeartbeatRequest)
}

func updateNqmAgentHeartbeatImpl(reqs []*nqmModel.HeartbeatRequest) {
	if len(reqs) == 0 {
		return
	}

	utils.BuildPanicCapture(
		func() {
			rdb.UpdateNqmAgentHeartbeat(reqs)
		},
		func(p interface{}) {
			logger.Errorf("[PANIC] Update heartbeats of NQM agent[#%d]: %v", len(reqs), p)
		},
	)()
}

var NqmCachedTargetList *nqmCachedTargetListService

func InitCachedTargetList(c *NqmCachedTargetListConfig) {
	NqmCachedTargetList = newNqmCachedTargetListService(c)
	logger.Infof("Target list service for agent. Timeout: %v. Queue Size: %v", c.Dur, c.Size)
	NqmCachedTargetList.Start()
}

func CloseCachedTargetList() {
	logger.Info("Closing NQM target list service...")
	NqmCachedTargetList.Stop()
	logger.Info("Finish.")
}

type NqmCachedTargetListConfig struct {
	Size int
	Dur  time.Duration
}

func newNqmCachedTargetListService(c *NqmCachedTargetListConfig) *nqmCachedTargetListService {
	return &nqmCachedTargetListService{
		agentIDQueueForRefreshCache: make(chan int32, c.Size),
		cacheTimeout:                c.Dur,
	}
}

type nqmCachedTargetListService struct {
	cacheTimeout                time.Duration
	agentIDQueueForRefreshCache chan int32
}

// Load get the current cached target list for an agent
func (s *nqmCachedTargetListService) Load(agentID int32) []*nqmModel.HeartbeatTarget {
	result, cacheLog := rdb.GetPingListFromCache(agentID, time.Now())

	go utils.BuildPanicCapture(
		func() {
			s.addRefreshCache(agentID, cacheLog)
		},
		func(p interface{}) {
			logger.Errorf("Cannot add agent id [%d] to queue: %v", agentID, p)
		},
	)()

	return result
}

// Start the service for refreshing cache of ping list
func (s *nqmCachedTargetListService) Start() {
	go func() {
		for agentID := range s.agentIDQueueForRefreshCache {
			logger.Debugf("Refresh cache for agent: [%d]", agentID)

			err := s.buildCacheOfPingList(agentID)
			if err != nil {
				logger.Errorf("Agent[%d]. Refresh has error: %v", agentID, err)
			}
		}
	}()
}

// Release resources of this service
func (s *nqmCachedTargetListService) Stop() {
	queueSize := len(s.agentIDQueueForRefreshCache)
	logger.Infof("Stopping nqmCachedTargetListService. Size of queue(refreshing cache): [%d]", queueSize)

	close(s.agentIDQueueForRefreshCache)

	/**
	 * Waiting for queue to be processed
	 */
	maxTimes := 30
	for queueSize > 0 && maxTimes > 0 {
		logger.Infof("Sleep for 2 seconds to wait queue to be processed... Current size: [%d]", queueSize)
		time.Sleep(2 * time.Second)
		maxTimes--
		queueSize = len(s.agentIDQueueForRefreshCache)
	}
	// :~)
}

func (s *nqmCachedTargetListService) addRefreshCache(agentID int32, cacheLog *model.PingListLog) {
	now := time.Now()

	/**
	 * If the timeout has reached, adds the id of agent into queue for refreshing cache
	 */
	diffDuration := now.Sub(cacheLog.RefreshTime)

	logger.Debugf(
		"Queue Size(refreshing cache of ping list): [%d]. Minutes: [%d]",
		len(s.agentIDQueueForRefreshCache),
		diffDuration/time.Minute,
	)

	if now.Sub(cacheLog.RefreshTime) >= s.cacheTimeout {
		s.agentIDQueueForRefreshCache <- agentID
	}
	// :~)
}

func (s *nqmCachedTargetListService) buildCacheOfPingList(agentID int32) (err error) {
	defer func() {
		p := recover()
		if p != nil {
			err = fmt.Errorf("%v", p)
		}
	}()

	rdb.BuildCacheOfPingList(agentID, time.Now())

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package server

import (
	"fmt"
	"net"
	"net/http"
	"path/filepath"
	"time"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/config"
	"github.com/cockroachdb/cockroach/gossip"
	"github.com/cockroachdb/cockroach/internal/client"
	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/rpc"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/sql/sqlbase"
	"github.com/cockroachdb/cockroach/storage"
	"github.com/cockroachdb/cockroach/storage/engine"
	"github.com/cockroachdb/cockroach/ts"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/metric"
	"github.com/cockroachdb/cockroach/util/retry"
	"github.com/cockroachdb/cockroach/util/stop"
	"github.com/pkg/errors"
)

const (
	// TestUser is a fixed user used in unittests.
	// It has valid embedded client certs.
	TestUser = "testuser"
	// initialSplitsTimeout is the amount of time to wait for initial splits to
	// occur on a freshly started server.
	// Note: this needs to be fairly high or tests become flaky.
	initialSplitsTimeout = 10 * time.Second
)

// makeTestContext returns a context for testing. It overrides the
// Certs with the test certs directory.
// We need to override the certs loader.
func makeTestContext() Context {
	ctx := MakeContext()

	// MaxOffset is the maximum offset for clocks in the cluster.
	// This is mostly irrelevant except when testing reads within
	// uncertainty intervals.
	ctx.MaxOffset = 50 * time.Millisecond

	// Test servers start in secure mode by default.
	ctx.Insecure = false

	// Load test certs. In addition, the tests requiring certs
	// need to call security.SetReadFileFn(securitytest.Asset)
	// in their init to mock out the file system calls for calls to AssetFS,
	// which has the test certs compiled in. Typically this is done
	// once per package, in main_test.go.
	ctx.SSLCA = filepath.Join(security.EmbeddedCertsDir, security.EmbeddedCACert)
	ctx.SSLCert = filepath.Join(security.EmbeddedCertsDir, security.EmbeddedNodeCert)
	ctx.SSLCertKey = filepath.Join(security.EmbeddedCertsDir, security.EmbeddedNodeKey)

	// Addr defaults to localhost with port set at time of call to
	// Start() to an available port.
	// Call TestServer.ServingAddr() for the full address (including bound port).
	ctx.Addr = "127.0.0.1:0"
	ctx.HTTPAddr = "127.0.0.1:0"
	// Set standard user for intra-cluster traffic.
	ctx.User = security.NodeUser

	return ctx
}

// makeTestContextFromParams creates a Context from a TestServerParams.
func makeTestContextFromParams(params base.TestServerArgs) Context {
	ctx := makeTestContext()
	ctx.TestingKnobs = params.Knobs
	if params.JoinAddr != "" {
		ctx.JoinUsing = params.JoinAddr
	}
	ctx.Insecure = params.Insecure
	ctx.SocketFile = params.SocketFile
	if params.MetricsSampleInterval != time.Duration(0) {
		ctx.MetricsSampleInterval = params.MetricsSampleInterval
	}
	if params.MaxOffset != time.Duration(0) {
		ctx.MaxOffset = params.MaxOffset
	}
	if params.ScanInterval != time.Duration(0) {
		ctx.ScanInterval = params.ScanInterval
	}
	if params.ScanMaxIdleTime != time.Duration(0) {
		ctx.ScanMaxIdleTime = params.ScanMaxIdleTime
	}
	if params.SSLCA != "" {
		ctx.SSLCA = params.SSLCA
	}
	if params.SSLCert != "" {
		ctx.SSLCert = params.SSLCert
	}
	if params.SSLCertKey != "" {
		ctx.SSLCertKey = params.SSLCertKey
	}
	ctx.JoinUsing = params.JoinAddr
	return ctx
}

// A TestServer encapsulates an in-memory instantiation of a cockroach node with
// a single store. It provides tests with access to Server internals.
// Where possible, it should be used through the
// testingshim.TestServerInterface.
//
// Example usage of a TestServer:
//
//   s, db, kvDB := sqlutils.SetupServer(t, testingshim.TestServerParams{})
//   defer s.Stopper().Stop()
//   // If really needed, in tests that can depend on server, downcast to
//   // server.TestServer:
//   ts := s.(*server.TestServer)
//
type TestServer struct {
	// Ctx is the context used by this server.
	Ctx *Context
	// server is the embedded Cockroach server struct.
	*Server
}

// Stopper returns the embedded server's Stopper.
func (ts *TestServer) Stopper() *stop.Stopper {
	return ts.stopper
}

// Gossip returns the gossip instance used by the TestServer.
func (ts *TestServer) Gossip() *gossip.Gossip {
	if ts != nil {
		return ts.gossip
	}
	return nil
}

// Clock returns the clock used by the TestServer.
func (ts *TestServer) Clock() *hlc.Clock {
	if ts != nil {
		return ts.clock
	}
	return nil
}

// RPCContext returns the rpc context used by the TestServer.
func (ts *TestServer) RPCContext() *rpc.Context {
	if ts != nil {
		return ts.rpcContext
	}
	return nil
}

// TsDB returns the ts.DB instance used by the TestServer.
func (ts *TestServer) TsDB() *ts.DB {
	if ts != nil {
		return ts.tsDB
	}
	return nil
}

// DB returns the client.DB instance used by the TestServer.
func (ts *TestServer) DB() *client.DB {
	if ts != nil {
		return ts.db
	}
	return nil
}

// Start starts the TestServer by bootstrapping an in-memory store
// (defaults to maximum of 100M). The server is started, launching the
// node RPC server and all HTTP endpoints. Use the value of
// TestServer.ServingAddr() after Start() for client connections.
// Use TestServer.Stopper.Stop() to shutdown the server after the test
// completes.
func (ts *TestServer) Start(params base.TestServerArgs) error {
	if ts.Ctx == nil {
		panic("Ctx not set")
	}

	// !!! I shouldn't change params.Stopper.
	if params.Stopper == nil {
		params.Stopper = stop.NewStopper()
	}

	if !params.PartOfCluster {
		// Change the replication requirements so we don't get log spam about ranges
		// not being replicated enough.
		cfg := config.DefaultZoneConfig()
		cfg.ReplicaAttrs = []roachpb.Attributes{{}}
		fn := config.TestingSetDefaultZoneConfig(cfg)
		params.Stopper.AddCloser(stop.CloserFn(fn))
	}

	// Needs to be called before NewServer to ensure resolvers are initialized.
	if err := ts.Ctx.InitNode(); err != nil {
		return err
	}

	// Ensure we have the correct number of engines. Add in-memory ones where
	// needed. There must be at least one store/engine.
	if params.StoresPerNode < 1 {
		params.StoresPerNode = 1
	}
	for i := len(ts.Ctx.Engines); i < params.StoresPerNode; i++ {
		ts.Ctx.Engines = append(ts.Ctx.Engines, engine.NewInMem(roachpb.Attributes{}, 100<<20, params.Stopper))
	}

	var err error
	ts.Server, err = NewServer(*ts.Ctx, params.Stopper)
	if err != nil {
		return err
	}
	// Our context must be shared with our server.
	ts.Ctx = &ts.Server.ctx

	if err := ts.Server.Start(); err != nil {
		return err
	}

	// If enabled, wait for initial splits to complete before returning control.
	// If initial splits do not complete, the server is stopped before
	// returning.
	if config.TestingTableSplitsDisabled() {
		return nil
	}
	if err := ts.WaitForInitialSplits(); err != nil {
		ts.Stop()
		return err
	}

	return nil
}

// ExpectedInitialRangeCount returns the expected number of ranges that should
// be on the server after initial (asynchronous) splits have been completed,
// assuming no additional information is added outside of the normal bootstrap
// process.
func ExpectedInitialRangeCount() int {
	return GetBootstrapSchema().DescriptorCount() - sqlbase.NumSystemDescriptors + 1
}

// WaitForInitialSplits waits for the server to complete its expected initial
// splits at startup. If the expected range count is not reached within a
// configured timeout, an error is returned.
func (ts *TestServer) WaitForInitialSplits() error {
	return WaitForInitialSplits(ts.DB())
}

// WaitForInitialSplits waits for the expected number of initial ranges to be
// populated in the meta2 table. If the expected range count is not reached
// within a configured timeout, an error is returned.
func WaitForInitialSplits(db *client.DB) error {
	expectedRanges := ExpectedInitialRangeCount()
	return util.RetryForDuration(initialSplitsTimeout, func() error {
		// Scan all keys in the Meta2Prefix; we only need a count.
		rows, err := db.Scan(keys.Meta2Prefix, keys.MetaMax, 0)
		if err != nil {
			return err
		}
		if a, e := len(rows), expectedRanges; a != e {
			return errors.Errorf("had %d ranges at startup, expected %d", a, e)
		}
		return nil
	})
}

// Stores returns the collection of stores from this TestServer's node.
func (ts *TestServer) Stores() *storage.Stores {
	return ts.node.stores
}

// ServingAddr returns the server's address. Should be used by clients.
func (ts *TestServer) ServingAddr() string {
	return ts.ctx.Addr
}

// ServingHost returns the host portion of the rpc server's address.
func (ts *TestServer) ServingHost() (string, error) {
	h, _, err := net.SplitHostPort(ts.ServingAddr())
	return h, err
}

// ServingPort returns the port portion of the rpc server's address.
func (ts *TestServer) ServingPort() (string, error) {
	_, p, err := net.SplitHostPort(ts.ServingAddr())
	return p, err
}

// SetRangeRetryOptions sets the retry options for stores in TestServer.
func (ts *TestServer) SetRangeRetryOptions(ro retry.Options) {
	if err := ts.node.stores.VisitStores(func(s *storage.Store) error {
		s.SetRangeRetryOptions(ro)
		return nil
	}); err != nil {
		panic(err)
	}
}

// WriteSummaries implements TestServerInterface.
func (ts *TestServer) WriteSummaries() error {
	return ts.node.writeSummaries()
}

// AdminURL implements TestServerInterface.
func (ts *TestServer) AdminURL() string {
	return ts.Ctx.AdminURL()
}

// GetHTTPClient implements TestServerInterface.
func (ts *TestServer) GetHTTPClient() (http.Client, error) {
	return ts.Ctx.GetHTTPClient()
}

// MustGetSQLCounter implements TestServerInterface.
func (ts *TestServer) MustGetSQLCounter(name string) int64 {
	var c int64
	var found bool

	ts.sqlExecutor.Registry().Each(func(n string, v interface{}) {
		if name == n {
			c = v.(*metric.Counter).Count()
			found = true
		}
	})
	if !found {
		panic(fmt.Sprintf("couldn't find metric %s", name))
	}
	return c
}

// MustGetSQLNetworkCounter implements TestServerInterface.
func (ts *TestServer) MustGetSQLNetworkCounter(name string) int64 {
	var c int64
	var found bool

	ts.pgServer.Registry().Each(func(n string, v interface{}) {
		if name == n {
			c = v.(*metric.Counter).Count()
			found = true
		}
	})
	if !found {
		panic(fmt.Sprintf("couldn't find metric %s", name))
	}
	return c
}

// KVClient is part of TestServerInterface.
func (ts *TestServer) KVClient() interface{} { return ts.db }

// KVDB is part of TestServerInterface.
func (ts *TestServer) KVDB() interface{} { return ts.kvDB }

// LeaseManager is part of TestServerInterface.
func (ts *TestServer) LeaseManager() interface{} {
	return ts.leaseMgr
}

type testServerFactoryImpl struct{}

// TestServerFactory can be passed to testingshim.InitTestServerFactory
var TestServerFactory = testServerFactoryImpl{}

// New is part of TestServerFactory interface.
func (testServerFactoryImpl) New(
	params base.TestServerArgs,
) interface{} {
	ctx := makeTestContextFromParams(params)
	return &TestServer{Ctx: &ctx}
}
<commit_msg>testserver: cleanup leftover comment<commit_after>// Copyright 2014 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package server

import (
	"fmt"
	"net"
	"net/http"
	"path/filepath"
	"time"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/config"
	"github.com/cockroachdb/cockroach/gossip"
	"github.com/cockroachdb/cockroach/internal/client"
	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/rpc"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/sql/sqlbase"
	"github.com/cockroachdb/cockroach/storage"
	"github.com/cockroachdb/cockroach/storage/engine"
	"github.com/cockroachdb/cockroach/ts"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/metric"
	"github.com/cockroachdb/cockroach/util/retry"
	"github.com/cockroachdb/cockroach/util/stop"
	"github.com/pkg/errors"
)

const (
	// TestUser is a fixed user used in unittests.
	// It has valid embedded client certs.
	TestUser = "testuser"
	// initialSplitsTimeout is the amount of time to wait for initial splits to
	// occur on a freshly started server.
	// Note: this needs to be fairly high or tests become flaky.
	initialSplitsTimeout = 10 * time.Second
)

// makeTestContext returns a context for testing. It overrides the
// Certs with the test certs directory.
// We need to override the certs loader.
func makeTestContext() Context {
	ctx := MakeContext()

	// MaxOffset is the maximum offset for clocks in the cluster.
	// This is mostly irrelevant except when testing reads within
	// uncertainty intervals.
	ctx.MaxOffset = 50 * time.Millisecond

	// Test servers start in secure mode by default.
	ctx.Insecure = false

	// Load test certs. In addition, the tests requiring certs
	// need to call security.SetReadFileFn(securitytest.Asset)
	// in their init to mock out the file system calls for calls to AssetFS,
	// which has the test certs compiled in. Typically this is done
	// once per package, in main_test.go.
	ctx.SSLCA = filepath.Join(security.EmbeddedCertsDir, security.EmbeddedCACert)
	ctx.SSLCert = filepath.Join(security.EmbeddedCertsDir, security.EmbeddedNodeCert)
	ctx.SSLCertKey = filepath.Join(security.EmbeddedCertsDir, security.EmbeddedNodeKey)

	// Addr defaults to localhost with port set at time of call to
	// Start() to an available port.
	// Call TestServer.ServingAddr() for the full address (including bound port).
	ctx.Addr = "127.0.0.1:0"
	ctx.HTTPAddr = "127.0.0.1:0"
	// Set standard user for intra-cluster traffic.
	ctx.User = security.NodeUser

	return ctx
}

// makeTestContextFromParams creates a Context from a TestServerParams.
func makeTestContextFromParams(params base.TestServerArgs) Context {
	ctx := makeTestContext()
	ctx.TestingKnobs = params.Knobs
	if params.JoinAddr != "" {
		ctx.JoinUsing = params.JoinAddr
	}
	ctx.Insecure = params.Insecure
	ctx.SocketFile = params.SocketFile
	if params.MetricsSampleInterval != time.Duration(0) {
		ctx.MetricsSampleInterval = params.MetricsSampleInterval
	}
	if params.MaxOffset != time.Duration(0) {
		ctx.MaxOffset = params.MaxOffset
	}
	if params.ScanInterval != time.Duration(0) {
		ctx.ScanInterval = params.ScanInterval
	}
	if params.ScanMaxIdleTime != time.Duration(0) {
		ctx.ScanMaxIdleTime = params.ScanMaxIdleTime
	}
	if params.SSLCA != "" {
		ctx.SSLCA = params.SSLCA
	}
	if params.SSLCert != "" {
		ctx.SSLCert = params.SSLCert
	}
	if params.SSLCertKey != "" {
		ctx.SSLCertKey = params.SSLCertKey
	}
	ctx.JoinUsing = params.JoinAddr
	return ctx
}

// A TestServer encapsulates an in-memory instantiation of a cockroach node with
// a single store. It provides tests with access to Server internals.
// Where possible, it should be used through the
// testingshim.TestServerInterface.
//
// Example usage of a TestServer:
//
//   s, db, kvDB := sqlutils.SetupServer(t, testingshim.TestServerParams{})
//   defer s.Stopper().Stop()
//   // If really needed, in tests that can depend on server, downcast to
//   // server.TestServer:
//   ts := s.(*server.TestServer)
//
type TestServer struct {
	// Ctx is the context used by this server.
	Ctx *Context
	// server is the embedded Cockroach server struct.
	*Server
}

// Stopper returns the embedded server's Stopper.
func (ts *TestServer) Stopper() *stop.Stopper {
	return ts.stopper
}

// Gossip returns the gossip instance used by the TestServer.
func (ts *TestServer) Gossip() *gossip.Gossip {
	if ts != nil {
		return ts.gossip
	}
	return nil
}

// Clock returns the clock used by the TestServer.
func (ts *TestServer) Clock() *hlc.Clock {
	if ts != nil {
		return ts.clock
	}
	return nil
}

// RPCContext returns the rpc context used by the TestServer.
func (ts *TestServer) RPCContext() *rpc.Context {
	if ts != nil {
		return ts.rpcContext
	}
	return nil
}

// TsDB returns the ts.DB instance used by the TestServer.
func (ts *TestServer) TsDB() *ts.DB {
	if ts != nil {
		return ts.tsDB
	}
	return nil
}

// DB returns the client.DB instance used by the TestServer.
func (ts *TestServer) DB() *client.DB {
	if ts != nil {
		return ts.db
	}
	return nil
}

// Start starts the TestServer by bootstrapping an in-memory store
// (defaults to maximum of 100M). The server is started, launching the
// node RPC server and all HTTP endpoints. Use the value of
// TestServer.ServingAddr() after Start() for client connections.
// Use TestServer.Stopper().Stop() to shutdown the server after the test
// completes.
func (ts *TestServer) Start(params base.TestServerArgs) error {
	if ts.Ctx == nil {
		panic("Ctx not set")
	}

	if params.Stopper == nil {
		params.Stopper = stop.NewStopper()
	}

	if !params.PartOfCluster {
		// Change the replication requirements so we don't get log spam about ranges
		// not being replicated enough.
		cfg := config.DefaultZoneConfig()
		cfg.ReplicaAttrs = []roachpb.Attributes{{}}
		fn := config.TestingSetDefaultZoneConfig(cfg)
		params.Stopper.AddCloser(stop.CloserFn(fn))
	}

	// Needs to be called before NewServer to ensure resolvers are initialized.
	if err := ts.Ctx.InitNode(); err != nil {
		return err
	}

	// Ensure we have the correct number of engines. Add in-memory ones where
	// needed. There must be at least one store/engine.
	if params.StoresPerNode < 1 {
		params.StoresPerNode = 1
	}
	for i := len(ts.Ctx.Engines); i < params.StoresPerNode; i++ {
		ts.Ctx.Engines = append(ts.Ctx.Engines, engine.NewInMem(roachpb.Attributes{}, 100<<20, params.Stopper))
	}

	var err error
	ts.Server, err = NewServer(*ts.Ctx, params.Stopper)
	if err != nil {
		return err
	}
	// Our context must be shared with our server.
	ts.Ctx = &ts.Server.ctx

	if err := ts.Server.Start(); err != nil {
		return err
	}

	// If enabled, wait for initial splits to complete before returning control.
	// If initial splits do not complete, the server is stopped before
	// returning.
	if config.TestingTableSplitsDisabled() {
		return nil
	}
	if err := ts.WaitForInitialSplits(); err != nil {
		ts.Stop()
		return err
	}

	return nil
}

// ExpectedInitialRangeCount returns the expected number of ranges that should
// be on the server after initial (asynchronous) splits have been completed,
// assuming no additional information is added outside of the normal bootstrap
// process.
func ExpectedInitialRangeCount() int {
	return GetBootstrapSchema().DescriptorCount() - sqlbase.NumSystemDescriptors + 1
}

// WaitForInitialSplits waits for the server to complete its expected initial
// splits at startup. If the expected range count is not reached within a
// configured timeout, an error is returned.
func (ts *TestServer) WaitForInitialSplits() error {
	return WaitForInitialSplits(ts.DB())
}

// WaitForInitialSplits waits for the expected number of initial ranges to be
// populated in the meta2 table. If the expected range count is not reached
// within a configured timeout, an error is returned.
func WaitForInitialSplits(db *client.DB) error {
	expectedRanges := ExpectedInitialRangeCount()
	return util.RetryForDuration(initialSplitsTimeout, func() error {
		// Scan all keys in the Meta2Prefix; we only need a count.
		rows, err := db.Scan(keys.Meta2Prefix, keys.MetaMax, 0)
		if err != nil {
			return err
		}
		if a, e := len(rows), expectedRanges; a != e {
			return errors.Errorf("had %d ranges at startup, expected %d", a, e)
		}
		return nil
	})
}

// Stores returns the collection of stores from this TestServer's node.
func (ts *TestServer) Stores() *storage.Stores {
	return ts.node.stores
}

// ServingAddr returns the server's address. Should be used by clients.
func (ts *TestServer) ServingAddr() string {
	return ts.ctx.Addr
}

// ServingHost returns the host portion of the rpc server's address.
func (ts *TestServer) ServingHost() (string, error) {
	h, _, err := net.SplitHostPort(ts.ServingAddr())
	return h, err
}

// ServingPort returns the port portion of the rpc server's address.
func (ts *TestServer) ServingPort() (string, error) {
	_, p, err := net.SplitHostPort(ts.ServingAddr())
	return p, err
}

// SetRangeRetryOptions sets the retry options for stores in TestServer.
func (ts *TestServer) SetRangeRetryOptions(ro retry.Options) {
	if err := ts.node.stores.VisitStores(func(s *storage.Store) error {
		s.SetRangeRetryOptions(ro)
		return nil
	}); err != nil {
		panic(err)
	}
}

// WriteSummaries implements TestServerInterface.
func (ts *TestServer) WriteSummaries() error {
	return ts.node.writeSummaries()
}

// AdminURL implements TestServerInterface.
func (ts *TestServer) AdminURL() string {
	return ts.Ctx.AdminURL()
}

// GetHTTPClient implements TestServerInterface.
func (ts *TestServer) GetHTTPClient() (http.Client, error) {
	return ts.Ctx.GetHTTPClient()
}

// MustGetSQLCounter implements TestServerInterface.
func (ts *TestServer) MustGetSQLCounter(name string) int64 {
	var c int64
	var found bool

	ts.sqlExecutor.Registry().Each(func(n string, v interface{}) {
		if name == n {
			c = v.(*metric.Counter).Count()
			found = true
		}
	})
	if !found {
		panic(fmt.Sprintf("couldn't find metric %s", name))
	}
	return c
}

// MustGetSQLNetworkCounter implements TestServerInterface.
func (ts *TestServer) MustGetSQLNetworkCounter(name string) int64 {
	var c int64
	var found bool

	ts.pgServer.Registry().Each(func(n string, v interface{}) {
		if name == n {
			c = v.(*metric.Counter).Count()
			found = true
		}
	})
	if !found {
		panic(fmt.Sprintf("couldn't find metric %s", name))
	}
	return c
}

// KVClient is part of TestServerInterface.
func (ts *TestServer) KVClient() interface{} { return ts.db }

// KVDB is part of TestServerInterface.
func (ts *TestServer) KVDB() interface{} { return ts.kvDB }

// LeaseManager is part of TestServerInterface.
func (ts *TestServer) LeaseManager() interface{} {
	return ts.leaseMgr
}

type testServerFactoryImpl struct{}

// TestServerFactory can be passed to testingshim.InitTestServerFactory
var TestServerFactory = testServerFactoryImpl{}

// New is part of TestServerFactory interface.
func (testServerFactoryImpl) New(
	params base.TestServerArgs,
) interface{} {
	ctx := makeTestContextFromParams(params)
	return &TestServer{Ctx: &ctx}
}
<|endoftext|>


<commit_before>package version

// VERSION ...
const VERSION = "1.1.19"
<commit_msg>v1.1.20<commit_after>package version

// VERSION ...
const VERSION = "1.1.20"
<|endoftext|>


<commit_before>// Copyright 2016 The Linux Foundation
//
// 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 version

import "fmt"

const (
	// VersionMajor is for an API incompatible changes
	VersionMajor = 0
	// VersionMinor is for functionality in a backwards-compatible manner
	VersionMinor = 3
	// VersionPatch is for backwards-compatible bug fixes
	VersionPatch = 0

	// VersionDev indicates development branch. Releases will be empty string.
	VersionDev = ""
)

// Version is the specification version that the package types support.
var Version = fmt.Sprintf("%d.%d.%d%s", VersionMajor, VersionMinor, VersionPatch, VersionDev)
<commit_msg>version: master back to -dev<commit_after>// Copyright 2016 The Linux Foundation
//
// 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 version

import "fmt"

const (
	// VersionMajor is for an API incompatible changes
	VersionMajor = 0
	// VersionMinor is for functionality in a backwards-compatible manner
	VersionMinor = 3
	// VersionPatch is for backwards-compatible bug fixes
	VersionPatch = 0

	// VersionDev indicates development branch. Releases will be empty string.
	VersionDev = "-dev"
)

// Version is the specification version that the package types support.
var Version = fmt.Sprintf("%d.%d.%d%s", VersionMajor, VersionMinor, VersionPatch, VersionDev)
<|endoftext|>


<commit_before>package version

// Version represents the hub version number
var Version = "2.14.0"
<commit_msg>hub 2.14.1<commit_after>package version

// Version represents the hub version number
var Version = "2.14.1"
<|endoftext|>


<commit_before>// The version package provides a location to set the release versions for all
// packages to consume, without creating import cycles.
//
// This package should not import any other terraform packages.
package version

import (
	"fmt"

	version "github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
var Version = "0.12.0"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
var Prerelease = "dev"

// SemVer is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVer *version.Version

func init() {
	SemVer = version.Must(version.NewVersion(Version))
}

// Header is the header name used to send the current terraform version
// in http requests.
const Header = "Terraform-Version"

// String returns the complete version string, including prerelease
func String() string {
	if Prerelease != "" {
		return fmt.Sprintf("%s-%s", Version, Prerelease)
	}
	return Version
}
<commit_msg>v0.12.0<commit_after>// The version package provides a location to set the release versions for all
// packages to consume, without creating import cycles.
//
// This package should not import any other terraform packages.
package version

import (
	"fmt"

	version "github.com/hashicorp/go-version"
)

// The main version number that is being run at the moment.
var Version = "0.12.0"

// A pre-release marker for the version. If this is "" (empty string)
// then it means that it is a final release. Otherwise, this is a pre-release
// such as "dev" (in development), "beta", "rc1", etc.
var Prerelease = ""

// SemVer is an instance of version.Version. This has the secondary
// benefit of verifying during tests and init time that our version is a
// proper semantic version, which should always be the case.
var SemVer *version.Version

func init() {
	SemVer = version.Must(version.NewVersion(Version))
}

// Header is the header name used to send the current terraform version
// in http requests.
const Header = "Terraform-Version"

// String returns the complete version string, including prerelease
func String() string {
	if Prerelease != "" {
		return fmt.Sprintf("%s-%s", Version, Prerelease)
	}
	return Version
}
<|endoftext|>


<commit_before>package version

import (
	"fmt"

	"github.com/github/hub/git"
)

var Version = "2.3.0-pre10"

func FullVersion() (string, error) {
	gitVersion, err := git.Version()
	if err != nil {
		gitVersion = "git version (unavailable)"
	}
	return fmt.Sprintf("%s\nhub version %s", gitVersion, Version), err
}
<commit_msg>hub 2.3.0<commit_after>package version

import (
	"fmt"

	"github.com/github/hub/git"
)

var Version = "2.3.0"

func FullVersion() (string, error) {
	gitVersion, err := git.Version()
	if err != nil {
		gitVersion = "git version (unavailable)"
	}
	return fmt.Sprintf("%s\nhub version %s", gitVersion, Version), err
}
<|endoftext|>


<commit_before>package version

const Version = "1.3.2"
<commit_msg>chore(version): 1.3.3<commit_after>package version

const Version = "1.3.3"
<|endoftext|>


<commit_before>package services

// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import (
	"context"
	"fmt"
	"regexp"
	"strings"

	"cloud.google.com/go/iam"
	"github.com/pkg/errors"
	crm "google.golang.org/api/cloudresourcemanager/v1"
)

type crmClient interface {
	GetAncestry(context.Context, string) (*crm.GetAncestryResponse, error)
	SetPolicyProject(context.Context, string, *crm.Policy) (*crm.Policy, error)
	GetPolicyProject(context.Context, string) (*crm.Policy, error)
	GetPolicyOrganization(context.Context, string) (*crm.Policy, error)
	SetPolicyOrganization(context.Context, string, *crm.Policy) (*crm.Policy, error)
	GetOrganization(context.Context, string) (*crm.Organization, error)
	SetPolicyProjectWithMask(context.Context, string, *crm.Policy, ...string) (*crm.Policy, error)
}

type storageClient interface {
	SetBucketPolicy(context.Context, string, *iam.Policy) error
	BucketPolicy(context.Context, string) (*iam.Policy, error)
	EnableBucketOnlyPolicy(context.Context, string) error
}

// Resource service.
type Resource struct {
	crm     crmClient
	storage storageClient
}

// NewResource returns a new resource service.
func NewResource(crm crmClient, s storageClient) *Resource {
	return &Resource{
		crm:     crm,
		storage: s,
	}
}

// RemoveUsersProject removes users from the project.
func (r *Resource) RemoveUsersProject(ctx context.Context, projectID string, remove []string) error {
	existingPolicy, err := r.crm.GetPolicyProject(ctx, projectID)
	if err != nil {
		return fmt.Errorf("failed to get project policy: %q", err)
	}
	policy := r.removeUsersFromPolicy(existingPolicy, remove)
	if _, err := r.crm.SetPolicyProject(ctx, projectID, policy); err != nil {
		return fmt.Errorf("failed to set project policy: %q", err)
	}
	return nil
}

// RemoveMembersFromBucket removes members from the bucket.
func (r *Resource) RemoveMembersFromBucket(ctx context.Context, bucketName string, members []string) error {
	p, err := r.storage.BucketPolicy(ctx, bucketName)
	if err != nil {
		return err
	}
	// Save what we need to remove in a map so we don't mutate a slice while we iterate over it.
	toRemove := make(map[iam.RoleName]map[string]bool)

	for _, role := range p.Roles() {
		for _, policyMember := range p.Members(role) {
			for _, m := range members {
				if policyMember != m {
					continue
				}
				if toRemove[role] == nil {
					toRemove[role] = make(map[string]bool)
				}
				toRemove[role][m] = true
			}
		}
	}

	for k, v := range toRemove {
		for kk := range v {
			p.Remove(kk, k)
		}
	}
	return r.storage.SetBucketPolicy(ctx, bucketName, p)
}

// EnableAuditLogs enable audit logs to all services and LogTypes.
func (r *Resource) EnableAuditLogs(ctx context.Context, projectID string) (*crm.Policy, error) {
	res, err := r.crm.GetPolicyProject(ctx, projectID)
	if err != nil {
		return nil, errors.Wrap(err, "failed to get project policy")
	}
	isDefault := false
	enableAll := &crm.AuditConfig{
		AuditLogConfigs: []*crm.AuditLogConfig{
			{LogType: "ADMIN_READ"},
			{LogType: "DATA_READ"},
			{LogType: "DATA_WRITE"},
		},
		Service: "allServices",
	}
	for _, conf := range res.AuditConfigs {
		if conf.Service == "allServices" {
			conf.AuditLogConfigs = enableAll.AuditLogConfigs
			isDefault = true
		}
	}
	if !isDefault {
		res.AuditConfigs = append(res.AuditConfigs, enableAll)
	}

	result, err := r.crm.SetPolicyProjectWithMask(ctx, projectID, res, "auditConfigs")
	if err != nil {
		return nil, errors.Wrap(err, "failed to update project policy")
	}
	return result, nil
}

// GetProjectAncestry returns a slice of the project's ancestry.
func (r *Resource) GetProjectAncestry(ctx context.Context, projectID string) ([]string, error) {
	resp, err := r.crm.GetAncestry(ctx, projectID)
	if err != nil {
		return nil, err
	}
	s := []string{}
	for _, a := range resp.Ancestor {
		s = append(s, a.ResourceId.Type+"s/"+a.ResourceId.Id)
	}
	return s, nil
}

// removeUsersFromPolicy removes a slice of users from a policy
func (r *Resource) removeUsersFromPolicy(policy *crm.Policy, users []string) *crm.Policy {
	for _, b := range policy.Bindings {
		members := []string{}
		for _, member := range b.Members {
			isUser := strings.HasPrefix(member, "user:")
			found := false
			for _, user := range users {
				if user == member {
					found = true
					break
				}
			}
			if !isUser || !found {
				members = append(members, member)
				continue
			}
		}
		b.Members = members
	}
	return policy

// removeMembersFromOrgPolicy removes Google account (user:) members that doesn't match the given regex.
func (r *Resource) removeMembersFromOrgPolicy(regex *regexp.Regexp, policy *crm.Policy) (*crm.Policy, []string) {
	membersToRemove := []string{}
	for _, b := range policy.Bindings {
		allowedMembers := []string{}
		for _, m := range b.Members {
			isUser := strings.HasPrefix(m, "user:")
			if !isUser || regex.MatchString(m) {
				allowedMembers = append(allowedMembers, m)
			} else {
				membersToRemove = append(membersToRemove, m)
			}
		}
		b.Members = allowedMembers
	}
	return policy, membersToRemove
}

// removeMembersFromPolicy removes members that match the given regex.
func (r *Resource) removeMembersFromPolicy(regex *regexp.Regexp, policy *crm.Policy) *crm.Policy {
	for _, b := range policy.Bindings {
		members := []string{}
		for _, m := range b.Members {
			if !regex.MatchString(m) {
				members = append(members, m)
			}
		}
		b.Members = members
	}
	return policy
}

// RemoveMembersOrganization removes the given members from the organization.
func (r *Resource) RemoveMembersOrganization(ctx context.Context, displayName, name string, allowed []string, p *crm.Policy) ([]string, error) {
	allowed = append(allowed, displayName)
	j := strings.Replace(strings.Join(allowed, "|"), ".", `\.`, -1)
	e, err := regexp.Compile("^.+@" + j + "$")
	if err != nil {
		return nil, fmt.Errorf("failed to compile regex: %q", err)
	}
	newPolicy, membersToRemove := r.removeMembersFromOrgPolicy(e, p)
	if _, err := r.crm.SetPolicyOrganization(ctx, name, newPolicy); err != nil {
		return membersToRemove, fmt.Errorf("failed to set project policy: %q", err)
	}
	return membersToRemove, nil
}

// PolicyOrganization returns the IAM policy for the given resource name.
func (r *Resource) PolicyOrganization(ctx context.Context, name string) (*crm.Policy, error) {
	return r.crm.GetPolicyOrganization(ctx, name)
}

// Organization returns the organization name for the given organization resource.
func (r *Resource) Organization(ctx context.Context, orgID string) (*crm.Organization, error) {
	return r.crm.GetOrganization(ctx, "organizations/"+orgID)
}

// EnableBucketOnlyPolicy enable bucket only policy for the given bucket
func (r *Resource) EnableBucketOnlyPolicy(ctx context.Context, bucketName string) error {
	return r.storage.EnableBucketOnlyPolicy(ctx, bucketName)
}

// IfProjectWithinResources executes the provided function if the project ID is an ancestor of any provided resources.
func (r *Resource) IfProjectWithinResources(ctx context.Context, conf *Resources, projectID string, fn func() error) error {
	if err := r.IfProjectInFolders(ctx, conf.FolderIDs, projectID, fn); err != nil {
		return err
	}
	if err := r.IfProjectInProjects(ctx, conf.ProjectIDs, projectID, fn); err != nil {
		return err
	}
	if err := r.IfProjectInOrg(ctx, conf.OrganizationID, projectID, fn); err != nil {
		return err
	}
	return nil
}

// IfProjectInFolders will apply the function if the project ID is within the folder IDs.
func (r *Resource) IfProjectInFolders(ctx context.Context, ids []string, projectID string, fn func() error) error {
	if len(ids) == 0 {
		return nil
	}
	ancestors, err := r.GetProjectAncestry(ctx, projectID)
	if err != nil {
		return errors.Wrap(err, "failed to get project ancestry")
	}
	for _, resource := range ancestors {
		for _, folderID := range ids {
			if resource != "folders/"+folderID {
				continue
			}
			if err := fn(); err != nil {
				return err
			}
		}
	}
	return nil
}

// IfProjectInProjects will apply the function if the project ID is within the project IDs.
func (r *Resource) IfProjectInProjects(ctx context.Context, ids []string, projectID string, fn func() error) error {
	if len(ids) == 0 {
		return nil
	}
	for _, v := range ids {
		if v != projectID {
			continue
		}
		if err := fn(); err != nil {
			return err
		}
	}
	return nil
}

// IfProjectInOrg will apply the function if the project ID is within the organization.
func (r *Resource) IfProjectInOrg(ctx context.Context, orgID, projectID string, fn func() error) error {
	if orgID == "" {
		return nil
	}
	ancestors, err := r.GetProjectAncestry(ctx, projectID)
	if err != nil {
		return errors.Wrap(err, "failed to get project ancestry")
	}
	for _, resource := range ancestors {
		if resource == "organizations/"+orgID {
			if err := fn(); err != nil {
				return err
			}
		}
	}
	return nil
}
<commit_msg>restore missing curly brace<commit_after>package services

// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import (
	"context"
	"fmt"
	"regexp"
	"strings"

	"cloud.google.com/go/iam"
	"github.com/pkg/errors"
	crm "google.golang.org/api/cloudresourcemanager/v1"
)

type crmClient interface {
	GetAncestry(context.Context, string) (*crm.GetAncestryResponse, error)
	SetPolicyProject(context.Context, string, *crm.Policy) (*crm.Policy, error)
	GetPolicyProject(context.Context, string) (*crm.Policy, error)
	GetPolicyOrganization(context.Context, string) (*crm.Policy, error)
	SetPolicyOrganization(context.Context, string, *crm.Policy) (*crm.Policy, error)
	GetOrganization(context.Context, string) (*crm.Organization, error)
	SetPolicyProjectWithMask(context.Context, string, *crm.Policy, ...string) (*crm.Policy, error)
}

type storageClient interface {
	SetBucketPolicy(context.Context, string, *iam.Policy) error
	BucketPolicy(context.Context, string) (*iam.Policy, error)
	EnableBucketOnlyPolicy(context.Context, string) error
}

// Resource service.
type Resource struct {
	crm     crmClient
	storage storageClient
}

// NewResource returns a new resource service.
func NewResource(crm crmClient, s storageClient) *Resource {
	return &Resource{
		crm:     crm,
		storage: s,
	}
}

// RemoveUsersProject removes users from the project.
func (r *Resource) RemoveUsersProject(ctx context.Context, projectID string, remove []string) error {
	existingPolicy, err := r.crm.GetPolicyProject(ctx, projectID)
	if err != nil {
		return fmt.Errorf("failed to get project policy: %q", err)
	}
	policy := r.removeUsersFromPolicy(existingPolicy, remove)
	if _, err := r.crm.SetPolicyProject(ctx, projectID, policy); err != nil {
		return fmt.Errorf("failed to set project policy: %q", err)
	}
	return nil
}

// RemoveMembersFromBucket removes members from the bucket.
func (r *Resource) RemoveMembersFromBucket(ctx context.Context, bucketName string, members []string) error {
	p, err := r.storage.BucketPolicy(ctx, bucketName)
	if err != nil {
		return err
	}
	// Save what we need to remove in a map so we don't mutate a slice while we iterate over it.
	toRemove := make(map[iam.RoleName]map[string]bool)

	for _, role := range p.Roles() {
		for _, policyMember := range p.Members(role) {
			for _, m := range members {
				if policyMember != m {
					continue
				}
				if toRemove[role] == nil {
					toRemove[role] = make(map[string]bool)
				}
				toRemove[role][m] = true
			}
		}
	}

	for k, v := range toRemove {
		for kk := range v {
			p.Remove(kk, k)
		}
	}
	return r.storage.SetBucketPolicy(ctx, bucketName, p)
}

// EnableAuditLogs enable audit logs to all services and LogTypes.
func (r *Resource) EnableAuditLogs(ctx context.Context, projectID string) (*crm.Policy, error) {
	res, err := r.crm.GetPolicyProject(ctx, projectID)
	if err != nil {
		return nil, errors.Wrap(err, "failed to get project policy")
	}
	isDefault := false
	enableAll := &crm.AuditConfig{
		AuditLogConfigs: []*crm.AuditLogConfig{
			{LogType: "ADMIN_READ"},
			{LogType: "DATA_READ"},
			{LogType: "DATA_WRITE"},
		},
		Service: "allServices",
	}
	for _, conf := range res.AuditConfigs {
		if conf.Service == "allServices" {
			conf.AuditLogConfigs = enableAll.AuditLogConfigs
			isDefault = true
		}
	}
	if !isDefault {
		res.AuditConfigs = append(res.AuditConfigs, enableAll)
	}

	result, err := r.crm.SetPolicyProjectWithMask(ctx, projectID, res, "auditConfigs")
	if err != nil {
		return nil, errors.Wrap(err, "failed to update project policy")
	}
	return result, nil
}

// GetProjectAncestry returns a slice of the project's ancestry.
func (r *Resource) GetProjectAncestry(ctx context.Context, projectID string) ([]string, error) {
	resp, err := r.crm.GetAncestry(ctx, projectID)
	if err != nil {
		return nil, err
	}
	s := []string{}
	for _, a := range resp.Ancestor {
		s = append(s, a.ResourceId.Type+"s/"+a.ResourceId.Id)
	}
	return s, nil
}

// removeUsersFromPolicy removes a slice of users from a policy
func (r *Resource) removeUsersFromPolicy(policy *crm.Policy, users []string) *crm.Policy {
	for _, b := range policy.Bindings {
		members := []string{}
		for _, member := range b.Members {
			isUser := strings.HasPrefix(member, "user:")
			found := false
			for _, user := range users {
				if user == member {
					found = true
					break
				}
			}
			if !isUser || !found {
				members = append(members, member)
				continue
			}
		}
		b.Members = members
	}
	return policy
}

// removeMembersFromOrgPolicy removes Google account (user:) members that doesn't match the given regex.
func (r *Resource) removeMembersFromOrgPolicy(regex *regexp.Regexp, policy *crm.Policy) (*crm.Policy, []string) {
	membersToRemove := []string{}
	for _, b := range policy.Bindings {
		allowedMembers := []string{}
		for _, m := range b.Members {
			isUser := strings.HasPrefix(m, "user:")
			if !isUser || regex.MatchString(m) {
				allowedMembers = append(allowedMembers, m)
			} else {
				membersToRemove = append(membersToRemove, m)
			}
		}
		b.Members = allowedMembers
	}
	return policy, membersToRemove
}

// removeMembersFromPolicy removes members that match the given regex.
func (r *Resource) removeMembersFromPolicy(regex *regexp.Regexp, policy *crm.Policy) *crm.Policy {
	for _, b := range policy.Bindings {
		members := []string{}
		for _, m := range b.Members {
			if !regex.MatchString(m) {
				members = append(members, m)
			}
		}
		b.Members = members
	}
	return policy
}

// RemoveMembersOrganization removes the given members from the organization.
func (r *Resource) RemoveMembersOrganization(ctx context.Context, displayName, name string, allowed []string, p *crm.Policy) ([]string, error) {
	allowed = append(allowed, displayName)
	j := strings.Replace(strings.Join(allowed, "|"), ".", `\.`, -1)
	e, err := regexp.Compile("^.+@" + j + "$")
	if err != nil {
		return nil, fmt.Errorf("failed to compile regex: %q", err)
	}
	newPolicy, membersToRemove := r.removeMembersFromOrgPolicy(e, p)
	if _, err := r.crm.SetPolicyOrganization(ctx, name, newPolicy); err != nil {
		return membersToRemove, fmt.Errorf("failed to set project policy: %q", err)
	}
	return membersToRemove, nil
}

// PolicyOrganization returns the IAM policy for the given resource name.
func (r *Resource) PolicyOrganization(ctx context.Context, name string) (*crm.Policy, error) {
	return r.crm.GetPolicyOrganization(ctx, name)
}

// Organization returns the organization name for the given organization resource.
func (r *Resource) Organization(ctx context.Context, orgID string) (*crm.Organization, error) {
	return r.crm.GetOrganization(ctx, "organizations/"+orgID)
}

// EnableBucketOnlyPolicy enable bucket only policy for the given bucket
func (r *Resource) EnableBucketOnlyPolicy(ctx context.Context, bucketName string) error {
	return r.storage.EnableBucketOnlyPolicy(ctx, bucketName)
}

// IfProjectWithinResources executes the provided function if the project ID is an ancestor of any provided resources.
func (r *Resource) IfProjectWithinResources(ctx context.Context, conf *Resources, projectID string, fn func() error) error {
	if err := r.IfProjectInFolders(ctx, conf.FolderIDs, projectID, fn); err != nil {
		return err
	}
	if err := r.IfProjectInProjects(ctx, conf.ProjectIDs, projectID, fn); err != nil {
		return err
	}
	if err := r.IfProjectInOrg(ctx, conf.OrganizationID, projectID, fn); err != nil {
		return err
	}
	return nil
}

// IfProjectInFolders will apply the function if the project ID is within the folder IDs.
func (r *Resource) IfProjectInFolders(ctx context.Context, ids []string, projectID string, fn func() error) error {
	if len(ids) == 0 {
		return nil
	}
	ancestors, err := r.GetProjectAncestry(ctx, projectID)
	if err != nil {
		return errors.Wrap(err, "failed to get project ancestry")
	}
	for _, resource := range ancestors {
		for _, folderID := range ids {
			if resource != "folders/"+folderID {
				continue
			}
			if err := fn(); err != nil {
				return err
			}
		}
	}
	return nil
}

// IfProjectInProjects will apply the function if the project ID is within the project IDs.
func (r *Resource) IfProjectInProjects(ctx context.Context, ids []string, projectID string, fn func() error) error {
	if len(ids) == 0 {
		return nil
	}
	for _, v := range ids {
		if v != projectID {
			continue
		}
		if err := fn(); err != nil {
			return err
		}
	}
	return nil
}

// IfProjectInOrg will apply the function if the project ID is within the organization.
func (r *Resource) IfProjectInOrg(ctx context.Context, orgID, projectID string, fn func() error) error {
	if orgID == "" {
		return nil
	}
	ancestors, err := r.GetProjectAncestry(ctx, projectID)
	if err != nil {
		return errors.Wrap(err, "failed to get project ancestry")
	}
	for _, resource := range ancestors {
		if resource == "organizations/"+orgID {
			if err := fn(); err != nil {
				return err
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package shared

import "github.com/uber/tchannel-go"

// The TChannel interface defines the dependencies for TChannel in Ringpop.
type TChannel interface {
	tchannel.Registrar
	PeerInfo() tchannel.LocalPeerInfo
	GetSubChannel(string, ...tchannel.SubChannelOption) *tchannel.SubChannel
}

// SubChannel represents a TChannel SubChannel as used in Ringpop.
type SubChannel interface {
	tchannel.Registrar
}
<commit_msg>Use narrower duck-typing instead of embedding the whole interface (#166)<commit_after>package shared

import "github.com/uber/tchannel-go"

// The TChannel interface defines the dependencies for TChannel in Ringpop.
type TChannel interface {
	Register(h tchannel.Handler, methodName string)
	PeerInfo() tchannel.LocalPeerInfo
	GetSubChannel(string, ...tchannel.SubChannelOption) *tchannel.SubChannel
}

// SubChannel represents a TChannel SubChannel as used in Ringpop.
type SubChannel interface {
	tchannel.Registrar
}
<|endoftext|>