<commit_before>package compact

import (
	"os"

	"github.com/kezhuw/leveldb/internal/file"
	"github.com/kezhuw/leveldb/internal/files"
	"github.com/kezhuw/leveldb/internal/keys"
	"github.com/kezhuw/leveldb/internal/memtable"
	"github.com/kezhuw/leveldb/internal/options"
	"github.com/kezhuw/leveldb/internal/table"
	"github.com/kezhuw/leveldb/internal/version"
)

func NewMemTableCompaction(dbname string, seq keys.Sequence, fileNumber uint64, maxLevel int, mem *memtable.MemTable, base *version.Version, opts *options.Options) Compactor {
	c := &memtableCompaction{
		mem:              mem,
		base:             base,
		maxLevel:         maxLevel,
		fileNumbers:      make([]uint64, 1),
		smallestSequence: seq,
		fs:               opts.FileSystem,
		options:          opts}
	c.fileNumbers[0] = fileNumber
	c.tableMeta.Number = fileNumber
	c.tableMeta.Size = 0
	c.tableName = files.TableFileName(dbname, fileNumber)
	return c
}

type memtableCompaction struct {
	mem              *memtable.MemTable
	base             *version.Version
	maxLevel         int
	smallestSequence keys.Sequence
	fileNumbers      []uint64

	fs      file.FileSystem
	options *options.Options

	tableMeta   version.FileMeta
	tableName   string
	tableWriter table.Writer
}

func (c *memtableCompaction) Level() int {
	return -1
}

func (c *memtableCompaction) Rewind() {
	c.tableMeta.Size = 0
}

func (c *memtableCompaction) FileNumbers() []uint64 {
	return c.fileNumbers
}

func (c *memtableCompaction) compact() (err error) {
	f, err := c.fs.Open(c.tableName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC)
	if err != nil {
		return err
	}
	defer func() {
		f.Close()
		if err != nil {
			c.fs.Remove(c.tableName)
		}
	}()

	it := c.mem.NewIterator()
	defer it.Release()

	if !it.First() {
		return it.Err()
	}

	w := &c.tableWriter
	w.Reset(f, c.options)

	w.Add(it.Key(), it.Value())
	c.tableMeta.Smallest = append(c.tableMeta.Smallest[:0], it.Key()...)
	c.tableMeta.Largest = append(c.tableMeta.Largest[:0], c.tableMeta.Smallest...)
	lastUserKey, lastSequence := c.tableMeta.Smallest.Split2()
	ucmp := c.options.Comparator.UserKeyComparator
	for it.Next() {
		key := it.Key()
		currentUserKey, currentSequence := keys.InternalKey(key).Split2()
		if ucmp.Compare(lastUserKey, currentUserKey) == 0 && lastSequence <= c.smallestSequence {
			continue
		}
		w.Add(key, it.Value())
		c.tableMeta.Largest = append(c.tableMeta.Largest[:0], key...)
		lastUserKey, lastSequence = c.tableMeta.Largest.UserKey(), currentSequence
	}
	if err := it.Err(); err != nil {
		return err
	}
	if err := w.Finish(); err != nil {
		return err
	}
	if err := f.Sync(); err != nil {
		return err
	}
	c.tableMeta.Size = uint64(w.FileSize())
	return nil
}

func (c *memtableCompaction) record(edit *version.Edit) {
	if c.tableMeta.Size == 0 {
		return
	}
	fmeta := version.FileMeta{
		Number:   c.tableMeta.Number,
		Size:     c.tableMeta.Size,
		Smallest: c.tableMeta.Smallest.Dup(),
		Largest:  c.tableMeta.Largest.Dup()}
	level := c.base.PickLevelForMemTableOutput(fmeta.Smallest, fmeta.Largest)
	if c.maxLevel > 0 && c.maxLevel < level {
		level = c.maxLevel
	}
	edit.AddedFiles = append(edit.AddedFiles[:0], version.LevelFileMeta{Level: level, FileMeta: fmeta})
}

func (c *memtableCompaction) Compact(edit *version.Edit) error {
	err := c.compact()
	if err != nil {
		return err
	}
	c.record(edit)
	return nil
}
<commit_msg>Fix wrong method to split internal key<commit_after>package compact

import (
	"os"

	"github.com/kezhuw/leveldb/internal/file"
	"github.com/kezhuw/leveldb/internal/files"
	"github.com/kezhuw/leveldb/internal/keys"
	"github.com/kezhuw/leveldb/internal/memtable"
	"github.com/kezhuw/leveldb/internal/options"
	"github.com/kezhuw/leveldb/internal/table"
	"github.com/kezhuw/leveldb/internal/version"
)

func NewMemTableCompaction(dbname string, seq keys.Sequence, fileNumber uint64, maxLevel int, mem *memtable.MemTable, base *version.Version, opts *options.Options) Compactor {
	c := &memtableCompaction{
		mem:              mem,
		base:             base,
		maxLevel:         maxLevel,
		fileNumbers:      make([]uint64, 1),
		smallestSequence: seq,
		fs:               opts.FileSystem,
		options:          opts}
	c.fileNumbers[0] = fileNumber
	c.tableMeta.Number = fileNumber
	c.tableMeta.Size = 0
	c.tableName = files.TableFileName(dbname, fileNumber)
	return c
}

type memtableCompaction struct {
	mem              *memtable.MemTable
	base             *version.Version
	maxLevel         int
	smallestSequence keys.Sequence
	fileNumbers      []uint64

	fs      file.FileSystem
	options *options.Options

	tableMeta   version.FileMeta
	tableName   string
	tableWriter table.Writer
}

func (c *memtableCompaction) Level() int {
	return -1
}

func (c *memtableCompaction) Rewind() {
	c.tableMeta.Size = 0
}

func (c *memtableCompaction) FileNumbers() []uint64 {
	return c.fileNumbers
}

func (c *memtableCompaction) compact() (err error) {
	f, err := c.fs.Open(c.tableName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC)
	if err != nil {
		return err
	}
	defer func() {
		f.Close()
		if err != nil {
			c.fs.Remove(c.tableName)
		}
	}()

	it := c.mem.NewIterator()
	defer it.Release()

	if !it.First() {
		return it.Err()
	}

	w := &c.tableWriter
	w.Reset(f, c.options)

	w.Add(it.Key(), it.Value())
	c.tableMeta.Smallest = append(c.tableMeta.Smallest[:0], it.Key()...)
	c.tableMeta.Largest = append(c.tableMeta.Largest[:0], c.tableMeta.Smallest...)
	lastUserKey, lastSequence, _ := c.tableMeta.Smallest.Split()
	ucmp := c.options.Comparator.UserKeyComparator
	for it.Next() {
		key := it.Key()
		currentUserKey, currentSequence, _ := keys.InternalKey(key).Split()
		if ucmp.Compare(lastUserKey, currentUserKey) == 0 && lastSequence <= c.smallestSequence {
			continue
		}
		w.Add(key, it.Value())
		c.tableMeta.Largest = append(c.tableMeta.Largest[:0], key...)
		lastUserKey, lastSequence = c.tableMeta.Largest.UserKey(), currentSequence
	}
	if err := it.Err(); err != nil {
		return err
	}
	if err := w.Finish(); err != nil {
		return err
	}
	if err := f.Sync(); err != nil {
		return err
	}
	c.tableMeta.Size = uint64(w.FileSize())
	return nil
}

func (c *memtableCompaction) record(edit *version.Edit) {
	if c.tableMeta.Size == 0 {
		return
	}
	fmeta := version.FileMeta{
		Number:   c.tableMeta.Number,
		Size:     c.tableMeta.Size,
		Smallest: c.tableMeta.Smallest.Dup(),
		Largest:  c.tableMeta.Largest.Dup()}
	level := c.base.PickLevelForMemTableOutput(fmeta.Smallest, fmeta.Largest)
	if c.maxLevel > 0 && c.maxLevel < level {
		level = c.maxLevel
	}
	edit.AddedFiles = append(edit.AddedFiles[:0], version.LevelFileMeta{Level: level, FileMeta: fmeta})
}

func (c *memtableCompaction) Compact(edit *version.Edit) error {
	err := c.compact()
	if err != nil {
		return err
	}
	c.record(edit)
	return nil
}
<|endoftext|>


<commit_before>package pool

import (
	"context"
	"fmt"
	"sync/atomic"
)

const (
	stateDefault = 0
	stateInited  = 1
	stateClosed  = 2
)

type BadConnError struct {
	wrapped error
}

var _ error = (*BadConnError)(nil)

func (e BadConnError) Error() string {
	return "pg: Conn is in a bad state"
}

func (e BadConnError) Unwrap() error {
	return e.wrapped
}

type SingleConnPool struct {
	pool Pooler

	state uint32 // atomic
	ch    chan *Conn

	level         int32 // atomic
	_badConnError atomic.Value
}

var _ Pooler = (*SingleConnPool)(nil)

func NewSingleConnPool(pool Pooler) *SingleConnPool {
	p, ok := pool.(*SingleConnPool)
	if !ok {
		p = &SingleConnPool{
			pool: pool,
			ch:   make(chan *Conn, 1),
		}
	}
	atomic.AddInt32(&p.level, 1)
	return p
}

func (p *SingleConnPool) SetConn(cn *Conn) {
	if atomic.CompareAndSwapUint32(&p.state, stateDefault, stateInited) {
		p.ch <- cn
	} else {
		panic("not reached")
	}
}

func (p *SingleConnPool) NewConn(c context.Context) (*Conn, error) {
	return p.pool.NewConn(c)
}

func (p *SingleConnPool) CloseConn(cn *Conn) error {
	return p.pool.CloseConn(cn)
}

func (p *SingleConnPool) Get(c context.Context) (*Conn, error) {
	// In worst case this races with Close which is not a very common operation.
	for i := 0; i < 1000; i++ {
		switch atomic.LoadUint32(&p.state) {
		case stateDefault:
			cn, err := p.pool.Get(c)
			if err != nil {
				return nil, err
			}
			if atomic.CompareAndSwapUint32(&p.state, stateDefault, stateInited) {
				return cn, nil
			}
			p.pool.Remove(cn, ErrClosed)
		case stateInited:
			if err := p.badConnError(); err != nil {
				return nil, err
			}
			cn, ok := <-p.ch
			if !ok {
				return nil, ErrClosed
			}
			return cn, nil
		case stateClosed:
			return nil, ErrClosed
		default:
			panic("not reached")
		}
	}
	return nil, fmt.Errorf("pg: SingleConnPool.Get: infinite loop")
}

func (p *SingleConnPool) Put(cn *Conn) {
	defer func() {
		if recover() != nil {
			p.freeConn(cn)
		}
	}()
	p.ch <- cn
}

func (p *SingleConnPool) freeConn(cn *Conn) {
	if err := p.badConnError(); err != nil {
		p.pool.Remove(cn, err)
	} else {
		p.pool.Put(cn)
	}
}

func (p *SingleConnPool) Remove(cn *Conn, reason error) {
	defer func() {
		if recover() != nil {
			p.pool.Remove(cn, ErrClosed)
		}
	}()
	p._badConnError.Store(BadConnError{wrapped: reason})
	p.ch <- cn
}

func (p *SingleConnPool) Len() int {
	switch atomic.LoadUint32(&p.state) {
	case stateDefault:
		return 0
	case stateInited:
		return 1
	case stateClosed:
		return 0
	default:
		panic("not reached")
	}
}

func (p *SingleConnPool) IdleLen() int {
	return len(p.ch)
}

func (p *SingleConnPool) Stats() *Stats {
	return &Stats{}
}

func (p *SingleConnPool) Close() error {
	level := atomic.AddInt32(&p.level, -1)
	if level > 0 {
		return nil
	}

	for i := 0; i < 1000; i++ {
		state := atomic.LoadUint32(&p.state)
		if state == stateClosed {
			return ErrClosed
		}
		if atomic.CompareAndSwapUint32(&p.state, state, stateClosed) {
			close(p.ch)
			cn, ok := <-p.ch
			if ok {
				p.freeConn(cn)
			}
			return nil
		}
	}

	return fmt.Errorf("pg: SingleConnPool.Close: infinite loop")
}

func (p *SingleConnPool) Reset() error {
	if p.badConnError() == nil {
		return nil
	}

	select {
	case cn, ok := <-p.ch:
		if !ok {
			return ErrClosed
		}
		p.pool.Remove(cn, ErrClosed)
		p._badConnError.Store(BadConnError{wrapped: nil})
	default:
		return fmt.Errorf("pg: SingleConnPool does not have a Conn")
	}

	if !atomic.CompareAndSwapUint32(&p.state, stateInited, stateDefault) {
		state := atomic.LoadUint32(&p.state)
		return fmt.Errorf("pg: invalid SingleConnPool state: %d", state)
	}

	return nil
}

func (p *SingleConnPool) badConnError() error {
	if v := p._badConnError.Load(); v != nil {
		err := v.(BadConnError)
		if err.wrapped != nil {
			return err
		}
	}
	return nil
}
<commit_msg>malign struct fields<commit_after>package pool

import (
	"context"
	"fmt"
	"sync/atomic"
)

const (
	stateDefault = 0
	stateInited  = 1
	stateClosed  = 2
)

type BadConnError struct {
	wrapped error
}

var _ error = (*BadConnError)(nil)

func (e BadConnError) Error() string {
	return "pg: Conn is in a bad state"
}

func (e BadConnError) Unwrap() error {
	return e.wrapped
}

type SingleConnPool struct {
	pool  Pooler
	level int32 // atomic

	state uint32 // atomic
	ch    chan *Conn

	_badConnError atomic.Value
}

var _ Pooler = (*SingleConnPool)(nil)

func NewSingleConnPool(pool Pooler) *SingleConnPool {
	p, ok := pool.(*SingleConnPool)
	if !ok {
		p = &SingleConnPool{
			pool: pool,
			ch:   make(chan *Conn, 1),
		}
	}
	atomic.AddInt32(&p.level, 1)
	return p
}

func (p *SingleConnPool) SetConn(cn *Conn) {
	if atomic.CompareAndSwapUint32(&p.state, stateDefault, stateInited) {
		p.ch <- cn
	} else {
		panic("not reached")
	}
}

func (p *SingleConnPool) NewConn(c context.Context) (*Conn, error) {
	return p.pool.NewConn(c)
}

func (p *SingleConnPool) CloseConn(cn *Conn) error {
	return p.pool.CloseConn(cn)
}

func (p *SingleConnPool) Get(c context.Context) (*Conn, error) {
	// In worst case this races with Close which is not a very common operation.
	for i := 0; i < 1000; i++ {
		switch atomic.LoadUint32(&p.state) {
		case stateDefault:
			cn, err := p.pool.Get(c)
			if err != nil {
				return nil, err
			}
			if atomic.CompareAndSwapUint32(&p.state, stateDefault, stateInited) {
				return cn, nil
			}
			p.pool.Remove(cn, ErrClosed)
		case stateInited:
			if err := p.badConnError(); err != nil {
				return nil, err
			}
			cn, ok := <-p.ch
			if !ok {
				return nil, ErrClosed
			}
			return cn, nil
		case stateClosed:
			return nil, ErrClosed
		default:
			panic("not reached")
		}
	}
	return nil, fmt.Errorf("pg: SingleConnPool.Get: infinite loop")
}

func (p *SingleConnPool) Put(cn *Conn) {
	defer func() {
		if recover() != nil {
			p.freeConn(cn)
		}
	}()
	p.ch <- cn
}

func (p *SingleConnPool) freeConn(cn *Conn) {
	if err := p.badConnError(); err != nil {
		p.pool.Remove(cn, err)
	} else {
		p.pool.Put(cn)
	}
}

func (p *SingleConnPool) Remove(cn *Conn, reason error) {
	defer func() {
		if recover() != nil {
			p.pool.Remove(cn, ErrClosed)
		}
	}()
	p._badConnError.Store(BadConnError{wrapped: reason})
	p.ch <- cn
}

func (p *SingleConnPool) Len() int {
	switch atomic.LoadUint32(&p.state) {
	case stateDefault:
		return 0
	case stateInited:
		return 1
	case stateClosed:
		return 0
	default:
		panic("not reached")
	}
}

func (p *SingleConnPool) IdleLen() int {
	return len(p.ch)
}

func (p *SingleConnPool) Stats() *Stats {
	return &Stats{}
}

func (p *SingleConnPool) Close() error {
	level := atomic.AddInt32(&p.level, -1)
	if level > 0 {
		return nil
	}

	for i := 0; i < 1000; i++ {
		state := atomic.LoadUint32(&p.state)
		if state == stateClosed {
			return ErrClosed
		}
		if atomic.CompareAndSwapUint32(&p.state, state, stateClosed) {
			close(p.ch)
			cn, ok := <-p.ch
			if ok {
				p.freeConn(cn)
			}
			return nil
		}
	}

	return fmt.Errorf("pg: SingleConnPool.Close: infinite loop")
}

func (p *SingleConnPool) Reset() error {
	if p.badConnError() == nil {
		return nil
	}

	select {
	case cn, ok := <-p.ch:
		if !ok {
			return ErrClosed
		}
		p.pool.Remove(cn, ErrClosed)
		p._badConnError.Store(BadConnError{wrapped: nil})
	default:
		return fmt.Errorf("pg: SingleConnPool does not have a Conn")
	}

	if !atomic.CompareAndSwapUint32(&p.state, stateInited, stateDefault) {
		state := atomic.LoadUint32(&p.state)
		return fmt.Errorf("pg: invalid SingleConnPool state: %d", state)
	}

	return nil
}

func (p *SingleConnPool) badConnError() error {
	if v := p._badConnError.Load(); v != nil {
		err := v.(BadConnError)
		if err.wrapped != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 *
 * Copyright 2020 gRPC 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.
 *
 */

// Binary client for xDS interop tests.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	testpb "google.golang.org/grpc/interop/grpc_testing"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
	_ "google.golang.org/grpc/xds"
)

type statsWatcherKey struct {
	startID int32
	endID   int32
}

// rpcInfo contains the rpc type and the hostname where the response is received
// from.
type rpcInfo struct {
	typ      string
	hostname string
}

type statsWatcher struct {
	rpcsByPeer    map[string]int32
	rpcsByType    map[string]map[string]int32
	numFailures   int32
	remainingRpcs int32
	chanHosts     chan *rpcInfo
}

func (watcher *statsWatcher) buildResp() *testpb.LoadBalancerStatsResponse {
	rpcsByType := make(map[string]*testpb.LoadBalancerStatsResponse_RpcsByPeer, len(watcher.rpcsByType))
	for t, rpcsByPeer := range watcher.rpcsByType {
		rpcsByType[t] = &testpb.LoadBalancerStatsResponse_RpcsByPeer{
			RpcsByPeer: rpcsByPeer,
		}
	}

	return &testpb.LoadBalancerStatsResponse{
		NumFailures:  watcher.numFailures + watcher.remainingRpcs,
		RpcsByPeer:   watcher.rpcsByPeer,
		RpcsByMethod: rpcsByType,
	}
}

var (
	failOnFailedRPC = flag.Bool("fail_on_failed_rpc", false, "Fail client if any RPCs fail")
	numChannels     = flag.Int("num_channels", 1, "Num of channels")
	printResponse   = flag.Bool("print_response", false, "Write RPC response to stdout")
	qps             = flag.Int("qps", 1, "QPS per channel, for each type of RPC")
	rpc             = flag.String("rpc", "UnaryCall", "Types of RPCs to make, ',' separated string. RPCs can be EmptyCall or UnaryCall")
	rpcMetadata     = flag.String("metadata", "", "The metadata to send with RPC, in format EmptyCall:key1:value1,UnaryCall:key2:value2")
	rpcTimeout      = flag.Duration("rpc_timeout", 20*time.Second, "Per RPC timeout")
	server          = flag.String("server", "localhost:8080", "Address of server to connect to")
	statsPort       = flag.Int("stats_port", 8081, "Port to expose peer distribution stats service")

	mu               sync.Mutex
	currentRequestID int32
	watchers         = make(map[statsWatcherKey]*statsWatcher)
)

type statsService struct {
	testpb.UnimplementedLoadBalancerStatsServiceServer
}

// Wait for the next LoadBalancerStatsRequest.GetNumRpcs to start and complete,
// and return the distribution of remote peers. This is essentially a clientside
// LB reporting mechanism that is designed to be queried by an external test
// driver when verifying that the client is distributing RPCs as expected.
func (s *statsService) GetClientStats(ctx context.Context, in *testpb.LoadBalancerStatsRequest) (*testpb.LoadBalancerStatsResponse, error) {
	mu.Lock()
	watcherKey := statsWatcherKey{currentRequestID, currentRequestID + in.GetNumRpcs()}
	watcher, ok := watchers[watcherKey]
	if !ok {
		watcher = &statsWatcher{
			rpcsByPeer:    make(map[string]int32),
			rpcsByType:    make(map[string]map[string]int32),
			numFailures:   0,
			remainingRpcs: in.GetNumRpcs(),
			chanHosts:     make(chan *rpcInfo),
		}
		watchers[watcherKey] = watcher
	}
	mu.Unlock()

	ctx, cancel := context.WithTimeout(ctx, time.Duration(in.GetTimeoutSec())*time.Second)
	defer cancel()

	defer func() {
		mu.Lock()
		delete(watchers, watcherKey)
		mu.Unlock()
	}()

	// Wait until the requested RPCs have all been recorded or timeout occurs.
	for {
		select {
		case info := <-watcher.chanHosts:
			if info != nil {
				watcher.rpcsByPeer[info.hostname]++

				rpcsByPeerForType := watcher.rpcsByType[info.typ]
				if rpcsByPeerForType == nil {
					rpcsByPeerForType = make(map[string]int32)
					watcher.rpcsByType[info.typ] = rpcsByPeerForType
				}
				rpcsByPeerForType[info.hostname]++
			} else {
				watcher.numFailures++
			}
			watcher.remainingRpcs--
			if watcher.remainingRpcs == 0 {
				return watcher.buildResp(), nil
			}
		case <-ctx.Done():
			grpclog.Info("Timed out, returning partial stats")
			return watcher.buildResp(), nil
		}
	}
}

const (
	unaryCall string = "UnaryCall"
	emptyCall string = "EmptyCall"
)

func parseRPCTypes(rpcStr string) (ret []string) {
	if len(rpcStr) == 0 {
		return []string{unaryCall}
	}

	rpcs := strings.Split(rpcStr, ",")
	for _, r := range rpcs {
		switch r {
		case unaryCall, emptyCall:
			ret = append(ret, r)
		default:
			flag.PrintDefaults()
			log.Fatalf("unsupported RPC type: %v", r)
		}
	}
	return
}

type rpcConfig struct {
	typ string
	md  metadata.MD
}

// parseRPCMetadata turns EmptyCall:key1:value1 into
//   {typ: emptyCall, md: {key1:value1}}.
func parseRPCMetadata(rpcMetadataStr string, rpcs []string) []*rpcConfig {
	rpcMetadataSplit := strings.Split(rpcMetadataStr, ",")
	rpcsToMD := make(map[string][]string)
	for _, rm := range rpcMetadataSplit {
		rmSplit := strings.Split(rm, ":")
		if len(rmSplit)%2 != 1 {
			log.Fatalf("invalid metadata config %v, want EmptyCall:key1:value1", rm)
		}
		rpcsToMD[rmSplit[0]] = append(rpcsToMD[rmSplit[0]], rmSplit[1:]...)
	}
	ret := make([]*rpcConfig, 0, len(rpcs))
	for _, rpcT := range rpcs {
		rpcC := &rpcConfig{
			typ: rpcT,
		}
		if md := rpcsToMD[string(rpcT)]; len(md) > 0 {
			rpcC.md = metadata.Pairs(md...)
		}
		ret = append(ret, rpcC)
	}
	return ret
}

func main() {
	flag.Parse()
	rpcCfgs := parseRPCMetadata(*rpcMetadata, parseRPCTypes(*rpc))

	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *statsPort))
	if err != nil {
		grpclog.Fatalf("failed to listen: %v", err)
	}
	s := grpc.NewServer()
	defer s.Stop()
	testpb.RegisterLoadBalancerStatsServiceServer(s, &statsService{})
	go s.Serve(lis)

	clients := make([]testpb.TestServiceClient, *numChannels)
	for i := 0; i < *numChannels; i++ {
		conn, err := grpc.DialContext(context.Background(), *server, grpc.WithInsecure())
		if err != nil {
			grpclog.Fatalf("Fail to dial: %v", err)
		}
		defer conn.Close()
		clients[i] = testpb.NewTestServiceClient(conn)
	}
	ticker := time.NewTicker(time.Second / time.Duration(*qps**numChannels))
	defer ticker.Stop()
	sendRPCs(clients, rpcCfgs, ticker)
}

func makeOneRPC(c testpb.TestServiceClient, cfg *rpcConfig) (*peer.Peer, *rpcInfo, error) {
	ctx, cancel := context.WithTimeout(context.Background(), *rpcTimeout)
	defer cancel()

	if len(cfg.md) != 0 {
		ctx = metadata.NewOutgoingContext(ctx, cfg.md)
	}
	info := rpcInfo{typ: cfg.typ}

	var (
		p      peer.Peer
		header metadata.MD
		err    error
	)
	switch cfg.typ {
	case unaryCall:
		var resp *testpb.SimpleResponse
		resp, err = c.UnaryCall(ctx, &testpb.SimpleRequest{FillServerId: true}, grpc.Peer(&p), grpc.Header(&header))
		// For UnaryCall, also read hostname from response, in case the server
		// isn't updated to send headers.
		if resp != nil {
			info.hostname = resp.Hostname
		}
	case emptyCall:
		_, err = c.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&p), grpc.Header(&header))
	}
	if err != nil {
		return nil, nil, err
	}

	hosts := header["hostname"]
	if len(hosts) > 0 {
		info.hostname = hosts[0]
	}
	return &p, &info, err
}

func sendRPCs(clients []testpb.TestServiceClient, cfgs []*rpcConfig, ticker *time.Ticker) {
	var i int
	for range ticker.C {
		go func(i int) {
			// Get and increment request ID, and save a list of watchers that
			// are interested in this RPC.
			mu.Lock()
			savedRequestID := currentRequestID
			currentRequestID++
			savedWatchers := []*statsWatcher{}
			for key, value := range watchers {
				if key.startID <= savedRequestID && savedRequestID < key.endID {
					savedWatchers = append(savedWatchers, value)
				}
			}
			mu.Unlock()

			c := clients[i]

			for _, cfg := range cfgs {
				p, info, err := makeOneRPC(c, cfg)

				for _, watcher := range savedWatchers {
					// This sends an empty string if the RPC failed.
					watcher.chanHosts <- info
				}
				if err != nil && *failOnFailedRPC {
					grpclog.Fatalf("RPC failed: %v", err)
				}
				if *printResponse {
					if err == nil {
						if cfg.typ == unaryCall {
							// Need to keep this format, because some tests are
							// relying on stdout.
							fmt.Printf("Greeting: Hello world, this is %s, from %v\n", info.hostname, p.Addr)
						} else {
							fmt.Printf("RPC %q, from host %s, addr %v\n", cfg.typ, info.hostname, p.Addr)
						}
					} else {
						fmt.Printf("RPC %q, failed with %v\n", cfg.typ, err)
					}
				}
			}
		}(i)
		i = (i + 1) % len(clients)
	}
}
<commit_msg>interop: update --fail_on_failed_rpc to wait for initial success (#3763)<commit_after>/*
 *
 * Copyright 2020 gRPC 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.
 *
 */

// Binary client for xDS interop tests.
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	testpb "google.golang.org/grpc/interop/grpc_testing"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
	_ "google.golang.org/grpc/xds"
)

type statsWatcherKey struct {
	startID int32
	endID   int32
}

// rpcInfo contains the rpc type and the hostname where the response is received
// from.
type rpcInfo struct {
	typ      string
	hostname string
}

type statsWatcher struct {
	rpcsByPeer    map[string]int32
	rpcsByType    map[string]map[string]int32
	numFailures   int32
	remainingRpcs int32
	chanHosts     chan *rpcInfo
}

func (watcher *statsWatcher) buildResp() *testpb.LoadBalancerStatsResponse {
	rpcsByType := make(map[string]*testpb.LoadBalancerStatsResponse_RpcsByPeer, len(watcher.rpcsByType))
	for t, rpcsByPeer := range watcher.rpcsByType {
		rpcsByType[t] = &testpb.LoadBalancerStatsResponse_RpcsByPeer{
			RpcsByPeer: rpcsByPeer,
		}
	}

	return &testpb.LoadBalancerStatsResponse{
		NumFailures:  watcher.numFailures + watcher.remainingRpcs,
		RpcsByPeer:   watcher.rpcsByPeer,
		RpcsByMethod: rpcsByType,
	}
}

var (
	failOnFailedRPC = flag.Bool("fail_on_failed_rpc", false, "Fail client if any RPCs fail after first success")
	numChannels     = flag.Int("num_channels", 1, "Num of channels")
	printResponse   = flag.Bool("print_response", false, "Write RPC response to stdout")
	qps             = flag.Int("qps", 1, "QPS per channel, for each type of RPC")
	rpc             = flag.String("rpc", "UnaryCall", "Types of RPCs to make, ',' separated string. RPCs can be EmptyCall or UnaryCall")
	rpcMetadata     = flag.String("metadata", "", "The metadata to send with RPC, in format EmptyCall:key1:value1,UnaryCall:key2:value2")
	rpcTimeout      = flag.Duration("rpc_timeout", 20*time.Second, "Per RPC timeout")
	server          = flag.String("server", "localhost:8080", "Address of server to connect to")
	statsPort       = flag.Int("stats_port", 8081, "Port to expose peer distribution stats service")

	mu               sync.Mutex
	currentRequestID int32
	watchers         = make(map[statsWatcherKey]*statsWatcher)

	// 0 or 1 representing an RPC has succeeded. Use hasRPCSucceeded and
	// setRPCSucceeded to access in a safe manner.
	rpcSucceeded uint32
)

type statsService struct {
	testpb.UnimplementedLoadBalancerStatsServiceServer
}

func hasRPCSucceeded() bool {
	return atomic.LoadUint32(&rpcSucceeded) > 0
}

func setRPCSucceeded() {
	atomic.StoreUint32(&rpcSucceeded, 1)
}

// Wait for the next LoadBalancerStatsRequest.GetNumRpcs to start and complete,
// and return the distribution of remote peers. This is essentially a clientside
// LB reporting mechanism that is designed to be queried by an external test
// driver when verifying that the client is distributing RPCs as expected.
func (s *statsService) GetClientStats(ctx context.Context, in *testpb.LoadBalancerStatsRequest) (*testpb.LoadBalancerStatsResponse, error) {
	mu.Lock()
	watcherKey := statsWatcherKey{currentRequestID, currentRequestID + in.GetNumRpcs()}
	watcher, ok := watchers[watcherKey]
	if !ok {
		watcher = &statsWatcher{
			rpcsByPeer:    make(map[string]int32),
			rpcsByType:    make(map[string]map[string]int32),
			numFailures:   0,
			remainingRpcs: in.GetNumRpcs(),
			chanHosts:     make(chan *rpcInfo),
		}
		watchers[watcherKey] = watcher
	}
	mu.Unlock()

	ctx, cancel := context.WithTimeout(ctx, time.Duration(in.GetTimeoutSec())*time.Second)
	defer cancel()

	defer func() {
		mu.Lock()
		delete(watchers, watcherKey)
		mu.Unlock()
	}()

	// Wait until the requested RPCs have all been recorded or timeout occurs.
	for {
		select {
		case info := <-watcher.chanHosts:
			if info != nil {
				watcher.rpcsByPeer[info.hostname]++

				rpcsByPeerForType := watcher.rpcsByType[info.typ]
				if rpcsByPeerForType == nil {
					rpcsByPeerForType = make(map[string]int32)
					watcher.rpcsByType[info.typ] = rpcsByPeerForType
				}
				rpcsByPeerForType[info.hostname]++
			} else {
				watcher.numFailures++
			}
			watcher.remainingRpcs--
			if watcher.remainingRpcs == 0 {
				return watcher.buildResp(), nil
			}
		case <-ctx.Done():
			grpclog.Info("Timed out, returning partial stats")
			return watcher.buildResp(), nil
		}
	}
}

const (
	unaryCall string = "UnaryCall"
	emptyCall string = "EmptyCall"
)

func parseRPCTypes(rpcStr string) (ret []string) {
	if len(rpcStr) == 0 {
		return []string{unaryCall}
	}

	rpcs := strings.Split(rpcStr, ",")
	for _, r := range rpcs {
		switch r {
		case unaryCall, emptyCall:
			ret = append(ret, r)
		default:
			flag.PrintDefaults()
			log.Fatalf("unsupported RPC type: %v", r)
		}
	}
	return
}

type rpcConfig struct {
	typ string
	md  metadata.MD
}

// parseRPCMetadata turns EmptyCall:key1:value1 into
//   {typ: emptyCall, md: {key1:value1}}.
func parseRPCMetadata(rpcMetadataStr string, rpcs []string) []*rpcConfig {
	rpcMetadataSplit := strings.Split(rpcMetadataStr, ",")
	rpcsToMD := make(map[string][]string)
	for _, rm := range rpcMetadataSplit {
		rmSplit := strings.Split(rm, ":")
		if len(rmSplit)%2 != 1 {
			log.Fatalf("invalid metadata config %v, want EmptyCall:key1:value1", rm)
		}
		rpcsToMD[rmSplit[0]] = append(rpcsToMD[rmSplit[0]], rmSplit[1:]...)
	}
	ret := make([]*rpcConfig, 0, len(rpcs))
	for _, rpcT := range rpcs {
		rpcC := &rpcConfig{
			typ: rpcT,
		}
		if md := rpcsToMD[string(rpcT)]; len(md) > 0 {
			rpcC.md = metadata.Pairs(md...)
		}
		ret = append(ret, rpcC)
	}
	return ret
}

func main() {
	flag.Parse()
	rpcCfgs := parseRPCMetadata(*rpcMetadata, parseRPCTypes(*rpc))

	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *statsPort))
	if err != nil {
		grpclog.Fatalf("failed to listen: %v", err)
	}
	s := grpc.NewServer()
	defer s.Stop()
	testpb.RegisterLoadBalancerStatsServiceServer(s, &statsService{})
	go s.Serve(lis)

	clients := make([]testpb.TestServiceClient, *numChannels)
	for i := 0; i < *numChannels; i++ {
		conn, err := grpc.DialContext(context.Background(), *server, grpc.WithInsecure())
		if err != nil {
			grpclog.Fatalf("Fail to dial: %v", err)
		}
		defer conn.Close()
		clients[i] = testpb.NewTestServiceClient(conn)
	}
	ticker := time.NewTicker(time.Second / time.Duration(*qps**numChannels))
	defer ticker.Stop()
	sendRPCs(clients, rpcCfgs, ticker)
}

func makeOneRPC(c testpb.TestServiceClient, cfg *rpcConfig) (*peer.Peer, *rpcInfo, error) {
	ctx, cancel := context.WithTimeout(context.Background(), *rpcTimeout)
	defer cancel()

	if len(cfg.md) != 0 {
		ctx = metadata.NewOutgoingContext(ctx, cfg.md)
	}
	info := rpcInfo{typ: cfg.typ}

	var (
		p      peer.Peer
		header metadata.MD
		err    error
	)
	switch cfg.typ {
	case unaryCall:
		var resp *testpb.SimpleResponse
		resp, err = c.UnaryCall(ctx, &testpb.SimpleRequest{FillServerId: true}, grpc.Peer(&p), grpc.Header(&header))
		// For UnaryCall, also read hostname from response, in case the server
		// isn't updated to send headers.
		if resp != nil {
			info.hostname = resp.Hostname
		}
	case emptyCall:
		_, err = c.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&p), grpc.Header(&header))
	}
	if err != nil {
		return nil, nil, err
	}

	hosts := header["hostname"]
	if len(hosts) > 0 {
		info.hostname = hosts[0]
	}
	return &p, &info, err
}

func sendRPCs(clients []testpb.TestServiceClient, cfgs []*rpcConfig, ticker *time.Ticker) {
	var i int
	for range ticker.C {
		go func(i int) {
			// Get and increment request ID, and save a list of watchers that
			// are interested in this RPC.
			mu.Lock()
			savedRequestID := currentRequestID
			currentRequestID++
			savedWatchers := []*statsWatcher{}
			for key, value := range watchers {
				if key.startID <= savedRequestID && savedRequestID < key.endID {
					savedWatchers = append(savedWatchers, value)
				}
			}
			mu.Unlock()

			c := clients[i]

			for _, cfg := range cfgs {
				p, info, err := makeOneRPC(c, cfg)

				for _, watcher := range savedWatchers {
					// This sends an empty string if the RPC failed.
					watcher.chanHosts <- info
				}
				if err != nil && *failOnFailedRPC && hasRPCSucceeded() {
					grpclog.Fatalf("RPC failed: %v", err)
				}
				if err == nil {
					setRPCSucceeded()
				}
				if *printResponse {
					if err == nil {
						if cfg.typ == unaryCall {
							// Need to keep this format, because some tests are
							// relying on stdout.
							fmt.Printf("Greeting: Hello world, this is %s, from %v\n", info.hostname, p.Addr)
						} else {
							fmt.Printf("RPC %q, from host %s, addr %v\n", cfg.typ, info.hostname, p.Addr)
						}
					} else {
						fmt.Printf("RPC %q, failed with %v\n", cfg.typ, err)
					}
				}
			}
		}(i)
		i = (i + 1) % len(clients)
	}
}
<|endoftext|>


<commit_before>// +build linux

/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package cm

import (
	"fmt"
	"io/ioutil"
	"regexp"
	"strconv"
	"time"

	"github.com/golang/glog"
	"github.com/opencontainers/runc/libcontainer/cgroups/fs"
	"github.com/opencontainers/runc/libcontainer/configs"
	"k8s.io/apimachinery/pkg/util/wait"
	kubecm "k8s.io/kubernetes/pkg/kubelet/cm"
	"k8s.io/kubernetes/pkg/kubelet/qos"
	utilversion "k8s.io/kubernetes/pkg/util/version"

	"k8s.io/kubernetes/pkg/kubelet/dockershim/libdocker"
)

const (
	// The percent of the machine memory capacity.
	dockerMemoryLimitThresholdPercent = kubecm.DockerMemoryLimitThresholdPercent

	// The minimum memory limit allocated to docker container.
	minDockerMemoryLimit = kubecm.MinDockerMemoryLimit

	// The Docker OOM score adjustment.
	dockerOOMScoreAdj = qos.DockerOOMScoreAdj
)

var (
	memoryCapacityRegexp = regexp.MustCompile(`MemTotal:\s*([0-9]+) kB`)
)

func NewContainerManager(cgroupsName string, client libdocker.Interface) ContainerManager {
	return &containerManager{
		cgroupsName: cgroupsName,
		client:      client,
	}
}

type containerManager struct {
	// Docker client.
	client libdocker.Interface
	// Name of the cgroups.
	cgroupsName string
	// Manager for the cgroups.
	cgroupsManager *fs.Manager
}

func (m *containerManager) Start() error {
	// TODO: check if the required cgroups are mounted.
	if len(m.cgroupsName) != 0 {
		manager, err := createCgroupManager(m.cgroupsName)
		if err != nil {
			return err
		}
		m.cgroupsManager = manager
	}
	go wait.Until(m.doWork, 5*time.Minute, wait.NeverStop)
	return nil
}

func (m *containerManager) doWork() {
	v, err := m.client.Version()
	if err != nil {
		glog.Errorf("Unable to get docker version: %v", err)
		return
	}
	version, err := utilversion.ParseGeneric(v.APIVersion)
	if err != nil {
		glog.Errorf("Unable to parse docker version %q: %v", v.APIVersion, err)
		return
	}
	// EnsureDockerInContainer does two things.
	//   1. Ensure processes run in the cgroups if m.cgroupsManager is not nil.
	//   2. Ensure processes have the OOM score applied.
	if err := kubecm.EnsureDockerInContainer(version, dockerOOMScoreAdj, m.cgroupsManager); err != nil {
		glog.Errorf("Unable to ensure the docker processes run in the desired containers")
	}
}

func createCgroupManager(name string) (*fs.Manager, error) {
	var memoryLimit uint64
	memoryCapacity, err := getMemoryCapacity()
	if err != nil || memoryCapacity*dockerMemoryLimitThresholdPercent/100 < minDockerMemoryLimit {
		memoryLimit = minDockerMemoryLimit
	}
	glog.V(2).Infof("Configure resource-only container %q with memory limit: %d", name, memoryLimit)

	allowAllDevices := true
	cm := &fs.Manager{
		Cgroups: &configs.Cgroup{
			Parent: "/",
			Name:   name,
			Resources: &configs.Resources{
				Memory:          int64(memoryLimit),
				MemorySwap:      -1,
				AllowAllDevices: &allowAllDevices,
			},
		},
	}
	return cm, nil
}

// getMemoryCapacity returns the memory capacity on the machine in bytes.
func getMemoryCapacity() (uint64, error) {
	out, err := ioutil.ReadFile("/proc/meminfo")
	if err != nil {
		return 0, err
	}
	return parseCapacity(out, memoryCapacityRegexp)
}

// parseCapacity matches a Regexp in a []byte, returning the resulting value in bytes.
// Assumes that the value matched by the Regexp is in KB.
func parseCapacity(b []byte, r *regexp.Regexp) (uint64, error) {
	matches := r.FindSubmatch(b)
	if len(matches) != 2 {
		return 0, fmt.Errorf("failed to match regexp in output: %q", string(b))
	}
	m, err := strconv.ParseUint(string(matches[1]), 10, 64)
	if err != nil {
		return 0, err
	}

	// Convert to bytes.
	return m * 1024, err
}
<commit_msg>Add detailed err in ensure docker process error<commit_after>// +build linux

/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package cm

import (
	"fmt"
	"io/ioutil"
	"regexp"
	"strconv"
	"time"

	"github.com/golang/glog"
	"github.com/opencontainers/runc/libcontainer/cgroups/fs"
	"github.com/opencontainers/runc/libcontainer/configs"
	"k8s.io/apimachinery/pkg/util/wait"
	kubecm "k8s.io/kubernetes/pkg/kubelet/cm"
	"k8s.io/kubernetes/pkg/kubelet/qos"
	utilversion "k8s.io/kubernetes/pkg/util/version"

	"k8s.io/kubernetes/pkg/kubelet/dockershim/libdocker"
)

const (
	// The percent of the machine memory capacity.
	dockerMemoryLimitThresholdPercent = kubecm.DockerMemoryLimitThresholdPercent

	// The minimum memory limit allocated to docker container.
	minDockerMemoryLimit = kubecm.MinDockerMemoryLimit

	// The Docker OOM score adjustment.
	dockerOOMScoreAdj = qos.DockerOOMScoreAdj
)

var (
	memoryCapacityRegexp = regexp.MustCompile(`MemTotal:\s*([0-9]+) kB`)
)

func NewContainerManager(cgroupsName string, client libdocker.Interface) ContainerManager {
	return &containerManager{
		cgroupsName: cgroupsName,
		client:      client,
	}
}

type containerManager struct {
	// Docker client.
	client libdocker.Interface
	// Name of the cgroups.
	cgroupsName string
	// Manager for the cgroups.
	cgroupsManager *fs.Manager
}

func (m *containerManager) Start() error {
	// TODO: check if the required cgroups are mounted.
	if len(m.cgroupsName) != 0 {
		manager, err := createCgroupManager(m.cgroupsName)
		if err != nil {
			return err
		}
		m.cgroupsManager = manager
	}
	go wait.Until(m.doWork, 5*time.Minute, wait.NeverStop)
	return nil
}

func (m *containerManager) doWork() {
	v, err := m.client.Version()
	if err != nil {
		glog.Errorf("Unable to get docker version: %v", err)
		return
	}
	version, err := utilversion.ParseGeneric(v.APIVersion)
	if err != nil {
		glog.Errorf("Unable to parse docker version %q: %v", v.APIVersion, err)
		return
	}
	// EnsureDockerInContainer does two things.
	//   1. Ensure processes run in the cgroups if m.cgroupsManager is not nil.
	//   2. Ensure processes have the OOM score applied.
	if err := kubecm.EnsureDockerInContainer(version, dockerOOMScoreAdj, m.cgroupsManager); err != nil {
		glog.Errorf("Unable to ensure the docker processes run in the desired containers: %v", err)
	}
}

func createCgroupManager(name string) (*fs.Manager, error) {
	var memoryLimit uint64
	memoryCapacity, err := getMemoryCapacity()
	if err != nil || memoryCapacity*dockerMemoryLimitThresholdPercent/100 < minDockerMemoryLimit {
		memoryLimit = minDockerMemoryLimit
	}
	glog.V(2).Infof("Configure resource-only container %q with memory limit: %d", name, memoryLimit)

	allowAllDevices := true
	cm := &fs.Manager{
		Cgroups: &configs.Cgroup{
			Parent: "/",
			Name:   name,
			Resources: &configs.Resources{
				Memory:          int64(memoryLimit),
				MemorySwap:      -1,
				AllowAllDevices: &allowAllDevices,
			},
		},
	}
	return cm, nil
}

// getMemoryCapacity returns the memory capacity on the machine in bytes.
func getMemoryCapacity() (uint64, error) {
	out, err := ioutil.ReadFile("/proc/meminfo")
	if err != nil {
		return 0, err
	}
	return parseCapacity(out, memoryCapacityRegexp)
}

// parseCapacity matches a Regexp in a []byte, returning the resulting value in bytes.
// Assumes that the value matched by the Regexp is in KB.
func parseCapacity(b []byte, r *regexp.Regexp) (uint64, error) {
	matches := r.FindSubmatch(b)
	if len(matches) != 2 {
		return 0, fmt.Errorf("failed to match regexp in output: %q", string(b))
	}
	m, err := strconv.ParseUint(string(matches[1]), 10, 64)
	if err != nil {
		return 0, err
	}

	// Convert to bytes.
	return m * 1024, err
}
<|endoftext|>


<commit_before>// Copyright 2021 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package telemetry

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"sync"
	"time"

	pingcapErrors "github.com/pingcap/errors"
	"github.com/pingcap/tidb/domain/infosync"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/util/logutil"
	pmodel "github.com/prometheus/common/model"
	"go.uber.org/zap"
)

type slowQueryStats struct {
	// Slow Query statistic buckets
	SQBInfo *SlowQueryBucket `json:"slowQueryBucket"`
}

// SlowQueryBucket records the statistic information of slow query buckets
// Buckets:   prometheus.ExponentialBuckets(0.001, 2, 28), // 1ms ~ 1.5days  // defined in metrics/server.go
type SlowQueryBucket map[string]int

func (bucketMap SlowQueryBucket) String() string {
	if bucketMap == nil {
		return "nil"
	}
	var retStr string = "{"
	for k, v := range bucketMap {
		retStr += k + ":" + strconv.Itoa(v) + ","
	}
	retStr = retStr[:len(retStr)-1]
	return retStr
}

const slowQueryBucketNum = 29 //prometheus.ExponentialBuckets(0.001, 2, 28), and 1 more +Inf

var (
	// lastSQBInfo records last statistic information of slow query buckets
	lastSQBInfo SlowQueryBucket
	// currentSQBInfo records current statitic information of slow query buckets
	currentSQBInfo SlowQueryBucket
	slowQueryLock  sync.Mutex
)

func getSlowQueryStats(ctx sessionctx.Context) (*slowQueryStats, error) {
	slowQueryBucket, err := getSlowQueryBucket(ctx)
	if err != nil {
		logutil.BgLogger().Info(err.Error())
		return nil, err
	}

	return &slowQueryStats{slowQueryBucket}, nil
}

// getSlowQueryBucket genenrates the delta SlowQueryBucket to report
func getSlowQueryBucket(ctx sessionctx.Context) (*SlowQueryBucket, error) {
	// update currentSQBInfo first, then gen delta
	if err := updateCurrentSQB(ctx); err != nil {
		return nil, err
	}
	delta := calculateDeltaSQB()
	return delta, nil
}

// updateCurrentSQB records current slow query buckets
func updateCurrentSQB(ctx sessionctx.Context) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = pingcapErrors.Errorf(fmt.Sprintln(r))
		}
	}()

	pQueryCtx, cancel := context.WithTimeout(context.Background(), time.Second*30)
	defer cancel()
	pQueryTs := time.Now().Add(-time.Minute)
	promQL := "tidb_server_slow_query_process_duration_seconds_bucket{sql_type=\"general\"}"
	value, err := querySQLMetric(pQueryCtx, pQueryTs, promQL)

	if err != nil && err != infosync.ErrPrometheusAddrIsNotSet {
		logutil.BgLogger().Info("querySlowQueryMetric got error")
		return err
	}
	if value == nil {
		return
	}
	if value.Type() != pmodel.ValVector {
		return errors.New("Prom vector expected, got " + value.Type().String())
	}
	promVec := value.(pmodel.Vector)
	slowQueryLock.Lock()
	for _, sample := range promVec {
		metric := sample.Metric
		bucketName := metric["le"] //hardcode bucket upper bound
		currentSQBInfo[string(bucketName)] = int(sample.Value)
	}
	slowQueryLock.Unlock()
	return nil
}

// calculateDeltaSQB calculate the delta between current slow query bucket and last slow query bucket
func calculateDeltaSQB() *SlowQueryBucket {
	deltaMap := make(SlowQueryBucket)
	slowQueryLock.Lock()
	for key, value := range currentSQBInfo {
		deltaMap[key] = value - (lastSQBInfo)[key]
	}
	slowQueryLock.Unlock()
	return &deltaMap
}

// init Init lastSQBInfo, follow the definition of metrics/server.go
// Buckets:   prometheus.ExponentialBuckets(0.001, 2, 28), // 1ms ~ 1.5days
func init() {
	lastSQBInfo = make(SlowQueryBucket)
	currentSQBInfo = make(SlowQueryBucket)

	bucketBase := 0.001 // From 0.001 to 134217.728, total 28 float number; the 29th is +Inf
	for i := 0; i < slowQueryBucketNum-1; i++ {
		lastSQBInfo[strconv.FormatFloat(bucketBase, 'f', 3, 32)] = 0
		currentSQBInfo[strconv.FormatFloat(bucketBase, 'f', 3, 32)] = 0
		bucketBase += bucketBase
	}
	lastSQBInfo["+Inf"] = 0
	currentSQBInfo["+Inf"] = 0

	logutil.BgLogger().Info("Telemetry slow query stats initialized", zap.String("currentSQBInfo", currentSQBInfo.String()), zap.String("lastSQBInfo", lastSQBInfo.String()))
}

// postReportSlowQueryStats copy currentSQBInfo to lastSQBInfo to be ready for next report
// this function is designed for being compatible with preview telemetry
func postReportSlowQueryStats() {
	slowQueryLock.Lock()
	lastSQBInfo = currentSQBInfo
	currentSQBInfo = make(SlowQueryBucket)
	slowQueryLock.Unlock()
	logutil.BgLogger().Info("Telemetry slow query stats, postReportSlowQueryStats finished")
}
<commit_msg>telemetry: Block logs when referenced by other components outside tidb (#25330)<commit_after>// Copyright 2021 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package telemetry

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"sync"
	"time"

	pingcapErrors "github.com/pingcap/errors"
	"github.com/pingcap/parser/mysql"
	"github.com/pingcap/tidb/domain/infosync"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/util/logutil"
	pmodel "github.com/prometheus/common/model"
	"go.uber.org/zap"
)

type slowQueryStats struct {
	// Slow Query statistic buckets
	SQBInfo *SlowQueryBucket `json:"slowQueryBucket"`
}

// SlowQueryBucket records the statistic information of slow query buckets
// Buckets:   prometheus.ExponentialBuckets(0.001, 2, 28), // 1ms ~ 1.5days  // defined in metrics/server.go
type SlowQueryBucket map[string]int

func (bucketMap SlowQueryBucket) String() string {
	if bucketMap == nil {
		return "nil"
	}
	var retStr string = "{"
	for k, v := range bucketMap {
		retStr += k + ":" + strconv.Itoa(v) + ","
	}
	retStr = retStr[:len(retStr)-1]
	return retStr
}

const slowQueryBucketNum = 29 //prometheus.ExponentialBuckets(0.001, 2, 28), and 1 more +Inf

var (
	// lastSQBInfo records last statistic information of slow query buckets
	lastSQBInfo SlowQueryBucket
	// currentSQBInfo records current statitic information of slow query buckets
	currentSQBInfo SlowQueryBucket
	slowQueryLock  sync.Mutex
)

func getSlowQueryStats(ctx sessionctx.Context) (*slowQueryStats, error) {
	slowQueryBucket, err := getSlowQueryBucket(ctx)
	if err != nil {
		logutil.BgLogger().Info(err.Error())
		return nil, err
	}

	return &slowQueryStats{slowQueryBucket}, nil
}

// getSlowQueryBucket genenrates the delta SlowQueryBucket to report
func getSlowQueryBucket(ctx sessionctx.Context) (*SlowQueryBucket, error) {
	// update currentSQBInfo first, then gen delta
	if err := updateCurrentSQB(ctx); err != nil {
		return nil, err
	}
	delta := calculateDeltaSQB()
	return delta, nil
}

// updateCurrentSQB records current slow query buckets
func updateCurrentSQB(ctx sessionctx.Context) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = pingcapErrors.Errorf(fmt.Sprintln(r))
		}
	}()

	pQueryCtx, cancel := context.WithTimeout(context.Background(), time.Second*30)
	defer cancel()
	pQueryTs := time.Now().Add(-time.Minute)
	promQL := "tidb_server_slow_query_process_duration_seconds_bucket{sql_type=\"general\"}"
	value, err := querySQLMetric(pQueryCtx, pQueryTs, promQL)

	if err != nil && err != infosync.ErrPrometheusAddrIsNotSet {
		logutil.BgLogger().Info("querySlowQueryMetric got error")
		return err
	}
	if value == nil {
		return
	}
	if value.Type() != pmodel.ValVector {
		return errors.New("Prom vector expected, got " + value.Type().String())
	}
	promVec := value.(pmodel.Vector)
	slowQueryLock.Lock()
	for _, sample := range promVec {
		metric := sample.Metric
		bucketName := metric["le"] //hardcode bucket upper bound
		currentSQBInfo[string(bucketName)] = int(sample.Value)
	}
	slowQueryLock.Unlock()
	return nil
}

// calculateDeltaSQB calculate the delta between current slow query bucket and last slow query bucket
func calculateDeltaSQB() *SlowQueryBucket {
	deltaMap := make(SlowQueryBucket)
	slowQueryLock.Lock()
	for key, value := range currentSQBInfo {
		deltaMap[key] = value - (lastSQBInfo)[key]
	}
	slowQueryLock.Unlock()
	return &deltaMap
}

// init Init lastSQBInfo, follow the definition of metrics/server.go
// Buckets:   prometheus.ExponentialBuckets(0.001, 2, 28), // 1ms ~ 1.5days
func init() {
	lastSQBInfo = make(SlowQueryBucket)
	currentSQBInfo = make(SlowQueryBucket)

	bucketBase := 0.001 // From 0.001 to 134217.728, total 28 float number; the 29th is +Inf
	for i := 0; i < slowQueryBucketNum-1; i++ {
		lastSQBInfo[strconv.FormatFloat(bucketBase, 'f', 3, 32)] = 0
		currentSQBInfo[strconv.FormatFloat(bucketBase, 'f', 3, 32)] = 0
		bucketBase += bucketBase
	}
	lastSQBInfo["+Inf"] = 0
	currentSQBInfo["+Inf"] = 0

	if mysql.TiDBReleaseVersion != "None" {
		logutil.BgLogger().Info("Telemetry slow query stats initialized", zap.String("currentSQBInfo", currentSQBInfo.String()), zap.String("lastSQBInfo", lastSQBInfo.String()))
	}
}

// postReportSlowQueryStats copy currentSQBInfo to lastSQBInfo to be ready for next report
// this function is designed for being compatible with preview telemetry
func postReportSlowQueryStats() {
	slowQueryLock.Lock()
	lastSQBInfo = currentSQBInfo
	currentSQBInfo = make(SlowQueryBucket)
	slowQueryLock.Unlock()
	logutil.BgLogger().Info("Telemetry slow query stats, postReportSlowQueryStats finished")
}
<|endoftext|>


<commit_before>// vSphere template
// The purpose of this is to abstract away all the values in a vSphere Salt Profile to be used
// elsewhere in pepper.

package vsphere

import (
	"os"
	"path/filepath"
	"text/template"

	"github.com/nextgearcapital/pepper/pkg/log"

	"github.com/spf13/viper"
)

// ProfileConfig :
type ProfileConfig struct {
	Platform     string   `mapstructure:"platform"`
	Environment  string   `mapstructure:"environment"`
	InstanceType string   `mapstructure:"instance_type"`
	Role         string   `mapstructure:"role"`
	Datacenter   string   `mapstructure:"datacenter"`
	Provider     string   `mapstructure:"provider"`
	Template     string   `mapstructure:"template"`
	CPU          int      `mapstructure:"cpu"`
	Memory       int      `mapstructure:"memory"`
	DiskSize     int      `mapstructure:"disksize"`
	DHCP         bool     `mapstructure:"dhcp"`
	Network      string   `mapstructure:"network"`
	IPAddress    string   `mapstructure:"ip_address"`
	Gateway      string   `mapstructure:"gateway"`
	Subnet       string   `mapstructure:"subnet"`
	Domain       string   `mapstructure:"domain"`
	DNSServers   []string `mapstructure:"dns_servers"`
	Cluster      string   `mapstructure:"cluster"`
	Folder       string   `mapstructure:"folder"`
	Datastore    string   `mapstructure:"datastore"`
	IsCoreOS     bool     `mapstructure:"is_coreos"`
	ConfigData   string   `mapstructure:"config_data"`
}

var (
	configPath   = "/etc/pepper/config.d/vmware"
	saltProfiles = "/etc/salt/cloud.profiles.d"
)

var (
	// Platform :
	Platform string
	// Environment :
	Environment string
	// InstanceType :
	InstanceType string
	// Template :
	Template string
	// IPAddress :
	IPAddress string
	// Role :
	Role string
	// Datacenter :
	Datacenter string
	// CPU :
	CPU int
	// Memory :
	Memory int
	// DiskSize :
	DiskSize int
)

// Prepare :
func (profile *ProfileConfig) Prepare() error {
	if err := os.MkdirAll(configPath, 0644); err != nil {
		return err
	}

	profile.Platform = Platform
	profile.Environment = Environment
	profile.InstanceType = InstanceType
	profile.Template = Template
	profile.IPAddress = IPAddress
	profile.Role = Role
	profile.Datacenter = Datacenter

	switch InstanceType {
	case "nano":
		profile.CPU = 1
		profile.Memory = 512
		profile.DiskSize = 20
	case "micro":
		profile.CPU = 1
		profile.Memory = 1024
		profile.DiskSize = 20
	case "small":
		profile.CPU = 1
		profile.Memory = 2048
		profile.DiskSize = 40
	case "medium":
		profile.CPU = 2
		profile.Memory = 4096
		profile.DiskSize = 60
	case "large":
		profile.CPU = 2
		profile.Memory = 8192
		profile.DiskSize = 80
	case "xlarge":
		profile.CPU = 4
		profile.Memory = 16384
		profile.DiskSize = 100
	case "ultra":
		profile.CPU = 8
		profile.Memory = 32768
		profile.DiskSize = 160
	case "mega":
		profile.CPU = 16
		profile.Memory = 65536
		profile.DiskSize = 200
	}

	viper.SetConfigName(Environment)
	viper.SetConfigType("yaml")
	viper.AddConfigPath(configPath)

	if err := viper.ReadInConfig(); err != nil {
		log.Err("Couldn't read the config! Did you put it in /etc/pepper/config.d/vmware?")
		return err
	}

	if err := viper.Unmarshal(profile); err != nil {
		return err
	}
	return nil
}

// Generate :
func (profile *ProfileConfig) Generate() error {
	compiled, err := template.New("vsphere_profile").Parse(vsphereTemplate)
	if err != nil {
		return err
	}

	profileName := profile.Platform + "-" + profile.Environment + "-" + profile.InstanceType + ".conf"
	f, err := os.OpenFile(filepath.Join(saltProfiles, profileName), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	if err := compiled.Execute(f, profile); err != nil {
		return err
	}
	return nil
}

const vsphereTemplate = `
{{- .Platform}}-{{.Environment}}-{{.InstanceType}}:
  provider: {{.Provider}}
  clonefrom: {{.Template}}
  num_cpus: {{.CPU}}
  memory: {{.Memory}}
  devices:
    cd:
      CD/DVD drive 1:
        device_type: client_device
        mode: atapi
    disk:
      Hard disk 1:
        size: {{.DiskSize}}
    network:
      Network adapter 1:
        name: {{.Network}}
        adapter_type: vmxnet3
        switch_type: distributed
		{{- if eq .DHCP false}}
        ip: {{.IPAddress}}
        gateway: {{.Gateway}}
        subnet_mask: {{.Subnet}}
	    {{- end}}
    scsi:
      SCSI controller 0:
        type: lsilogic
  {{- if .Domain}}
  domain: {{.Domain}}
  {{- end}}
  {{- if .DNSServers}}
  dns_servers:
  {{- range $value := .DNSServers}}
  - {{$value}}
  {{- end}}
  {{- end}}
  cluster: {{.Cluster}}
  folder: {{.Folder}}
  datastore: {{.Datastore}}
  template: False
  power_on: True
  deploy: False
  minion:
    grains:
	  env: {{.Environment}}
	  {{- if .Datacenter}}
	  datacenter: {{.Datacenter}}
	  {{- end}}
	  {{- if .Role}}
	  role: {{.Role}}
      {{- end}}
  extra_config:
    cpu.hotadd: 'yes'
    mem.hotadd: 'yes'
	{{- if .IsCoreOS}}
	guestinfo.coreos.config.data: {{.ConfigData}}
	guestinfo.coreos.config.data.encoding: base64
	{{- end}}
`
<commit_msg>Stupid tabs<commit_after>// vSphere template
// The purpose of this is to abstract away all the values in a vSphere Salt Profile to be used
// elsewhere in pepper.

package vsphere

import (
	"os"
	"path/filepath"
	"text/template"

	"github.com/nextgearcapital/pepper/pkg/log"

	"github.com/spf13/viper"
)

// ProfileConfig :
type ProfileConfig struct {
	Platform     string   `mapstructure:"platform"`
	Environment  string   `mapstructure:"environment"`
	InstanceType string   `mapstructure:"instance_type"`
	Role         string   `mapstructure:"role"`
	Datacenter   string   `mapstructure:"datacenter"`
	Provider     string   `mapstructure:"provider"`
	Template     string   `mapstructure:"template"`
	CPU          int      `mapstructure:"cpu"`
	Memory       int      `mapstructure:"memory"`
	DiskSize     int      `mapstructure:"disksize"`
	DHCP         bool     `mapstructure:"dhcp"`
	Network      string   `mapstructure:"network"`
	IPAddress    string   `mapstructure:"ip_address"`
	Gateway      string   `mapstructure:"gateway"`
	Subnet       string   `mapstructure:"subnet"`
	Domain       string   `mapstructure:"domain"`
	DNSServers   []string `mapstructure:"dns_servers"`
	Cluster      string   `mapstructure:"cluster"`
	Folder       string   `mapstructure:"folder"`
	Datastore    string   `mapstructure:"datastore"`
	IsCoreOS     bool     `mapstructure:"is_coreos"`
	ConfigData   string   `mapstructure:"config_data"`
}

var (
	configPath   = "/etc/pepper/config.d/vmware"
	saltProfiles = "/etc/salt/cloud.profiles.d"
)

var (
	// Platform :
	Platform string
	// Environment :
	Environment string
	// InstanceType :
	InstanceType string
	// Template :
	Template string
	// IPAddress :
	IPAddress string
	// Role :
	Role string
	// Datacenter :
	Datacenter string
	// CPU :
	CPU int
	// Memory :
	Memory int
	// DiskSize :
	DiskSize int
)

// Prepare :
func (profile *ProfileConfig) Prepare() error {
	if err := os.MkdirAll(configPath, 0644); err != nil {
		return err
	}

	profile.Platform = Platform
	profile.Environment = Environment
	profile.InstanceType = InstanceType
	profile.Template = Template
	profile.IPAddress = IPAddress
	profile.Role = Role
	profile.Datacenter = Datacenter

	switch InstanceType {
	case "nano":
		profile.CPU = 1
		profile.Memory = 512
		profile.DiskSize = 20
	case "micro":
		profile.CPU = 1
		profile.Memory = 1024
		profile.DiskSize = 20
	case "small":
		profile.CPU = 1
		profile.Memory = 2048
		profile.DiskSize = 40
	case "medium":
		profile.CPU = 2
		profile.Memory = 4096
		profile.DiskSize = 60
	case "large":
		profile.CPU = 2
		profile.Memory = 8192
		profile.DiskSize = 80
	case "xlarge":
		profile.CPU = 4
		profile.Memory = 16384
		profile.DiskSize = 100
	case "ultra":
		profile.CPU = 8
		profile.Memory = 32768
		profile.DiskSize = 160
	case "mega":
		profile.CPU = 16
		profile.Memory = 65536
		profile.DiskSize = 200
	}

	viper.SetConfigName(Environment)
	viper.SetConfigType("yaml")
	viper.AddConfigPath(configPath)

	if err := viper.ReadInConfig(); err != nil {
		log.Err("Couldn't read the config! Did you put it in /etc/pepper/config.d/vmware?")
		return err
	}

	if err := viper.Unmarshal(profile); err != nil {
		return err
	}
	return nil
}

// Generate :
func (profile *ProfileConfig) Generate() error {
	compiled, err := template.New("vsphere_profile").Parse(vsphereTemplate)
	if err != nil {
		return err
	}

	profileName := profile.Platform + "-" + profile.Environment + "-" + profile.InstanceType + ".conf"
	f, err := os.OpenFile(filepath.Join(saltProfiles, profileName), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}

	if err := compiled.Execute(f, profile); err != nil {
		return err
	}
	return nil
}

const vsphereTemplate = `
{{- .Platform}}-{{.Environment}}-{{.InstanceType}}:
  provider: {{.Provider}}
  clonefrom: {{.Template}}
  num_cpus: {{.CPU}}
  memory: {{.Memory}}
  devices:
    cd:
      CD/DVD drive 1:
        device_type: client_device
        mode: atapi
    disk:
      Hard disk 1:
        size: {{.DiskSize}}
    network:
      Network adapter 1:
        name: {{.Network}}
        adapter_type: vmxnet3
        switch_type: distributed
		{{- if eq .DHCP false}}
        ip: {{.IPAddress}}
        gateway: {{.Gateway}}
        subnet_mask: {{.Subnet}}
	    {{- end}}
    scsi:
      SCSI controller 0:
        type: lsilogic
  {{- if .Domain}}
  domain: {{.Domain}}
  {{- end}}
  {{- if .DNSServers}}
  dns_servers:
  {{- range $value := .DNSServers}}
  - {{$value}}
  {{- end}}
  {{- end}}
  cluster: {{.Cluster}}
  folder: {{.Folder}}
  datastore: {{.Datastore}}
  template: False
  power_on: True
  deploy: False
  minion:
    grains:
      env: {{.Environment}}
      {{- if .Datacenter}}
      datacenter: {{.Datacenter}}
      {{- end}}
      {{- if .Role}}
      role: {{.Role}}
      {{- end}}
  extra_config:
    cpu.hotadd: 'yes'
    mem.hotadd: 'yes'
	{{- if .IsCoreOS}}
	guestinfo.coreos.config.data: {{.ConfigData}}
	guestinfo.coreos.config.data.encoding: base64
	{{- end}}
`
<|endoftext|>


<commit_before>// Copyright 2012-2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package hook

type listSource struct {
	NoUpdates
	hooks []Info
}

func (q *listSource) Empty() bool {
	return len(q.hooks) == 0
}

func (q *listSource) Next() Info {
	if q.Empty() {
		panic("Source is empty")
	}
	return q.hooks[0]
}

func (q *listSource) Pop() {
	if q.Empty() {
		panic("Source is empty")
	}
	q.hooks = q.hooks[1:]
}

// NewListSource returns a Source that generates only the supplied hooks, in
// order; and which cannot be updated.
func NewListSource(list []Info) Source {
	source := &listSource{hooks: make([]Info, len(list))}
	copy(source.hooks, list)
	return source
}
<commit_msg>Add doc comments to ListSource methods<commit_after>// Copyright 2012-2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package hook

type listSource struct {
	NoUpdates
	hooks []Info
}

// Empty is defined in Source.
func (q *listSource) Empty() bool {
	return len(q.hooks) == 0
}

// Next is defined in Source.
func (q *listSource) Next() Info {
	if q.Empty() {
		panic("Source is empty")
	}
	return q.hooks[0]
}

// Pop is defined in Source.
func (q *listSource) Pop() {
	if q.Empty() {
		panic("Source is empty")
	}
	q.hooks = q.hooks[1:]
}

// NewListSource returns a Source that generates only the supplied hooks, in
// order; and which cannot be updated.
func NewListSource(list []Info) Source {
	source := &listSource{hooks: make([]Info, len(list))}
	copy(source.hooks, list)
	return source
}
<|endoftext|>


<commit_before>package workflows

import (
	"encoding/base64"
	"github.com/stelligent/mu/common"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"io"
	"testing"
)

func TestServiceLoader_FromConfig(t *testing.T) {
	assert := assert.New(t)

	ctx := new(common.Context)
	ctx.Config.Repo.Name = "myrepo"
	ctx.Config.Repo.Slug = "foo/myrepo"
	ctx.Config.Repo.Revision = "1.0.0"
	ctx.Config.Service.Name = "myservice"

	workflow := new(serviceWorkflow)
	err := workflow.serviceLoader(ctx, "2.0.0", "ecr")()
	assert.Nil(err)
	assert.Equal("myservice", workflow.serviceName)
	assert.Equal("2.0.0", workflow.serviceTag)
	assert.Equal(common.ArtifactProviderEcr, workflow.artifactProvider)
}

func TestServiceLoader_FromRepo(t *testing.T) {
	assert := assert.New(t)

	ctx := new(common.Context)
	ctx.Config.Repo.Name = "myrepo"
	ctx.Config.Repo.Slug = "foo/myrepo"
	ctx.Config.Repo.Revision = "1.0.0"

	workflow := new(serviceWorkflow)
	err := workflow.serviceLoader(ctx, "", "ecr")()
	assert.Nil(err)
	assert.Equal("myrepo", workflow.serviceName)
	assert.Equal("1.0.0", workflow.serviceTag)
	assert.Equal(common.ArtifactProviderEcr, workflow.artifactProvider)
}

type mockedRepositoryAuthenticator struct {
	mock.Mock
}

func (m *mockedRepositoryAuthenticator) AuthenticateRepository(repoURL string) (string, error) {
	args := m.Called()
	return args.String(0), args.Error(1)
}

func TestServiceRegistryAuthenticator(t *testing.T) {
	assert := assert.New(t)

	authn := new(mockedRepositoryAuthenticator)
	authn.On("AuthenticateRepository").Return(base64.StdEncoding.EncodeToString([]byte("user:pass")), nil)

	workflow := new(serviceWorkflow)
	err := workflow.serviceRegistryAuthenticator(authn)()

	assert.Nil(err)
	assert.NotNil(workflow.registryAuth)

	authJSON, err := base64.StdEncoding.DecodeString(workflow.registryAuth)
	assert.Nil(err)
	assert.Equal("{\"username\":\"user\", \"password\":\"pass\"}", string(authJSON))

	authn.AssertExpectations(t)
	authn.AssertNumberOfCalls(t, "AuthenticateRepository", 1)
}

type mockedStackManagerForService struct {
	mock.Mock
}

func (m *mockedStackManagerForService) AwaitFinalStatus(stackName string) *common.Stack {
	args := m.Called(stackName)
	stack := args.Get(0)
	if stack == nil {
		return nil
	}
	return stack.(*common.Stack)
}
func (m *mockedStackManagerForService) UpsertStack(stackName string, templateBodyReader io.Reader, stackParameters map[string]string, stackTags map[string]string) error {
	args := m.Called(stackName)
	return args.Error(0)
}
func (m *mockedStackManagerForService) DeleteStack(stackName string) error {
	args := m.Called(stackName)
	return args.Error(0)
}

func TestServiceRepoUpserter(t *testing.T) {
	assert := assert.New(t)

	svc := new(common.Service)

	workflow := new(serviceWorkflow)
	workflow.serviceName = "foo"

	stackManager := new(mockedStackManagerForUpsert)
	stackManager.On("AwaitFinalStatus", "mu-repo-foo").Return(&common.Stack{Status: common.StackStatusCreateComplete})
	stackManager.On("UpsertStack", "mu-repo-foo", mock.AnythingOfType("map[string]string")).Return(nil)

	err := workflow.serviceRepoUpserter(svc, stackManager, stackManager)()
	assert.Nil(err)

	stackManager.AssertExpectations(t)
	stackManager.AssertNumberOfCalls(t, "AwaitFinalStatus", 1)
	stackManager.AssertNumberOfCalls(t, "UpsertStack", 1)
}

func TestEnvironmentTags(t *testing.T) {
	assert := assert.New(t)
	yamlConfig :=
		`
---
environments:
  - name: dev
    tags:
      mytag: first-tag
      foo: bar
`
	config, err := loadYamlConfig(yamlConfig)
	assert.Nil(err)
	assert.Equal(config.Environments[0].Name, "dev")

	joinedMap, err := concatTagMaps(config.Environments[0].Tags, buildEnvironmentTags(config.Environments[0].Name, config.Environments[0].Provider, "StackTypeVpc", "CodeRevision", "RepoName"), EnvironmentTags)
	assert.Nil(err)
	assert.Equal(len(joinedMap), 7)
	assert.NotNil(joinedMap["mytag"])
	assert.Equal(joinedMap["foo"], "bar")
}

func TestNoTagOverride(t *testing.T) {
	assert := assert.New(t)
	yamlConfig :=
		`
---
environments:
  - name: dev
    tags:
      environment: this-should-break
      foo: bar
`

	config, err := loadYamlConfig(yamlConfig)
	assert.Nil(err)
	_, maperr := concatTagMaps(config.Environments[0].Tags, buildEnvironmentTags(config.Environments[0].Name, config.Environments[0].Provider, "StackTypeVpc", "CodeRevision", "RepoName"), EnvironmentTags)

	assert.NotNil(maperr)
}
<commit_msg>update tests<commit_after>package workflows

import (
	"encoding/base64"
	"github.com/stelligent/mu/common"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"io"
	"testing"
)

func TestServiceLoader_FromConfig(t *testing.T) {
	assert := assert.New(t)

	ctx := new(common.Context)
	ctx.Config.Repo.Name = "myrepo"
	ctx.Config.Repo.Slug = "foo/myrepo"
	ctx.Config.Repo.Revision = "1.0.0"
	ctx.Config.Service.Name = "myservice"

	workflow := new(serviceWorkflow)
	err := workflow.serviceLoader(ctx, "2.0.0", "ecr")()
	assert.Nil(err)
	assert.Equal("myservice", workflow.serviceName)
	assert.Equal("2.0.0", workflow.serviceTag)
	assert.Equal(common.ArtifactProviderEcr, workflow.artifactProvider)
}

func TestServiceLoader_FromRepo(t *testing.T) {
	assert := assert.New(t)

	ctx := new(common.Context)
	ctx.Config.Repo.Name = "myrepo"
	ctx.Config.Repo.Slug = "foo/myrepo"
	ctx.Config.Repo.Revision = "1.0.0"

	workflow := new(serviceWorkflow)
	err := workflow.serviceLoader(ctx, "", "ecr")()
	assert.Nil(err)
	assert.Equal("myrepo", workflow.serviceName)
	assert.Equal("1.0.0", workflow.serviceTag)
	assert.Equal(common.ArtifactProviderEcr, workflow.artifactProvider)
}

type mockedRepositoryAuthenticator struct {
	mock.Mock
}

func (m *mockedRepositoryAuthenticator) AuthenticateRepository(repoURL string) (string, error) {
	args := m.Called()
	return args.String(0), args.Error(1)
}

func TestServiceRegistryAuthenticator(t *testing.T) {
	assert := assert.New(t)

	authn := new(mockedRepositoryAuthenticator)
	authn.On("AuthenticateRepository").Return(base64.StdEncoding.EncodeToString([]byte("user:pass")), nil)

	workflow := new(serviceWorkflow)
	err := workflow.serviceRegistryAuthenticator(authn)()

	assert.Nil(err)
	assert.NotNil(workflow.registryAuth)

	authJSON, err := base64.StdEncoding.DecodeString(workflow.registryAuth)
	assert.Nil(err)
	assert.Equal("{\"username\":\"user\", \"password\":\"pass\"}", string(authJSON))

	authn.AssertExpectations(t)
	authn.AssertNumberOfCalls(t, "AuthenticateRepository", 1)
}

type mockedStackManagerForService struct {
	mock.Mock
}

func (m *mockedStackManagerForService) AwaitFinalStatus(stackName string) *common.Stack {
	args := m.Called(stackName)
	stack := args.Get(0)
	if stack == nil {
		return nil
	}
	return stack.(*common.Stack)
}
func (m *mockedStackManagerForService) UpsertStack(stackName string, templateBodyReader io.Reader, stackParameters map[string]string, stackTags map[string]string) error {
	args := m.Called(stackName)
	return args.Error(0)
}
func (m *mockedStackManagerForService) DeleteStack(stackName string) error {
	args := m.Called(stackName)
	return args.Error(0)
}

func TestServiceRepoUpserter(t *testing.T) {
	assert := assert.New(t)

	svc := new(common.Service)

	workflow := new(serviceWorkflow)
	workflow.serviceName = "foo"

	stackManager := new(mockedStackManagerForUpsert)
	stackManager.On("AwaitFinalStatus", "mu-repo-foo").Return(&common.Stack{Status: common.StackStatusCreateComplete})
	stackManager.On("UpsertStack", "mu-repo-foo", mock.AnythingOfType("map[string]string")).Return(nil)

	err := workflow.serviceRepoUpserter(svc, stackManager, stackManager)()
	assert.Nil(err)

	stackManager.AssertExpectations(t)
	stackManager.AssertNumberOfCalls(t, "AwaitFinalStatus", 1)
	stackManager.AssertNumberOfCalls(t, "UpsertStack", 1)
}

func TestEnvironmentTags(t *testing.T) {
	assert := assert.New(t)
	yamlConfig :=
		`
---
environments:
  - name: dev
    tags:
      mytag: first-tag
      foo: bar
`
	config, err := loadYamlConfig(yamlConfig)
	assert.Nil(err)
	assert.Equal(config.Environments[0].Name, "dev")

	var envTags TagInterface = &EnvironmentT{
		Environment: config.Environments[0].Name,
		Type: "StackType",
		Provider: string(config.Environments[0].Provider),
		Revision: "Revision",
		Repo: "Repo",
	}
	joinedMap, err := concatTags(config.Environments[0].Tags, envTags)
	assert.Nil(err)
	assert.Equal(len(joinedMap), 7)
	assert.NotNil(joinedMap["mytag"])
	assert.Equal(joinedMap["foo"], "bar")
}

func TestNoTagOverride(t *testing.T) {
	assert := assert.New(t)
	yamlConfig :=
		`
---
environments:
  - name: dev
    tags:
      environment: this-should-break
      foo: bar
`

	config, err := loadYamlConfig(yamlConfig)
	assert.Nil(err)

	var envTags TagInterface = &EnvironmentT{
		Environment: config.Environments[0].Name,
		Type: "StackType",
		Provider: string(config.Environments[0].Provider),
		Revision: "Revision",
		Repo: "Repo",
	}
	_, maperr := concatTags(config.Environments[0].Tags, envTags)

	assert.NotNil(maperr)
}
<|endoftext|>


<commit_before>package fezzik_test

import (
	"flag"
	"fmt"
	"net/url"
	"runtime"

	"github.com/cloudfoundry-incubator/bbs"
	"github.com/cloudfoundry-incubator/consuladapter"
	"github.com/cloudfoundry-incubator/fezzik"
	"github.com/cloudfoundry-incubator/locket"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/say"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager/lagertest"

	"testing"
	"time"
)

var bbsAddress string
var bbsCA string
var bbsClientCert string
var bbsClientKey string
var consulAddress string

var publiclyAccessibleIP string
var numCells int

var bbsClient bbs.Client
var locketClient locket.Client
var domain, rootFS, guid string
var startTime time.Time

func init() {
	flag.StringVar(&bbsAddress, "bbs-address", "http://10.244.16.130:8889", "http address for the bbs (required)")
	flag.StringVar(&bbsCA, "bbs-ca", "", "bbs ca cert")
	flag.StringVar(&bbsClientCert, "bbs-client-cert", "", "bbs client ssl certificate")
	flag.StringVar(&bbsClientKey, "bbs-client-key", "", "bbs client ssl key")
	flag.StringVar(&consulAddress, "consul-address", "http://127.0.0.1:8500", "http address for the consul agent (required)")
	flag.StringVar(&publiclyAccessibleIP, "publicly-accessible-ip", "10.0.2.2", "a publicly accessible IP for the host the test is running on (necssary to run a local server that containers can phone home to)")
	flag.IntVar(&numCells, "num-cells", 0, "number of cells")
	flag.Parse()
}

func TestFezzik(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Fezzik Suite")
}

var _ = BeforeSuite(func() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error
	bbsClient = initializeBBSClient()

	consulClient, err := consuladapter.NewClient(consulAddress)
	Expect(err).NotTo(HaveOccurred())

	sessionMgr := consuladapter.NewSessionManager(consulClient)
	consulSession, err := consuladapter.NewSession("fezzik", 10*time.Second, consulClient, sessionMgr)
	Expect(err).NotTo(HaveOccurred())

	logger := lagertest.NewTestLogger("fezzik")

	locketClient = locket.NewClient(consulSession, clock.NewClock(), logger)

	domain = "fezzik"
	rootFS = "preloaded:cflinuxfs2"

	if numCells == 0 {
		cells, err := locketClient.Cells()
		Expect(err).NotTo(HaveOccurred())
		numCells = len(cells)
	}

	SetDefaultEventuallyPollingInterval(100 * time.Millisecond)

	say.Println(0, say.Green("Running Fezzik scaled to %d Cells", numCells))
})

var _ = BeforeEach(func() {
	startTime = time.Now()
	guid = fezzik.NewGuid(fmt.Sprintf("%s-%d", domain, GinkgoParallelNode()))
})

var _ = AfterEach(func() {
	endTime := time.Now()
	fmt.Fprint(
		GinkgoWriter,
		say.Cyan(
			"\n%s\nThis test referenced GUID %s\nStart time: %s (%d)\nEnd time: %s (%d)\n",
			CurrentGinkgoTestDescription().FullTestText,
			guid,
			startTime,
			startTime.Unix(),
			endTime,
			endTime.Unix(),
		),
	)
})

func initializeBBSClient() bbs.Client {
	bbsURL, err := url.Parse(bbsAddress)
	Expect(err).NotTo(HaveOccurred())

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(bbsAddress, bbsCA, bbsClientCert, bbsClientKey)
	Expect(err).NotTo(HaveOccurred())
	return bbsClient
}
<commit_msg>Add arguments to new bbs client constructor<commit_after>package fezzik_test

import (
	"flag"
	"fmt"
	"net/url"
	"runtime"

	"github.com/cloudfoundry-incubator/bbs"
	"github.com/cloudfoundry-incubator/consuladapter"
	"github.com/cloudfoundry-incubator/fezzik"
	"github.com/cloudfoundry-incubator/locket"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/say"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager/lagertest"

	"testing"
	"time"
)

var bbsAddress string
var bbsCA string
var bbsClientCert string
var bbsClientKey string
var consulAddress string

var publiclyAccessibleIP string
var numCells int

var bbsClient bbs.Client
var locketClient locket.Client
var domain, rootFS, guid string
var startTime time.Time

func init() {
	flag.StringVar(&bbsAddress, "bbs-address", "http://10.244.16.130:8889", "http address for the bbs (required)")
	flag.StringVar(&bbsCA, "bbs-ca", "", "bbs ca cert")
	flag.StringVar(&bbsClientCert, "bbs-client-cert", "", "bbs client ssl certificate")
	flag.StringVar(&bbsClientKey, "bbs-client-key", "", "bbs client ssl key")
	flag.StringVar(&consulAddress, "consul-address", "http://127.0.0.1:8500", "http address for the consul agent (required)")
	flag.StringVar(&publiclyAccessibleIP, "publicly-accessible-ip", "10.0.2.2", "a publicly accessible IP for the host the test is running on (necssary to run a local server that containers can phone home to)")
	flag.IntVar(&numCells, "num-cells", 0, "number of cells")
	flag.Parse()
}

func TestFezzik(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Fezzik Suite")
}

var _ = BeforeSuite(func() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error
	bbsClient = initializeBBSClient()

	consulClient, err := consuladapter.NewClient(consulAddress)
	Expect(err).NotTo(HaveOccurred())

	sessionMgr := consuladapter.NewSessionManager(consulClient)
	consulSession, err := consuladapter.NewSession("fezzik", 10*time.Second, consulClient, sessionMgr)
	Expect(err).NotTo(HaveOccurred())

	logger := lagertest.NewTestLogger("fezzik")

	locketClient = locket.NewClient(consulSession, clock.NewClock(), logger)

	domain = "fezzik"
	rootFS = "preloaded:cflinuxfs2"

	if numCells == 0 {
		cells, err := locketClient.Cells()
		Expect(err).NotTo(HaveOccurred())
		numCells = len(cells)
	}

	SetDefaultEventuallyPollingInterval(100 * time.Millisecond)

	say.Println(0, say.Green("Running Fezzik scaled to %d Cells", numCells))
})

var _ = BeforeEach(func() {
	startTime = time.Now()
	guid = fezzik.NewGuid(fmt.Sprintf("%s-%d", domain, GinkgoParallelNode()))
})

var _ = AfterEach(func() {
	endTime := time.Now()
	fmt.Fprint(
		GinkgoWriter,
		say.Cyan(
			"\n%s\nThis test referenced GUID %s\nStart time: %s (%d)\nEnd time: %s (%d)\n",
			CurrentGinkgoTestDescription().FullTestText,
			guid,
			startTime,
			startTime.Unix(),
			endTime,
			endTime.Unix(),
		),
	)
})

func initializeBBSClient() bbs.Client {
	bbsURL, err := url.Parse(bbsAddress)
	Expect(err).NotTo(HaveOccurred())

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(bbsAddress, bbsCA, bbsClientCert, bbsClientKey, 0, 0)
	Expect(err).NotTo(HaveOccurred())
	return bbsClient
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"encoding/base64"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/gorilla/mux"
)

type ApplicationConfig struct {
	Hostname           string `toml:"current_host"`
	Host               string
	Port               int
	TokenKey           string `toml:"token_key"`
	MaxConnections     int    `toml:"max_connections"`
	UseAwsHost         bool   `toml:"use_aws_host"`
	SslCertFile        string `toml:"ssl_cert_file"`
	SslKeyFile         string `toml:"ssl_key_file"`
	ClientMinPing      string `toml:"client_min_ping_interval"`
	ClientHelloTimeout string `toml:"client_hello_timeout"`
	PushLongPongs      bool   `toml:"push_long_pongs"`
	Gcm                GCMConfig
}

type GCMConfig struct {
	TTL         int    `toml:"ttl"`
	CollapseKey string `toml:"collapse_key"`
	ProjectId   string `toml:"project_id"`
	DryRun      bool   `toml:"dry_run"`
	ApiKey      string `toml:"api_key"`
	Url         string `toml:"url"`
}

type Application struct {
	hostname           string
	fullHostname       string
	host               string
	port               int
	maxConnnections    int
	clientMinPing      time.Duration
	clientHelloTimeout time.Duration
	pushLongPongs      bool
	tokenKey           []byte
	sslCertFile        string
	sslKeyFile         string
	log                *SimpleLogger
	metrics            *Metrics
	clients            map[string]*Client
	clientMux          *sync.RWMutex
	server             *Serv
	storage            *Storage
	router             *Router
	handlers           *Handler
	gcm                *GCMConfig
}

func (a *Application) ConfigStruct() interface{} {
	defaultHost, _ := os.Hostname()
	return &ApplicationConfig{
		Hostname:           defaultHost,
		Host:               "0.0.0.0",
		Port:               8080,
		MaxConnections:     1000,
		UseAwsHost:         false,
		ClientMinPing:      "20s",
		ClientHelloTimeout: "30s",
		Gcm: GCMConfig{
			TTL:         259200,
			CollapseKey: "simplepush",
			ProjectId:   "simplepush-gcm",
			DryRun:      false,
			Url:         "https://android.googleapis.com/gcm/send",
		},
	}
}

// Fully initialize the application, this initializes all the other components
// as well.
// Note: We implement the Init method to comply with the interface, so the app
// passed here will be nil.
func (a *Application) Init(app *Application, config interface{}) (err error) {
	conf := config.(*ApplicationConfig)

	if conf.UseAwsHost {
		if a.hostname, err = GetAWSPublicHostname(); err != nil {
			return
		}
	} else {
		a.hostname = conf.Hostname
	}

	if len(conf.TokenKey) > 0 {
		if a.tokenKey, err = base64.URLEncoding.DecodeString(conf.TokenKey); err != nil {
			return
		}
	}

	usingSSL := len(conf.SslCertFile) > 0 && len(conf.SslKeyFile) > 0
	if usingSSL && conf.Port == 443 {
		a.fullHostname = fmt.Sprintf("https://%s", a.hostname)
	} else if usingSSL {
		a.fullHostname = fmt.Sprintf("https://%s:%d", a.hostname, conf.Port)
	} else if conf.Port == 80 {
		a.fullHostname = fmt.Sprintf("http://%s", a.hostname)
	} else {
		a.fullHostname = fmt.Sprintf("http://%s:%d", a.hostname, conf.Port)
	}

	a.gcm = &conf.Gcm
	a.host = conf.Host
	a.port = conf.Port
	if a.clientMinPing, err = time.ParseDuration(conf.ClientMinPing); err != nil {
		err = fmt.Errorf("Unable to parse 'client_min_ping_interval: %s",
			err.Error())
		return
	}
	if a.clientHelloTimeout, err = time.ParseDuration(conf.ClientHelloTimeout); err != nil {
		err = fmt.Errorf("Unable to parse 'client_hello_timeout: %s",
			err.Error())
		return
	}
	a.pushLongPongs = conf.PushLongPongs
	a.sslCertFile = conf.SslCertFile
	a.sslKeyFile = conf.SslKeyFile
	a.maxConnnections = conf.MaxConnections
	a.clients = make(map[string]*Client)
	a.clientMux = new(sync.RWMutex)
	return
}

// Set a logger
func (a *Application) SetLogger(logger Logger) (err error) {
	a.log, err = NewLogger(logger)
	return
}

func (a *Application) SetMetrics(metrics *Metrics) error {
	a.metrics = metrics
	return nil
}

func (a *Application) SetStorage(storage *Storage) error {
	a.storage = storage
	return nil
}

func (a *Application) SetRouter(router *Router) error {
	a.router = router
	return nil
}

func (a *Application) SetServer(server *Serv) error {
	a.server = server
	return nil
}

func (a *Application) SetHandlers(handlers *Handler) error {
	a.handlers = handlers
	return nil
}

// Start the application
func (a *Application) Run() (errChan chan error) {
	errChan = make(chan error)

	RESTMux := mux.NewRouter()
	RouteMux := mux.NewRouter()

	RESTMux.HandleFunc("/update/{key}", a.handlers.UpdateHandler)
	RESTMux.HandleFunc("/status/", a.handlers.StatusHandler)
	RESTMux.HandleFunc("/realstatus/", a.handlers.RealStatusHandler)
	RESTMux.HandleFunc("/metrics/", a.handlers.MetricsHandler)
	RESTMux.Handle("/", websocket.Handler(a.handlers.PushSocketHandler))

	RouteMux.HandleFunc("/route/{uaid}", a.handlers.RouteHandler)

	a.log.Info("app", fmt.Sprintf("listening on %s:%d", a.host, a.port), nil)

	// Weigh the anchor!
	go func() {
		addr := a.host + ":" + strconv.Itoa(a.port)
		if len(a.sslCertFile) > 0 && len(a.sslKeyFile) > 0 {
			a.log.Info("app", "Using TLS", nil)
			errChan <- http.ListenAndServeTLS(addr, a.sslCertFile, a.sslKeyFile, RESTMux)
		} else {
			errChan <- http.ListenAndServe(addr, RESTMux)
		}
	}()

	go func() {
		addr := ":" + strconv.Itoa(a.router.port)
		a.log.Info("app", "Starting Router", LogFields{"port": addr})
		errChan <- http.ListenAndServe(addr, RouteMux)
	}()

	return errChan
}

func (a *Application) Hostname() string {
	return a.hostname
}

func (a *Application) FullHostname() string {
	return a.fullHostname
}

func (a *Application) MaxConnections() int {
	return a.maxConnnections
}

func (a *Application) Logger() *SimpleLogger {
	return a.log
}

func (a *Application) Storage() *Storage {
	return a.storage
}

func (a *Application) Metrics() *Metrics {
	return a.metrics
}

func (a *Application) Router() *Router {
	return a.router
}

func (a *Application) Server() *Serv {
	return a.server
}

func (a *Application) TokenKey() []byte {
	return a.tokenKey
}

func (a *Application) ClientCount() (count int) {
	a.clientMux.RLock()
	count = len(a.clients)
	a.clientMux.RUnlock()
	return
}

func (a *Application) ClientExists(uaid string) (collision bool) {
	_, collision = a.GetClient(uaid)
	return
}

func (a *Application) GetClient(uaid string) (client *Client, ok bool) {
	a.clientMux.RLock()
	client, ok = a.clients[uaid]
	a.clientMux.RUnlock()
	return
}

func (a *Application) AddClient(uaid string, client *Client) {
	a.clientMux.Lock()
	a.clients[uaid] = client
	a.clientMux.Unlock()
}

func (a *Application) RemoveClient(uaid string) {
	a.clientMux.Lock()
	delete(a.clients, uaid)
	a.clientMux.Unlock()
}
<commit_msg>Optimization for ClientCount to avoid locking the map just to determine how many clients are connected.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package simplepush

import (
	"encoding/base64"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/gorilla/mux"
)

type ApplicationConfig struct {
	Hostname           string `toml:"current_host"`
	Host               string
	Port               int
	TokenKey           string `toml:"token_key"`
	MaxConnections     int    `toml:"max_connections"`
	UseAwsHost         bool   `toml:"use_aws_host"`
	SslCertFile        string `toml:"ssl_cert_file"`
	SslKeyFile         string `toml:"ssl_key_file"`
	ClientMinPing      string `toml:"client_min_ping_interval"`
	ClientHelloTimeout string `toml:"client_hello_timeout"`
	PushLongPongs      bool   `toml:"push_long_pongs"`
	Gcm                GCMConfig
}

type GCMConfig struct {
	TTL         int    `toml:"ttl"`
	CollapseKey string `toml:"collapse_key"`
	ProjectId   string `toml:"project_id"`
	DryRun      bool   `toml:"dry_run"`
	ApiKey      string `toml:"api_key"`
	Url         string `toml:"url"`
}

type Application struct {
	hostname           string
	fullHostname       string
	host               string
	port               int
	maxConnnections    int
	clientMinPing      time.Duration
	clientHelloTimeout time.Duration
	pushLongPongs      bool
	tokenKey           []byte
	sslCertFile        string
	sslKeyFile         string
	log                *SimpleLogger
	metrics            *Metrics
	clients            map[string]*Client
	clientMux          *sync.RWMutex
	clientCount        *int32
	server             *Serv
	storage            *Storage
	router             *Router
	handlers           *Handler
	gcm                *GCMConfig
}

func (a *Application) ConfigStruct() interface{} {
	defaultHost, _ := os.Hostname()
	return &ApplicationConfig{
		Hostname:           defaultHost,
		Host:               "0.0.0.0",
		Port:               8080,
		MaxConnections:     1000,
		UseAwsHost:         false,
		ClientMinPing:      "20s",
		ClientHelloTimeout: "30s",
		Gcm: GCMConfig{
			TTL:         259200,
			CollapseKey: "simplepush",
			ProjectId:   "simplepush-gcm",
			DryRun:      false,
			Url:         "https://android.googleapis.com/gcm/send",
		},
	}
}

// Fully initialize the application, this initializes all the other components
// as well.
// Note: We implement the Init method to comply with the interface, so the app
// passed here will be nil.
func (a *Application) Init(app *Application, config interface{}) (err error) {
	conf := config.(*ApplicationConfig)

	if conf.UseAwsHost {
		if a.hostname, err = GetAWSPublicHostname(); err != nil {
			return
		}
	} else {
		a.hostname = conf.Hostname
	}

	if len(conf.TokenKey) > 0 {
		if a.tokenKey, err = base64.URLEncoding.DecodeString(conf.TokenKey); err != nil {
			return
		}
	}

	usingSSL := len(conf.SslCertFile) > 0 && len(conf.SslKeyFile) > 0
	if usingSSL && conf.Port == 443 {
		a.fullHostname = fmt.Sprintf("https://%s", a.hostname)
	} else if usingSSL {
		a.fullHostname = fmt.Sprintf("https://%s:%d", a.hostname, conf.Port)
	} else if conf.Port == 80 {
		a.fullHostname = fmt.Sprintf("http://%s", a.hostname)
	} else {
		a.fullHostname = fmt.Sprintf("http://%s:%d", a.hostname, conf.Port)
	}

	a.gcm = &conf.Gcm
	a.host = conf.Host
	a.port = conf.Port
	if a.clientMinPing, err = time.ParseDuration(conf.ClientMinPing); err != nil {
		err = fmt.Errorf("Unable to parse 'client_min_ping_interval: %s",
			err.Error())
		return
	}
	if a.clientHelloTimeout, err = time.ParseDuration(conf.ClientHelloTimeout); err != nil {
		err = fmt.Errorf("Unable to parse 'client_hello_timeout: %s",
			err.Error())
		return
	}
	a.pushLongPongs = conf.PushLongPongs
	a.sslCertFile = conf.SslCertFile
	a.sslKeyFile = conf.SslKeyFile
	a.maxConnnections = conf.MaxConnections
	a.clients = make(map[string]*Client)
	a.clientMux = new(sync.RWMutex)
	count := int32(0)
	a.clientCount = &count
	return
}

// Set a logger
func (a *Application) SetLogger(logger Logger) (err error) {
	a.log, err = NewLogger(logger)
	return
}

func (a *Application) SetMetrics(metrics *Metrics) error {
	a.metrics = metrics
	return nil
}

func (a *Application) SetStorage(storage *Storage) error {
	a.storage = storage
	return nil
}

func (a *Application) SetRouter(router *Router) error {
	a.router = router
	return nil
}

func (a *Application) SetServer(server *Serv) error {
	a.server = server
	return nil
}

func (a *Application) SetHandlers(handlers *Handler) error {
	a.handlers = handlers
	return nil
}

// Start the application
func (a *Application) Run() (errChan chan error) {
	errChan = make(chan error)

	RESTMux := mux.NewRouter()
	RouteMux := mux.NewRouter()

	RESTMux.HandleFunc("/update/{key}", a.handlers.UpdateHandler)
	RESTMux.HandleFunc("/status/", a.handlers.StatusHandler)
	RESTMux.HandleFunc("/realstatus/", a.handlers.RealStatusHandler)
	RESTMux.HandleFunc("/metrics/", a.handlers.MetricsHandler)
	RESTMux.Handle("/", websocket.Handler(a.handlers.PushSocketHandler))

	RouteMux.HandleFunc("/route/{uaid}", a.handlers.RouteHandler)

	a.log.Info("app", fmt.Sprintf("listening on %s:%d", a.host, a.port), nil)

	// Weigh the anchor!
	go func() {
		addr := a.host + ":" + strconv.Itoa(a.port)
		if len(a.sslCertFile) > 0 && len(a.sslKeyFile) > 0 {
			a.log.Info("app", "Using TLS", nil)
			errChan <- http.ListenAndServeTLS(addr, a.sslCertFile, a.sslKeyFile, RESTMux)
		} else {
			errChan <- http.ListenAndServe(addr, RESTMux)
		}
	}()

	go func() {
		addr := ":" + strconv.Itoa(a.router.port)
		a.log.Info("app", "Starting Router", LogFields{"port": addr})
		errChan <- http.ListenAndServe(addr, RouteMux)
	}()

	return errChan
}

func (a *Application) Hostname() string {
	return a.hostname
}

func (a *Application) FullHostname() string {
	return a.fullHostname
}

func (a *Application) MaxConnections() int {
	return a.maxConnnections
}

func (a *Application) Logger() *SimpleLogger {
	return a.log
}

func (a *Application) Storage() *Storage {
	return a.storage
}

func (a *Application) Metrics() *Metrics {
	return a.metrics
}

func (a *Application) Router() *Router {
	return a.router
}

func (a *Application) Server() *Serv {
	return a.server
}

func (a *Application) TokenKey() []byte {
	return a.tokenKey
}

func (a *Application) ClientCount() (count int) {
	return int(atomic.LoadInt32(a.clientCount))
}

func (a *Application) ClientExists(uaid string) (collision bool) {
	_, collision = a.GetClient(uaid)
	return
}

func (a *Application) GetClient(uaid string) (client *Client, ok bool) {
	a.clientMux.RLock()
	client, ok = a.clients[uaid]
	a.clientMux.RUnlock()
	return
}

func (a *Application) AddClient(uaid string, client *Client) {
	a.clientMux.Lock()
	a.clients[uaid] = client
	a.clientMux.Unlock()
	atomic.AddInt32(a.clientCount, 1)
}

func (a *Application) RemoveClient(uaid string) {
	a.clientMux.Lock()
	delete(a.clients, uaid)
	a.clientMux.Unlock()
	atomic.AddInt32(a.clientCount, -1)
}
<|endoftext|>


<commit_before>package module

import (
	"fmt"
	"github.com/briandowns/openweathermap"
	"github.com/davidscholberg/irkbot/lib/configure"
	"github.com/davidscholberg/irkbot/lib/message"
	"os"
	"strings"
)

func HelpWeather() []string {
	s := "weather <location> - display current weather for the given location"
	return []string{s}
}

func Weather(cfg *configure.Config, in *message.InboundMsg, actions *Actions) {
	if !strings.HasPrefix(in.Src, "#") {
		actions.Say("weather searches not allowed in PMs")
		return
	}

	msg := strings.Join(in.MsgArgs[1:], " ")
	//fetch API key from config
	apiKey := cfg.Modules["weather"]["api_key"]

	w, err := openweathermap.NewCurrent("c", "en", apiKey)
	if err != nil {
		actions.Say("error initializing weather search :(")
		fmt.Fprintln(os.Stderr, err)
		return
	}

	err = w.CurrentByName(msg)
	if err != nil {
		actions.Say("no results returned :(")
		fmt.Fprintln(os.Stderr, err)
		return
	}

	conditions := ""
	for i, condition := range w.Weather {
		if i > 0 {
			conditions += ", "
		}
		conditions += condition.Description
	}

	actions.Say(
		fmt.Sprintf(
			"current weather for %s, %s: %.2f°C, %d%% humidity, %s",
			w.Name,
			w.Sys.Country,
			w.Main.Temp,
			w.Main.Humidity,
			conditions,
		),
	)
}
<commit_msg>weather: updated messages to clarify supported search types<commit_after>package module

import (
	"fmt"
	"github.com/briandowns/openweathermap"
	"github.com/davidscholberg/irkbot/lib/configure"
	"github.com/davidscholberg/irkbot/lib/message"
	"os"
	"strings"
)

func HelpWeather() []string {
	s := "weather <location> - display current weather for the given location (only <city> or <city,country> searches are supported)"
	return []string{s}
}

func Weather(cfg *configure.Config, in *message.InboundMsg, actions *Actions) {
	if !strings.HasPrefix(in.Src, "#") {
		actions.Say("weather searches not allowed in PMs")
		return
	}

	msg := strings.Join(in.MsgArgs[1:], " ")
	//fetch API key from config
	apiKey := cfg.Modules["weather"]["api_key"]

	w, err := openweathermap.NewCurrent("c", "en", apiKey)
	if err != nil {
		actions.Say("error initializing weather search :(")
		fmt.Fprintln(os.Stderr, err)
		return
	}

	err = w.CurrentByName(msg)
	if err != nil {
		actions.Say("No results returned. Only <city> or <city,country> searches are supported.")
		fmt.Fprintln(os.Stderr, err)
		return
	}

	conditions := ""
	for i, condition := range w.Weather {
		if i > 0 {
			conditions += ", "
		}
		conditions += condition.Description
	}

	actions.Say(
		fmt.Sprintf(
			"current weather for %s, %s: %.2f°C, %d%% humidity, %s",
			w.Name,
			w.Sys.Country,
			w.Main.Temp,
			w.Main.Humidity,
			conditions,
		),
	)
}
<|endoftext|>