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

package pd

import (
	"context"
	"fmt"
	"reflect"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/failpoint"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/errs"
	"github.com/tikv/pd/pkg/grpcutil"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

// baseClient is a basic client for all other complex client.
type baseClient struct {
	urls      []string
	clusterID uint64
	// PD leader URL
	leader atomic.Value // Store as string
	// PD follower URLs
	followers atomic.Value // Store as []string
	// dc-location -> TSO allocator leader gRPC connection
	clientConns sync.Map // Store as map[string]*grpc.ClientConn
	// dc-location -> TSO allocator leader URL
	allocators sync.Map // Store as map[string]string

	checkLeaderCh        chan struct{}
	checkTSODispatcherCh chan struct{}

	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc

	security SecurityOption

	gRPCDialOptions  []grpc.DialOption
	timeout          time.Duration
	maxRetryTimes    int
	enableForwarding bool
}

// SecurityOption records options about tls
type SecurityOption struct {
	CAPath   string
	CertPath string
	KeyPath  string
}

// ClientOption configures client.
type ClientOption func(c *baseClient)

// WithGRPCDialOptions configures the client with gRPC dial options.
func WithGRPCDialOptions(opts ...grpc.DialOption) ClientOption {
	return func(c *baseClient) {
		c.gRPCDialOptions = append(c.gRPCDialOptions, opts...)
	}
}

// WithCustomTimeoutOption configures the client with timeout option.
func WithCustomTimeoutOption(timeout time.Duration) ClientOption {
	return func(c *baseClient) {
		c.timeout = timeout
	}
}

// WithForwardingOption configures the client with forwarding option.
func WithForwardingOption(enableForwarding bool) ClientOption {
	return func(c *baseClient) {
		c.enableForwarding = enableForwarding
	}
}

// WithMaxErrorRetry configures the client max retry times when connect meets error.
func WithMaxErrorRetry(count int) ClientOption {
	return func(c *baseClient) {
		c.maxRetryTimes = count
	}
}

// newBaseClient returns a new baseClient.
func newBaseClient(ctx context.Context, urls []string, security SecurityOption, opts ...ClientOption) (*baseClient, error) {
	ctx1, cancel := context.WithCancel(ctx)
	c := &baseClient{
		urls:                 urls,
		checkLeaderCh:        make(chan struct{}, 1),
		checkTSODispatcherCh: make(chan struct{}, 1),
		ctx:                  ctx1,
		cancel:               cancel,
		security:             security,
		timeout:              defaultPDTimeout,
		maxRetryTimes:        maxInitClusterRetries,
	}
	for _, opt := range opts {
		opt(c)
	}

	if err := c.initRetry(c.initClusterID); err != nil {
		c.cancel()
		return nil, err
	}
	if err := c.initRetry(c.updateMember); err != nil {
		c.cancel()
		return nil, err
	}
	log.Info("[pd] init cluster id", zap.Uint64("cluster-id", c.clusterID))

	c.wg.Add(1)
	go c.memberLoop()

	return c, nil
}

func (c *baseClient) initRetry(f func() error) error {
	var err error
	for i := 0; i < c.maxRetryTimes; i++ {
		if err = f(); err == nil {
			return nil
		}
		select {
		case <-c.ctx.Done():
			return err
		case <-time.After(time.Second):
		}
	}
	return errors.WithStack(err)
}

func (c *baseClient) memberLoop() {
	defer c.wg.Done()

	ctx, cancel := context.WithCancel(c.ctx)
	defer cancel()

	for {
		select {
		case <-c.checkLeaderCh:
		case <-time.After(time.Minute):
		case <-ctx.Done():
			return
		}
		failpoint.Inject("skipUpdateMember", func() {
			failpoint.Continue()
		})
		if err := c.updateMember(); err != nil {
			log.Error("[pd] failed updateMember", errs.ZapError(err))
		}
	}
}

// ScheduleCheckLeader is used to check leader.
func (c *baseClient) ScheduleCheckLeader() {
	select {
	case c.checkLeaderCh <- struct{}{}:
	default:
	}
}

func (c *baseClient) scheduleCheckTSODispatcher() {
	select {
	case c.checkTSODispatcherCh <- struct{}{}:
	default:
	}
}

// GetClusterID returns the ClusterID.
func (c *baseClient) GetClusterID(context.Context) uint64 {
	return c.clusterID
}

// GetLeaderAddr returns the leader address.
func (c *baseClient) GetLeaderAddr() string {
	leaderAddr := c.leader.Load()
	if leaderAddr == nil {
		return ""
	}
	return leaderAddr.(string)
}

// GetLeaderAddr returns the follower address.
func (c *baseClient) GetFollowerAddr() []string {
	followerAddrs := c.followers.Load()
	if followerAddrs == nil {
		return []string{}
	}
	return followerAddrs.([]string)
}

// GetURLs returns the URLs.
// For testing use. It should only be called when the client is closed.
func (c *baseClient) GetURLs() []string {
	return c.urls
}

func (c *baseClient) GetAllocatorLeaderURLs() map[string]string {
	allocatorLeader := make(map[string]string)
	c.allocators.Range(func(dcLocation, url interface{}) bool {
		allocatorLeader[dcLocation.(string)] = url.(string)
		return true
	})
	return allocatorLeader
}

func (c *baseClient) getAllocatorLeaderAddrByDCLocation(dcLocation string) (string, bool) {
	url, exist := c.allocators.Load(dcLocation)
	if !exist {
		return "", false
	}
	return url.(string), true
}

func (c *baseClient) getAllocatorClientConnByDCLocation(dcLocation string) (*grpc.ClientConn, string) {
	url, ok := c.allocators.Load(dcLocation)
	if !ok {
		panic(fmt.Sprintf("the allocator leader in %s should exist", dcLocation))
	}
	cc, ok := c.clientConns.Load(url)
	if !ok {
		panic(fmt.Sprintf("the client connection of %s in %s should exist", url, dcLocation))
	}
	return cc.(*grpc.ClientConn), url.(string)
}

const globalDCLocation = "global"

func (c *baseClient) gcAllocatorLeaderAddr(curAllocatorMap map[string]*pdpb.Member) {
	// Clean up the old TSO allocators
	c.allocators.Range(func(dcLocationKey, _ interface{}) bool {
		dcLocation := dcLocationKey.(string)
		// Skip the Global TSO Allocator
		if dcLocation == globalDCLocation {
			return true
		}
		if _, exist := curAllocatorMap[dcLocation]; !exist {
			log.Info("[pd] delete unused tso allocator", zap.String("dc-location", dcLocation))
			c.allocators.Delete(dcLocation)
		}
		return true
	})
}

func (c *baseClient) initClusterID() error {
	ctx, cancel := context.WithCancel(c.ctx)
	defer cancel()
	for _, u := range c.urls {
		timeoutCtx, timeoutCancel := context.WithTimeout(ctx, c.timeout)
		members, err := c.getMembers(timeoutCtx, u)
		timeoutCancel()
		if err != nil || members.GetHeader() == nil {
			log.Warn("[pd] failed to get cluster id", zap.String("url", u), errs.ZapError(err))
			continue
		}
		c.clusterID = members.GetHeader().GetClusterId()
		return nil
	}
	return errors.WithStack(errFailInitClusterID)
}

func (c *baseClient) updateMember() error {
	for _, u := range c.urls {
		ctx, cancel := context.WithTimeout(c.ctx, updateMemberTimeout)
		members, err := c.getMembers(ctx, u)
		if err != nil {
			log.Warn("[pd] cannot update member", zap.String("address", u), errs.ZapError(err))
		}
		cancel()
		if err := c.switchTSOAllocatorLeader(members.GetTsoAllocatorLeaders()); err != nil {
			return err
		}
		if err != nil || members.GetLeader() == nil || len(members.GetLeader().GetClientUrls()) == 0 {
			select {
			case <-c.ctx.Done():
				return errors.WithStack(err)
			default:
				continue
			}
		}
		c.updateURLs(members.GetMembers())
		c.updateFollowers(members.GetMembers(), members.GetLeader())
		if err := c.switchLeader(members.GetLeader().GetClientUrls()); err != nil {
			return err
		}
		c.scheduleCheckTSODispatcher()
		return nil
	}
	return errs.ErrClientGetLeader.FastGenByArgs(c.urls)
}

func (c *baseClient) getMembers(ctx context.Context, url string) (*pdpb.GetMembersResponse, error) {
	cc, err := c.getOrCreateGRPCConn(url)
	if err != nil {
		return nil, err
	}
	members, err := pdpb.NewPDClient(cc).GetMembers(ctx, &pdpb.GetMembersRequest{})
	if err != nil {
		attachErr := errors.Errorf("error:%s target:%s status:%s", err, cc.Target(), cc.GetState().String())
		return nil, errs.ErrClientGetMember.Wrap(attachErr).GenWithStackByCause()
	}
	return members, nil
}

func (c *baseClient) updateURLs(members []*pdpb.Member) {
	urls := make([]string, 0, len(members))
	for _, m := range members {
		urls = append(urls, m.GetClientUrls()...)
	}

	sort.Strings(urls)
	// the url list is same.
	if reflect.DeepEqual(c.urls, urls) {
		return
	}

	log.Info("[pd] update member urls", zap.Strings("old-urls", c.urls), zap.Strings("new-urls", urls))
	c.urls = urls
}

func (c *baseClient) switchLeader(addrs []string) error {
	// FIXME: How to safely compare leader urls? For now, only allows one client url.
	addr := addrs[0]
	oldLeader := c.GetLeaderAddr()
	if addr == oldLeader {
		return nil
	}

	log.Info("[pd] switch leader", zap.String("new-leader", addr), zap.String("old-leader", oldLeader))
	if _, err := c.getOrCreateGRPCConn(addr); err != nil {
		return err
	}
	// Set PD leader and Global TSO Allocator (which is also the PD leader)
	c.leader.Store(addr)
	c.allocators.Store(globalDCLocation, addr)
	return nil
}

func (c *baseClient) updateFollowers(members []*pdpb.Member, leader *pdpb.Member) {
	var addrs []string
	for _, member := range members {
		if member.GetMemberId() != leader.GetMemberId() {
			if len(member.GetClientUrls()) > 0 {
				addrs = append(addrs, member.GetClientUrls()...)
			}
		}
	}
	c.followers.Store(addrs)
}

func (c *baseClient) switchTSOAllocatorLeader(allocatorMap map[string]*pdpb.Member) error {
	if len(allocatorMap) == 0 {
		return nil
	}
	// Switch to the new one
	for dcLocation, member := range allocatorMap {
		if len(member.GetClientUrls()) == 0 {
			continue
		}
		addr := member.GetClientUrls()[0]
		oldAddr, exist := c.getAllocatorLeaderAddrByDCLocation(dcLocation)
		if exist && addr == oldAddr {
			continue
		}
		log.Info("[pd] switch dc tso allocator leader",
			zap.String("dc-location", dcLocation),
			zap.String("new-leader", addr),
			zap.String("old-leader", oldAddr))
		if _, err := c.getOrCreateGRPCConn(addr); err != nil {
			return err
		}
		c.allocators.Store(dcLocation, addr)
	}
	// Garbage collection of the old TSO allocator leaders
	c.gcAllocatorLeaderAddr(allocatorMap)
	return nil
}

func (c *baseClient) getOrCreateGRPCConn(addr string) (*grpc.ClientConn, error) {
	conn, ok := c.clientConns.Load(addr)
	if ok {
		return conn.(*grpc.ClientConn), nil
	}
	tlsCfg, err := grpcutil.TLSConfig{
		CAPath:   c.security.CAPath,
		CertPath: c.security.CertPath,
		KeyPath:  c.security.KeyPath,
	}.ToTLSConfig()
	if err != nil {
		return nil, err
	}
	dCtx, cancel := context.WithTimeout(c.ctx, dialTimeout)
	defer cancel()
	cc, err := grpcutil.GetClientConn(dCtx, addr, tlsCfg, c.gRPCDialOptions...)
	if err != nil {
		return nil, err
	}
	if old, ok := c.clientConns.Load(addr); ok {
		cc.Close()
		log.Debug("use old connection", zap.String("target", cc.Target()), zap.String("state", cc.GetState().String()))
		return old.(*grpc.ClientConn), nil
	}
	c.clientConns.Store(addr, cc)
	return cc, nil
}
<commit_msg>client: adjust the log level in updateMember (#3630)<commit_after>// Copyright 2019 TiKV Project 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package pd

import (
	"context"
	"fmt"
	"reflect"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/failpoint"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/log"
	"github.com/tikv/pd/pkg/errs"
	"github.com/tikv/pd/pkg/grpcutil"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

// baseClient is a basic client for all other complex client.
type baseClient struct {
	urls      []string
	clusterID uint64
	// PD leader URL
	leader atomic.Value // Store as string
	// PD follower URLs
	followers atomic.Value // Store as []string
	// dc-location -> TSO allocator leader gRPC connection
	clientConns sync.Map // Store as map[string]*grpc.ClientConn
	// dc-location -> TSO allocator leader URL
	allocators sync.Map // Store as map[string]string

	checkLeaderCh        chan struct{}
	checkTSODispatcherCh chan struct{}

	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc

	security SecurityOption

	gRPCDialOptions  []grpc.DialOption
	timeout          time.Duration
	maxRetryTimes    int
	enableForwarding bool
}

// SecurityOption records options about tls
type SecurityOption struct {
	CAPath   string
	CertPath string
	KeyPath  string
}

// ClientOption configures client.
type ClientOption func(c *baseClient)

// WithGRPCDialOptions configures the client with gRPC dial options.
func WithGRPCDialOptions(opts ...grpc.DialOption) ClientOption {
	return func(c *baseClient) {
		c.gRPCDialOptions = append(c.gRPCDialOptions, opts...)
	}
}

// WithCustomTimeoutOption configures the client with timeout option.
func WithCustomTimeoutOption(timeout time.Duration) ClientOption {
	return func(c *baseClient) {
		c.timeout = timeout
	}
}

// WithForwardingOption configures the client with forwarding option.
func WithForwardingOption(enableForwarding bool) ClientOption {
	return func(c *baseClient) {
		c.enableForwarding = enableForwarding
	}
}

// WithMaxErrorRetry configures the client max retry times when connect meets error.
func WithMaxErrorRetry(count int) ClientOption {
	return func(c *baseClient) {
		c.maxRetryTimes = count
	}
}

// newBaseClient returns a new baseClient.
func newBaseClient(ctx context.Context, urls []string, security SecurityOption, opts ...ClientOption) (*baseClient, error) {
	ctx1, cancel := context.WithCancel(ctx)
	c := &baseClient{
		urls:                 urls,
		checkLeaderCh:        make(chan struct{}, 1),
		checkTSODispatcherCh: make(chan struct{}, 1),
		ctx:                  ctx1,
		cancel:               cancel,
		security:             security,
		timeout:              defaultPDTimeout,
		maxRetryTimes:        maxInitClusterRetries,
	}
	for _, opt := range opts {
		opt(c)
	}

	if err := c.initRetry(c.initClusterID); err != nil {
		c.cancel()
		return nil, err
	}
	if err := c.initRetry(c.updateMember); err != nil {
		c.cancel()
		return nil, err
	}
	log.Info("[pd] init cluster id", zap.Uint64("cluster-id", c.clusterID))

	c.wg.Add(1)
	go c.memberLoop()

	return c, nil
}

func (c *baseClient) initRetry(f func() error) error {
	var err error
	for i := 0; i < c.maxRetryTimes; i++ {
		if err = f(); err == nil {
			return nil
		}
		select {
		case <-c.ctx.Done():
			return err
		case <-time.After(time.Second):
		}
	}
	return errors.WithStack(err)
}

func (c *baseClient) memberLoop() {
	defer c.wg.Done()

	ctx, cancel := context.WithCancel(c.ctx)
	defer cancel()

	for {
		select {
		case <-c.checkLeaderCh:
		case <-time.After(time.Minute):
		case <-ctx.Done():
			return
		}
		failpoint.Inject("skipUpdateMember", func() {
			failpoint.Continue()
		})
		if err := c.updateMember(); err != nil {
			log.Error("[pd] failed updateMember", errs.ZapError(err))
		}
	}
}

// ScheduleCheckLeader is used to check leader.
func (c *baseClient) ScheduleCheckLeader() {
	select {
	case c.checkLeaderCh <- struct{}{}:
	default:
	}
}

func (c *baseClient) scheduleCheckTSODispatcher() {
	select {
	case c.checkTSODispatcherCh <- struct{}{}:
	default:
	}
}

// GetClusterID returns the ClusterID.
func (c *baseClient) GetClusterID(context.Context) uint64 {
	return c.clusterID
}

// GetLeaderAddr returns the leader address.
func (c *baseClient) GetLeaderAddr() string {
	leaderAddr := c.leader.Load()
	if leaderAddr == nil {
		return ""
	}
	return leaderAddr.(string)
}

// GetLeaderAddr returns the follower address.
func (c *baseClient) GetFollowerAddr() []string {
	followerAddrs := c.followers.Load()
	if followerAddrs == nil {
		return []string{}
	}
	return followerAddrs.([]string)
}

// GetURLs returns the URLs.
// For testing use. It should only be called when the client is closed.
func (c *baseClient) GetURLs() []string {
	return c.urls
}

func (c *baseClient) GetAllocatorLeaderURLs() map[string]string {
	allocatorLeader := make(map[string]string)
	c.allocators.Range(func(dcLocation, url interface{}) bool {
		allocatorLeader[dcLocation.(string)] = url.(string)
		return true
	})
	return allocatorLeader
}

func (c *baseClient) getAllocatorLeaderAddrByDCLocation(dcLocation string) (string, bool) {
	url, exist := c.allocators.Load(dcLocation)
	if !exist {
		return "", false
	}
	return url.(string), true
}

func (c *baseClient) getAllocatorClientConnByDCLocation(dcLocation string) (*grpc.ClientConn, string) {
	url, ok := c.allocators.Load(dcLocation)
	if !ok {
		panic(fmt.Sprintf("the allocator leader in %s should exist", dcLocation))
	}
	cc, ok := c.clientConns.Load(url)
	if !ok {
		panic(fmt.Sprintf("the client connection of %s in %s should exist", url, dcLocation))
	}
	return cc.(*grpc.ClientConn), url.(string)
}

const globalDCLocation = "global"

func (c *baseClient) gcAllocatorLeaderAddr(curAllocatorMap map[string]*pdpb.Member) {
	// Clean up the old TSO allocators
	c.allocators.Range(func(dcLocationKey, _ interface{}) bool {
		dcLocation := dcLocationKey.(string)
		// Skip the Global TSO Allocator
		if dcLocation == globalDCLocation {
			return true
		}
		if _, exist := curAllocatorMap[dcLocation]; !exist {
			log.Info("[pd] delete unused tso allocator", zap.String("dc-location", dcLocation))
			c.allocators.Delete(dcLocation)
		}
		return true
	})
}

func (c *baseClient) initClusterID() error {
	ctx, cancel := context.WithCancel(c.ctx)
	defer cancel()
	for _, u := range c.urls {
		timeoutCtx, timeoutCancel := context.WithTimeout(ctx, c.timeout)
		members, err := c.getMembers(timeoutCtx, u)
		timeoutCancel()
		if err != nil || members.GetHeader() == nil {
			log.Warn("[pd] failed to get cluster id", zap.String("url", u), errs.ZapError(err))
			continue
		}
		c.clusterID = members.GetHeader().GetClusterId()
		return nil
	}
	return errors.WithStack(errFailInitClusterID)
}

func (c *baseClient) updateMember() error {
	for _, u := range c.urls {
		ctx, cancel := context.WithTimeout(c.ctx, updateMemberTimeout)
		members, err := c.getMembers(ctx, u)
		cancel()

		var errTSO error
		if err == nil {
			if members.GetLeader() == nil || len(members.GetLeader().GetClientUrls()) == 0 {
				err = errs.ErrClientGetLeader.FastGenByArgs("leader address don't exist")
			}
			// Still need to update TsoAllocatorLeadersEven, even if there is no PD leader
			errTSO = c.switchTSOAllocatorLeader(members.GetTsoAllocatorLeaders())
		}

		// Failed to get PD leader
		if err != nil {
			log.Info("[pd] cannot update member from this address",
				zap.String("address", u),
				errs.ZapError(err))
			select {
			case <-c.ctx.Done():
				return errors.WithStack(err)
			default:
				continue
			}
		}

		c.updateURLs(members.GetMembers())
		c.updateFollowers(members.GetMembers(), members.GetLeader())
		if err := c.switchLeader(members.GetLeader().GetClientUrls()); err != nil {
			return err
		}
		c.scheduleCheckTSODispatcher()

		// If `switchLeader` succeeds but `switchTSOAllocatorLeader` has an error,
		// the error of `switchTSOAllocatorLeader` will be returned.
		return errTSO
	}
	return errs.ErrClientGetLeader.FastGenByArgs(c.urls)
}

func (c *baseClient) getMembers(ctx context.Context, url string) (*pdpb.GetMembersResponse, error) {
	cc, err := c.getOrCreateGRPCConn(url)
	if err != nil {
		return nil, err
	}
	members, err := pdpb.NewPDClient(cc).GetMembers(ctx, &pdpb.GetMembersRequest{})
	if err != nil {
		attachErr := errors.Errorf("error:%s target:%s status:%s", err, cc.Target(), cc.GetState().String())
		return nil, errs.ErrClientGetMember.Wrap(attachErr).GenWithStackByCause()
	}
	return members, nil
}

func (c *baseClient) updateURLs(members []*pdpb.Member) {
	urls := make([]string, 0, len(members))
	for _, m := range members {
		urls = append(urls, m.GetClientUrls()...)
	}

	sort.Strings(urls)
	// the url list is same.
	if reflect.DeepEqual(c.urls, urls) {
		return
	}

	log.Info("[pd] update member urls", zap.Strings("old-urls", c.urls), zap.Strings("new-urls", urls))
	c.urls = urls
}

func (c *baseClient) switchLeader(addrs []string) error {
	// FIXME: How to safely compare leader urls? For now, only allows one client url.
	addr := addrs[0]
	oldLeader := c.GetLeaderAddr()
	if addr == oldLeader {
		return nil
	}

	if _, err := c.getOrCreateGRPCConn(addr); err != nil {
		log.Warn("[pd] failed to connect leader", zap.String("leader", addr), errs.ZapError(err))
		return err
	}
	// Set PD leader and Global TSO Allocator (which is also the PD leader)
	c.leader.Store(addr)
	c.allocators.Store(globalDCLocation, addr)
	log.Info("[pd] switch leader", zap.String("new-leader", addr), zap.String("old-leader", oldLeader))
	return nil
}

func (c *baseClient) updateFollowers(members []*pdpb.Member, leader *pdpb.Member) {
	var addrs []string
	for _, member := range members {
		if member.GetMemberId() != leader.GetMemberId() {
			if len(member.GetClientUrls()) > 0 {
				addrs = append(addrs, member.GetClientUrls()...)
			}
		}
	}
	c.followers.Store(addrs)
}

func (c *baseClient) switchTSOAllocatorLeader(allocatorMap map[string]*pdpb.Member) error {
	if len(allocatorMap) == 0 {
		return nil
	}
	// Switch to the new one
	for dcLocation, member := range allocatorMap {
		if len(member.GetClientUrls()) == 0 {
			continue
		}
		addr := member.GetClientUrls()[0]
		oldAddr, exist := c.getAllocatorLeaderAddrByDCLocation(dcLocation)
		if exist && addr == oldAddr {
			continue
		}
		if _, err := c.getOrCreateGRPCConn(addr); err != nil {
			log.Warn("[pd] failed to connect dc tso allocator leader",
				zap.String("dc-location", dcLocation),
				zap.String("leader", addr),
				errs.ZapError(err))
			return err
		}
		c.allocators.Store(dcLocation, addr)
		log.Info("[pd] switch dc tso allocator leader",
			zap.String("dc-location", dcLocation),
			zap.String("new-leader", addr),
			zap.String("old-leader", oldAddr))
	}
	// Garbage collection of the old TSO allocator leaders
	c.gcAllocatorLeaderAddr(allocatorMap)
	return nil
}

func (c *baseClient) getOrCreateGRPCConn(addr string) (*grpc.ClientConn, error) {
	conn, ok := c.clientConns.Load(addr)
	if ok {
		return conn.(*grpc.ClientConn), nil
	}
	tlsCfg, err := grpcutil.TLSConfig{
		CAPath:   c.security.CAPath,
		CertPath: c.security.CertPath,
		KeyPath:  c.security.KeyPath,
	}.ToTLSConfig()
	if err != nil {
		return nil, err
	}
	dCtx, cancel := context.WithTimeout(c.ctx, dialTimeout)
	defer cancel()
	cc, err := grpcutil.GetClientConn(dCtx, addr, tlsCfg, c.gRPCDialOptions...)
	if err != nil {
		return nil, err
	}
	if old, ok := c.clientConns.Load(addr); ok {
		cc.Close()
		log.Debug("use old connection", zap.String("target", cc.Target()), zap.String("state", cc.GetState().String()))
		return old.(*grpc.ClientConn), nil
	}
	c.clientConns.Store(addr, cc)
	return cc, nil
}
<|endoftext|>


<commit_before>package announcer

import (
	"math"
	"net"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/tracker"
)

type Status int

const (
	NotContactedYet Status = iota
	Contacting
	Working
	NotWorking
)

type PeriodicalAnnouncer struct {
	Tracker        tracker.Tracker
	status         Status
	statsCommandC  chan statsRequest
	numWant        int
	interval       time.Duration
	minInterval    time.Duration
	seeders        int
	leechers       int
	lastError      error
	log            logger.Logger
	completedC     chan struct{}
	newPeers       chan []*net.TCPAddr
	backoff        backoff.BackOff
	requests       chan *Request
	lastAnnounce   time.Time
	HasAnnounced   bool
	needMorePeersC chan bool
	closeC         chan struct{}
	doneC          chan struct{}
}

type Request struct {
	Response chan Response
	Cancel   chan struct{}
}

type Response struct {
	Torrent tracker.Torrent
}

func NewPeriodicalAnnouncer(trk tracker.Tracker, numWant int, minInterval time.Duration, requests chan *Request, completedC chan struct{}, newPeers chan []*net.TCPAddr, l logger.Logger) *PeriodicalAnnouncer {
	return &PeriodicalAnnouncer{
		Tracker:        trk,
		status:         NotContactedYet,
		statsCommandC:  make(chan statsRequest),
		numWant:        numWant,
		minInterval:    minInterval,
		log:            l,
		completedC:     completedC,
		newPeers:       newPeers,
		requests:       requests,
		needMorePeersC: make(chan bool),
		closeC:         make(chan struct{}),
		doneC:          make(chan struct{}),
		backoff: &backoff.ExponentialBackOff{
			InitialInterval:     5 * time.Second,
			RandomizationFactor: 0.5,
			Multiplier:          2,
			MaxInterval:         30 * time.Minute,
			MaxElapsedTime:      0, // never stop
			Clock:               backoff.SystemClock,
		},
	}
}

func (a *PeriodicalAnnouncer) Close() {
	close(a.closeC)
	<-a.doneC
}

type statsRequest struct {
	Response chan Stats
}

func (a *PeriodicalAnnouncer) Stats() Stats {
	var stats Stats
	req := statsRequest{Response: make(chan Stats, 1)}
	select {
	case a.statsCommandC <- req:
	case <-a.closeC:
	}
	select {
	case stats = <-req.Response:
	case <-a.closeC:
	}
	return stats
}

func (a *PeriodicalAnnouncer) NeedMorePeers(val bool) {
	select {
	case a.needMorePeersC <- val:
	case <-a.doneC:
	}
}

func (a *PeriodicalAnnouncer) Run() {
	defer close(a.doneC)
	a.backoff.Reset()

	timer := time.NewTimer(math.MaxInt64)
	defer timer.Stop()

	var needMorePeers bool

	ca := newCancelableAnnouncer(a.Tracker, a.requests, a.newPeers)
	defer ca.Cancel()

	ca.Announce(tracker.EventStarted, a.numWant)
	for {
		select {
		case <-timer.C:
			a.status = Contacting
			ca.Announce(tracker.EventNone, a.numWant)
		case resp := <-ca.ResponseC:
			ca.announcing = false
			a.lastAnnounce = time.Now()
			a.seeders = int(resp.Seeders)
			a.leechers = int(resp.Leechers)
			a.interval = resp.Interval
			if resp.MinInterval > 0 {
				a.minInterval = resp.MinInterval
			}
			a.HasAnnounced = true
			a.lastError = nil
			a.status = Working
			a.backoff.Reset()
			if needMorePeers {
				timer.Reset(a.minInterval)
			} else {
				timer.Reset(a.interval)
			}
		case a.lastError = <-ca.ErrorC:
			ca.announcing = false
			a.status = NotWorking
			a.log.Debugln("announce error:", a.lastError)
			timer.Reset(a.backoff.NextBackOff())
		case needMorePeers = <-a.needMorePeersC:
			if ca.announcing {
				break
			}
			if needMorePeers {
				timer.Reset(time.Until(a.lastAnnounce.Add(a.minInterval)))
			} else {
				timer.Reset(time.Until(a.lastAnnounce.Add(a.interval)))
			}
		case <-a.completedC:
			ca.Cancel()
			a.status = Contacting
			ca.Announce(tracker.EventCompleted, 0)
			a.completedC = nil
		case req := <-a.statsCommandC:
			req.Response <- a.stats()
		case <-a.closeC:
			return
		}
	}
}

type Stats struct {
	Status   Status
	Error    error
	Seeders  int
	Leechers int
}

func (a *PeriodicalAnnouncer) stats() Stats {
	return Stats{
		Status:   a.status,
		Error:    a.lastError,
		Seeders:  a.seeders,
		Leechers: a.leechers,
	}
}
<commit_msg>do not get stuck on NeedMorePeers<commit_after>package announcer

import (
	"math"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/tracker"
)

type Status int

const (
	NotContactedYet Status = iota
	Contacting
	Working
	NotWorking
)

type PeriodicalAnnouncer struct {
	Tracker       tracker.Tracker
	status        Status
	statsCommandC chan statsRequest
	numWant       int
	interval      time.Duration
	minInterval   time.Duration
	seeders       int
	leechers      int
	lastError     error
	log           logger.Logger
	completedC    chan struct{}
	newPeers      chan []*net.TCPAddr
	backoff       backoff.BackOff
	requests      chan *Request
	lastAnnounce  time.Time
	HasAnnounced  bool
	closeC        chan struct{}
	doneC         chan struct{}

	needMorePeers  bool
	mNeedMorePeers sync.RWMutex
	needMorePeersC chan struct{}
}

type Request struct {
	Response chan Response
	Cancel   chan struct{}
}

type Response struct {
	Torrent tracker.Torrent
}

func NewPeriodicalAnnouncer(trk tracker.Tracker, numWant int, minInterval time.Duration, requests chan *Request, completedC chan struct{}, newPeers chan []*net.TCPAddr, l logger.Logger) *PeriodicalAnnouncer {
	return &PeriodicalAnnouncer{
		Tracker:        trk,
		status:         NotContactedYet,
		statsCommandC:  make(chan statsRequest),
		numWant:        numWant,
		minInterval:    minInterval,
		log:            l,
		completedC:     completedC,
		newPeers:       newPeers,
		requests:       requests,
		needMorePeersC: make(chan struct{}, 1),
		closeC:         make(chan struct{}),
		doneC:          make(chan struct{}),
		backoff: &backoff.ExponentialBackOff{
			InitialInterval:     5 * time.Second,
			RandomizationFactor: 0.5,
			Multiplier:          2,
			MaxInterval:         30 * time.Minute,
			MaxElapsedTime:      0, // never stop
			Clock:               backoff.SystemClock,
		},
	}
}

func (a *PeriodicalAnnouncer) Close() {
	close(a.closeC)
	<-a.doneC
}

type statsRequest struct {
	Response chan Stats
}

func (a *PeriodicalAnnouncer) Stats() Stats {
	var stats Stats
	req := statsRequest{Response: make(chan Stats, 1)}
	select {
	case a.statsCommandC <- req:
	case <-a.closeC:
	}
	select {
	case stats = <-req.Response:
	case <-a.closeC:
	}
	return stats
}

func (a *PeriodicalAnnouncer) NeedMorePeers(val bool) {
	a.mNeedMorePeers.Lock()
	a.needMorePeers = val
	a.mNeedMorePeers.Unlock()
	select {
	case a.needMorePeersC <- struct{}{}:
	case <-a.doneC:
	default:
	}
}

func (a *PeriodicalAnnouncer) Run() {
	defer close(a.doneC)
	a.backoff.Reset()

	timer := time.NewTimer(math.MaxInt64)
	defer timer.Stop()

	ca := newCancelableAnnouncer(a.Tracker, a.requests, a.newPeers)
	defer ca.Cancel()

	ca.Announce(tracker.EventStarted, a.numWant)
	for {
		select {
		case <-timer.C:
			a.status = Contacting
			ca.Announce(tracker.EventNone, a.numWant)
		case resp := <-ca.ResponseC:
			ca.announcing = false
			a.lastAnnounce = time.Now()
			a.seeders = int(resp.Seeders)
			a.leechers = int(resp.Leechers)
			a.interval = resp.Interval
			if resp.MinInterval > 0 {
				a.minInterval = resp.MinInterval
			}
			a.HasAnnounced = true
			a.lastError = nil
			a.status = Working
			a.backoff.Reset()
			a.mNeedMorePeers.RLock()
			needMorePeers := a.needMorePeers
			a.mNeedMorePeers.RUnlock()
			if needMorePeers {
				timer.Reset(a.minInterval)
			} else {
				timer.Reset(a.interval)
			}
		case a.lastError = <-ca.ErrorC:
			ca.announcing = false
			a.lastAnnounce = time.Now()
			a.status = NotWorking
			a.log.Debugln("announce error:", a.lastError)
			timer.Reset(a.backoff.NextBackOff())
		case <-a.needMorePeersC:
			if ca.announcing {
				break
			}
			a.mNeedMorePeers.RLock()
			needMorePeers := a.needMorePeers
			a.mNeedMorePeers.RUnlock()
			if needMorePeers {
				timer.Reset(time.Until(a.lastAnnounce.Add(a.minInterval)))
			} else {
				timer.Reset(time.Until(a.lastAnnounce.Add(a.interval)))
			}
		case <-a.completedC:
			ca.Cancel()
			a.status = Contacting
			ca.Announce(tracker.EventCompleted, 0)
			a.completedC = nil
		case req := <-a.statsCommandC:
			req.Response <- a.stats()
		case <-a.closeC:
			return
		}
	}
}

type Stats struct {
	Status   Status
	Error    error
	Seeders  int
	Leechers int
}

func (a *PeriodicalAnnouncer) stats() Stats {
	return Stats{
		Status:   a.status,
		Error:    a.lastError,
		Seeders:  a.seeders,
		Leechers: a.leechers,
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 oglematchers_test

import (
	"errors"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

type ErrorTest struct {
	matcherCalled bool
	suppliedCandidate interface{}
	wrappedResult bool
	wrappedError error

	matcher Matcher
}

func init() { RegisterTestSuite(&ErrorTest{}) }

func (t *ErrorTest) SetUp() {
	wrapped := &fakeMatcher{
		func(c interface{}) (bool, error) {
			t.matcherCalled = true
			t.suppliedCandidate = c
			return t.wrappedResult, t.wrappedError
		},
		"is foo",
	}

	t.matcher = Error(wrapped)
}

func isFatal(err error) bool {
	_, isFatal := err.(*FatalError)
	return isFatal
}

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func (t *ErrorTest) Description() {
	ExpectThat(t.matcher.Description(), Equals("error is foo"))
}

func (t *ErrorTest) CandidateIsNil() {
	res, err := t.matcher.Matches(nil)

	ExpectThat(t.matcherCalled, Equals(false))
	ExpectThat(res, Equals(false))
	ExpectThat(err.Error(), Equals("which is not an error"))
	ExpectTrue(isFatal(err))
}

func (t *ErrorTest) CandidateIsString() {
	res, err := t.matcher.Matches("taco")

	ExpectThat(t.matcherCalled, Equals(false))
	ExpectThat(res, Equals(false))
	ExpectThat(err.Error(), Equals("which is not an error"))
	ExpectTrue(isFatal(err))
}

func (t *ErrorTest) CallsWrappedMatcher() {
	candidate := errors.New("taco")
	t.matcher.Matches(candidate)

	ExpectThat(t.matcherCalled, Equals(true))
	ExpectThat(t.suppliedCandidate, Equals("taco"))
}

func (t *ErrorTest) ReturnsWrappedMatcherResult() {
	t.wrappedResult = true
	t.wrappedError = errors.New("burrito")

	res, err := t.matcher.Matches(errors.New(""))

	ExpectThat(res, Equals(true))
	ExpectThat(err, Equals(t.wrappedError))
}
<commit_msg>Fixed a runtime panic.<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 oglematchers_test

import (
	"errors"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////

type ErrorTest struct {
	matcherCalled bool
	suppliedCandidate interface{}
	wrappedResult bool
	wrappedError error

	matcher Matcher
}

func init() { RegisterTestSuite(&ErrorTest{}) }

func (t *ErrorTest) SetUp(i *TestInfo) {
	wrapped := &fakeMatcher{
		func(c interface{}) (bool, error) {
			t.matcherCalled = true
			t.suppliedCandidate = c
			return t.wrappedResult, t.wrappedError
		},
		"is foo",
	}

	t.matcher = Error(wrapped)
}

func isFatal(err error) bool {
	_, isFatal := err.(*FatalError)
	return isFatal
}

////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////

func (t *ErrorTest) Description() {
	ExpectThat(t.matcher.Description(), Equals("error is foo"))
}

func (t *ErrorTest) CandidateIsNil() {
	res, err := t.matcher.Matches(nil)

	ExpectThat(t.matcherCalled, Equals(false))
	ExpectThat(res, Equals(false))
	ExpectThat(err.Error(), Equals("which is not an error"))
	ExpectTrue(isFatal(err))
}

func (t *ErrorTest) CandidateIsString() {
	res, err := t.matcher.Matches("taco")

	ExpectThat(t.matcherCalled, Equals(false))
	ExpectThat(res, Equals(false))
	ExpectThat(err.Error(), Equals("which is not an error"))
	ExpectTrue(isFatal(err))
}

func (t *ErrorTest) CallsWrappedMatcher() {
	candidate := errors.New("taco")
	t.matcher.Matches(candidate)

	ExpectThat(t.matcherCalled, Equals(true))
	ExpectThat(t.suppliedCandidate, Equals("taco"))
}

func (t *ErrorTest) ReturnsWrappedMatcherResult() {
	t.wrappedResult = true
	t.wrappedError = errors.New("burrito")

	res, err := t.matcher.Matches(errors.New(""))

	ExpectThat(res, Equals(true))
	ExpectThat(err, Equals(t.wrappedError))
}
<|endoftext|>


<commit_before>package garden

import (
	"io"
	"time"
)

//go:generate counterfeiter . Container

type Container interface {
	Handle() string

	// Stop stops a container.
	//
	// If kill is false, garden stops a container by sending the processes running inside it the SIGTERM signal.
	// It then waits for the processes to terminate before returning a response.
	// If one or more processes do not terminate within 10 seconds,
	// garden sends these processes the SIGKILL signal, killing them ungracefully.
	//
	// If kill is true, garden stops a container by sending the processing running inside it a SIGKILL signal.
	//
	// It is possible to copy files in to and out of a stopped container.
	// It is only when a container is destroyed that its filesystem is cleaned up.
	//
	// Errors:
	// * None.
	Stop(kill bool) error

	// Returns information about a container.
	Info() (ContainerInfo, error)

	// StreamIn streams data into a file in a container.
	//
	// Errors:
	// *  TODO.
	StreamIn(spec StreamInSpec) error

	// StreamOut streams a file out of a container.
	//
	// Errors:
	// * TODO.
	StreamOut(spec StreamOutSpec) (io.ReadCloser, error)

	// Returns the current bandwidth limits set for the container.
	CurrentBandwidthLimits() (BandwidthLimits, error)

	// Returns the current CPU limts set for the container.
	CurrentCPULimits() (CPULimits, error)

	// Returns the current disk limts set for the container.
	CurrentDiskLimits() (DiskLimits, error)

	// Returns the current memory limts set for the container.
	CurrentMemoryLimits() (MemoryLimits, error)

	// Map a port on the host to a port in the container so that traffic to the
	// host port is forwarded to the container port. This is deprecated in
	// favour of passing NetIn configuration in the ContainerSpec at creation
	// time.
	//
	// If a host port is not given, a port will be acquired from the server's port
	// pool.
	//
	// If a container port is not given, the port will be the same as the
	// container port.
	//
	// The resulting host and container ports are returned in that order.
	//
	// Errors:
	// * When no port can be acquired from the server's port pool.
	NetIn(hostPort, containerPort uint32) (uint32, uint32, error)

	// Whitelist outbound network traffic. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// If the configuration directive deny_networks is not used,
	// all networks are already whitelisted and this command is effectively a no-op.
	//
	// Later NetOut calls take precedence over earlier calls, which is
	// significant only in relation to logging.
	//
	// Errors:
	// * An error is returned if the NetOut call fails.
	NetOut(netOutRule NetOutRule) error

	// A Bulk call for NetOut. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// Errors:
	// * An error is returned if any of the NetOut calls fail.
	BulkNetOut(netOutRules []NetOutRule) error

	// Run a script inside a container.
	//
	// The root user will be mapped to a non-root UID in the host unless the container (not this process) was created with 'privileged' true.
	//
	// Errors:
	// * TODO.
	Run(ProcessSpec, ProcessIO) (Process, error)

	// Attach starts streaming the output back to the client from a specified process.
	//
	// Errors:
	// * processID does not refer to a running process.
	Attach(processID string, io ProcessIO) (Process, error)

	// Metrics returns the current set of metrics for a container
	Metrics() (Metrics, error)

	// Sets the grace time.
	SetGraceTime(graceTime time.Duration) error

	// Properties returns the current set of properties
	Properties() (Properties, error)

	// Property returns the value of the property with the specified name.
	//
	// Errors:
	// * When the property does not exist on the container.
	Property(name string) (string, error)

	// Set a named property on a container to a specified value.
	//
	// Errors:
	// * None.
	SetProperty(name string, value string) error

	// Remove a property with the specified name from a container.
	//
	// Errors:
	// * None.
	RemoveProperty(name string) error
}

// ProcessSpec contains parameters for running a script inside a container.
type ProcessSpec struct {
	// ID for the process. If empty, an ID will be generated.
	ID string `json:"id,omitempty"`

	// Path to command to execute.
	Path string `json:"path,omitempty"`

	// Arguments to pass to command.
	Args []string `json:"args,omitempty"`

	// Environment variables.
	Env []string `json:"env,omitempty"`

	// Working directory (default: home directory).
	Dir string `json:"dir,omitempty"`

	// The name of a user in the container to run the process as.
	// This must either be a username, or uid:gid.
	User string `json:"user,omitempty"`

	// Resource limits
	Limits ResourceLimits `json:"rlimits,omitempty"`

	// Limits to be applied to the newly created process
	OverrideContainerLimits *ProcessLimits `json:"limits,omitempty"`

	// Execute with a TTY for stdio.
	TTY *TTYSpec `json:"tty,omitempty"`

	// Execute process in own root filesystem, different from the other processes
	// in the container.
	Image ImageRef `json:"image,omitempty"`

	// Bind mounts to be applied to the process's filesystem
	// An error is returned if ProcessSpec.Image is not also set.
	BindMounts []BindMount `json:"bind_mounts,omitempty"`
}

type TTYSpec struct {
	WindowSize *WindowSize `json:"window_size,omitempty"`
}

type WindowSize struct {
	Columns int `json:"columns,omitempty"`
	Rows    int `json:"rows,omitempty"`
}

type ProcessIO struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

//go:generate counterfeiter . Process

type Process interface {
	ID() string
	Wait() (int, error)
	SetTTY(TTYSpec) error
	Signal(Signal) error
}

type Signal int

const (
	SignalTerminate Signal = iota
	SignalKill
)

type PortMapping struct {
	HostPort      uint32
	ContainerPort uint32
}

type StreamInSpec struct {
	Path      string
	User      string
	TarStream io.Reader
}

type StreamOutSpec struct {
	Path string
	User string
}

// ContainerInfo holds information about a container.
type ContainerInfo struct {
	State         string        // Either "active" or "stopped".
	Events        []string      // List of events that occurred for the container. It currently includes only "oom" (Out Of Memory) event if it occurred.
	HostIP        string        // The IP address of the gateway which controls the host side of the container's virtual ethernet pair.
	ContainerIP   string        // The IP address of the container side of the container's virtual ethernet pair.
	ExternalIP    string        //
	ContainerPath string        // The path to the directory holding the container's files (both its control scripts and filesystem).
	ProcessIDs    []string      // List of running processes.
	Properties    Properties    // List of properties defined for the container.
	MappedPorts   []PortMapping //
}

type ContainerInfoEntry struct {
	Info ContainerInfo
	Err  *Error
}

type Metrics struct {
	MemoryStat  ContainerMemoryStat
	CPUStat     ContainerCPUStat
	DiskStat    ContainerDiskStat
	NetworkStat ContainerNetworkStat
}

type ContainerMetricsEntry struct {
	Metrics Metrics
	Err     *Error
}

type ContainerMemoryStat struct {
	ActiveAnon              uint64 `json:"active_anon"`
	ActiveFile              uint64 `json:"active_file"`
	Cache                   uint64 `json:"cache"`
	HierarchicalMemoryLimit uint64 `json:"hierarchical_memory_limit"`
	InactiveAnon            uint64 `json:"inactive_anon"`
	InactiveFile            uint64 `json:"inactive_file"`
	MappedFile              uint64 `json:"mapped_file"`
	Pgfault                 uint64 `json:"pgfault"`
	Pgmajfault              uint64 `json:"pgmajfault"`
	Pgpgin                  uint64 `json:"pgpgin"`
	Pgpgout                 uint64 `json:"pgpgout"`
	Rss                     uint64 `json:"rss"`
	TotalActiveAnon         uint64 `json:"total_active_anon"`
	TotalActiveFile         uint64 `json:"total_active_file"`
	TotalCache              uint64 `json:"total_cache"`
	TotalInactiveAnon       uint64 `json:"total_inactive_anon"`
	TotalInactiveFile       uint64 `json:"total_inactive_file"`
	TotalMappedFile         uint64 `json:"total_mapped_file"`
	TotalPgfault            uint64 `json:"total_pgfault"`
	TotalPgmajfault         uint64 `json:"total_pgmajfault"`
	TotalPgpgin             uint64 `json:"total_pgpgin"`
	TotalPgpgout            uint64 `json:"total_pgpgout"`
	TotalRss                uint64 `json:"total_rss"`
	TotalUnevictable        uint64 `json:"total_unevictable"`
	Unevictable             uint64 `json:"unevictable"`
	Swap                    uint64 `json:"swap"`
	HierarchicalMemswLimit  uint64 `json:"hierarchical_memsw_limit"`
	TotalSwap               uint64 `json:"total_swap"`
	// A memory usage total which reports memory usage in the same way that limits are enforced.
	// This value includes memory consumed by nested containers.
	TotalUsageTowardLimit uint64
}

type ContainerCPUStat struct {
	Usage  uint64
	User   uint64
	System uint64
}

type ContainerDiskStat struct {
	TotalBytesUsed      uint64
	TotalInodesUsed     uint64
	ExclusiveBytesUsed  uint64
	ExclusiveInodesUsed uint64
}

type ContainerBandwidthStat struct {
	InRate   uint64
	InBurst  uint64
	OutRate  uint64
	OutBurst uint64
}

type ContainerNetworkStat struct {
	RxBytes uint64
	TxBytes uint64
}

type BandwidthLimits struct {
	RateInBytesPerSecond      uint64 `json:"rate,omitempty"`
	BurstRateInBytesPerSecond uint64 `json:"burst,omitempty"`
}

type ProcessLimits struct {
	CPU    CPULimits    `json:"cpu_limits,omitempty"`
	Memory MemoryLimits `json:"memory_limits,omitempty"`
}

type DiskLimits struct {
	InodeSoft uint64 `json:"inode_soft,omitempty"`
	InodeHard uint64 `json:"inode_hard,omitempty"`

	ByteSoft uint64 `json:"byte_soft,omitempty"`
	ByteHard uint64 `json:"byte_hard,omitempty"`

	Scope DiskLimitScope `json:"scope,omitempty"`
}

type MemoryLimits struct {
	//	Memory usage limit in bytes.
	LimitInBytes uint64 `json:"limit_in_bytes,omitempty"`
}

type CPULimits struct {
	LimitInShares uint64 `json:"limit_in_shares,omitempty"`
}

type PidLimits struct {
	// Limits the number of pids a container may create before new forks or clones are disallowed to processes in the container.
	// Note: this may only be enforced when a process attempts to fork, so it does not guarantee that a new container.Run(ProcessSpec)
	// will not succeed even if the limit has been exceeded, but the process will not be able to spawn further processes or threads.
	Max uint64 `json:"max,omitempty"`
}

// Resource limits.
//
// Please refer to the manual page of getrlimit for a description of the individual fields:
// http://www.kernel.org/doc/man-pages/online/pages/man2/getrlimit.2.html
type ResourceLimits struct {
	As         *uint64 `json:"as,omitempty"`
	Core       *uint64 `json:"core,omitempty"`
	Cpu        *uint64 `json:"cpu,omitempty"`
	Data       *uint64 `json:"data,omitempty"`
	Fsize      *uint64 `json:"fsize,omitempty"`
	Locks      *uint64 `json:"locks,omitempty"`
	Memlock    *uint64 `json:"memlock,omitempty"`
	Msgqueue   *uint64 `json:"msgqueue,omitempty"`
	Nice       *uint64 `json:"nice,omitempty"`
	Nofile     *uint64 `json:"nofile,omitempty"`
	Nproc      *uint64 `json:"nproc,omitempty"`
	Rss        *uint64 `json:"rss,omitempty"`
	Rtprio     *uint64 `json:"rtprio,omitempty"`
	Sigpending *uint64 `json:"sigpending,omitempty"`
	Stack      *uint64 `json:"stack,omitempty"`
}

type DiskLimitScope uint8

const DiskLimitScopeTotal DiskLimitScope = 0
const DiskLimitScopeExclusive DiskLimitScope = 1
<commit_msg>fix a wrong comment.<commit_after>package garden

import (
	"io"
	"time"
)

//go:generate counterfeiter . Container

type Container interface {
	Handle() string

	// Stop stops a container.
	//
	// If kill is false, garden stops a container by sending the processes running inside it the SIGTERM signal.
	// It then waits for the processes to terminate before returning a response.
	// If one or more processes do not terminate within 10 seconds,
	// garden sends these processes the SIGKILL signal, killing them ungracefully.
	//
	// If kill is true, garden stops a container by sending the processing running inside it a SIGKILL signal.
	//
	// It is possible to copy files in to and out of a stopped container.
	// It is only when a container is destroyed that its filesystem is cleaned up.
	//
	// Errors:
	// * None.
	Stop(kill bool) error

	// Returns information about a container.
	Info() (ContainerInfo, error)

	// StreamIn streams data into a file in a container.
	//
	// Errors:
	// *  TODO.
	StreamIn(spec StreamInSpec) error

	// StreamOut streams a file out of a container.
	//
	// Errors:
	// * TODO.
	StreamOut(spec StreamOutSpec) (io.ReadCloser, error)

	// Returns the current bandwidth limits set for the container.
	CurrentBandwidthLimits() (BandwidthLimits, error)

	// Returns the current CPU limts set for the container.
	CurrentCPULimits() (CPULimits, error)

	// Returns the current disk limts set for the container.
	CurrentDiskLimits() (DiskLimits, error)

	// Returns the current memory limts set for the container.
	CurrentMemoryLimits() (MemoryLimits, error)

	// Map a port on the host to a port in the container so that traffic to the
	// host port is forwarded to the container port. This is deprecated in
	// favour of passing NetIn configuration in the ContainerSpec at creation
	// time.
	//
	// If a host port is not given, a port will be acquired from the server's port
	// pool.
	//
	// If a container port is not given, the port will be the same as the
	// host port.
	//
	// The resulting host and container ports are returned in that order.
	//
	// Errors:
	// * When no port can be acquired from the server's port pool.
	NetIn(hostPort, containerPort uint32) (uint32, uint32, error)

	// Whitelist outbound network traffic. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// If the configuration directive deny_networks is not used,
	// all networks are already whitelisted and this command is effectively a no-op.
	//
	// Later NetOut calls take precedence over earlier calls, which is
	// significant only in relation to logging.
	//
	// Errors:
	// * An error is returned if the NetOut call fails.
	NetOut(netOutRule NetOutRule) error

	// A Bulk call for NetOut. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// Errors:
	// * An error is returned if any of the NetOut calls fail.
	BulkNetOut(netOutRules []NetOutRule) error

	// Run a script inside a container.
	//
	// The root user will be mapped to a non-root UID in the host unless the container (not this process) was created with 'privileged' true.
	//
	// Errors:
	// * TODO.
	Run(ProcessSpec, ProcessIO) (Process, error)

	// Attach starts streaming the output back to the client from a specified process.
	//
	// Errors:
	// * processID does not refer to a running process.
	Attach(processID string, io ProcessIO) (Process, error)

	// Metrics returns the current set of metrics for a container
	Metrics() (Metrics, error)

	// Sets the grace time.
	SetGraceTime(graceTime time.Duration) error

	// Properties returns the current set of properties
	Properties() (Properties, error)

	// Property returns the value of the property with the specified name.
	//
	// Errors:
	// * When the property does not exist on the container.
	Property(name string) (string, error)

	// Set a named property on a container to a specified value.
	//
	// Errors:
	// * None.
	SetProperty(name string, value string) error

	// Remove a property with the specified name from a container.
	//
	// Errors:
	// * None.
	RemoveProperty(name string) error
}

// ProcessSpec contains parameters for running a script inside a container.
type ProcessSpec struct {
	// ID for the process. If empty, an ID will be generated.
	ID string `json:"id,omitempty"`

	// Path to command to execute.
	Path string `json:"path,omitempty"`

	// Arguments to pass to command.
	Args []string `json:"args,omitempty"`

	// Environment variables.
	Env []string `json:"env,omitempty"`

	// Working directory (default: home directory).
	Dir string `json:"dir,omitempty"`

	// The name of a user in the container to run the process as.
	// This must either be a username, or uid:gid.
	User string `json:"user,omitempty"`

	// Resource limits
	Limits ResourceLimits `json:"rlimits,omitempty"`

	// Limits to be applied to the newly created process
	OverrideContainerLimits *ProcessLimits `json:"limits,omitempty"`

	// Execute with a TTY for stdio.
	TTY *TTYSpec `json:"tty,omitempty"`

	// Execute process in own root filesystem, different from the other processes
	// in the container.
	Image ImageRef `json:"image,omitempty"`

	// Bind mounts to be applied to the process's filesystem
	// An error is returned if ProcessSpec.Image is not also set.
	BindMounts []BindMount `json:"bind_mounts,omitempty"`
}

type TTYSpec struct {
	WindowSize *WindowSize `json:"window_size,omitempty"`
}

type WindowSize struct {
	Columns int `json:"columns,omitempty"`
	Rows    int `json:"rows,omitempty"`
}

type ProcessIO struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

//go:generate counterfeiter . Process

type Process interface {
	ID() string
	Wait() (int, error)
	SetTTY(TTYSpec) error
	Signal(Signal) error
}

type Signal int

const (
	SignalTerminate Signal = iota
	SignalKill
)

type PortMapping struct {
	HostPort      uint32
	ContainerPort uint32
}

type StreamInSpec struct {
	Path      string
	User      string
	TarStream io.Reader
}

type StreamOutSpec struct {
	Path string
	User string
}

// ContainerInfo holds information about a container.
type ContainerInfo struct {
	State         string        // Either "active" or "stopped".
	Events        []string      // List of events that occurred for the container. It currently includes only "oom" (Out Of Memory) event if it occurred.
	HostIP        string        // The IP address of the gateway which controls the host side of the container's virtual ethernet pair.
	ContainerIP   string        // The IP address of the container side of the container's virtual ethernet pair.
	ExternalIP    string        //
	ContainerPath string        // The path to the directory holding the container's files (both its control scripts and filesystem).
	ProcessIDs    []string      // List of running processes.
	Properties    Properties    // List of properties defined for the container.
	MappedPorts   []PortMapping //
}

type ContainerInfoEntry struct {
	Info ContainerInfo
	Err  *Error
}

type Metrics struct {
	MemoryStat  ContainerMemoryStat
	CPUStat     ContainerCPUStat
	DiskStat    ContainerDiskStat
	NetworkStat ContainerNetworkStat
}

type ContainerMetricsEntry struct {
	Metrics Metrics
	Err     *Error
}

type ContainerMemoryStat struct {
	ActiveAnon              uint64 `json:"active_anon"`
	ActiveFile              uint64 `json:"active_file"`
	Cache                   uint64 `json:"cache"`
	HierarchicalMemoryLimit uint64 `json:"hierarchical_memory_limit"`
	InactiveAnon            uint64 `json:"inactive_anon"`
	InactiveFile            uint64 `json:"inactive_file"`
	MappedFile              uint64 `json:"mapped_file"`
	Pgfault                 uint64 `json:"pgfault"`
	Pgmajfault              uint64 `json:"pgmajfault"`
	Pgpgin                  uint64 `json:"pgpgin"`
	Pgpgout                 uint64 `json:"pgpgout"`
	Rss                     uint64 `json:"rss"`
	TotalActiveAnon         uint64 `json:"total_active_anon"`
	TotalActiveFile         uint64 `json:"total_active_file"`
	TotalCache              uint64 `json:"total_cache"`
	TotalInactiveAnon       uint64 `json:"total_inactive_anon"`
	TotalInactiveFile       uint64 `json:"total_inactive_file"`
	TotalMappedFile         uint64 `json:"total_mapped_file"`
	TotalPgfault            uint64 `json:"total_pgfault"`
	TotalPgmajfault         uint64 `json:"total_pgmajfault"`
	TotalPgpgin             uint64 `json:"total_pgpgin"`
	TotalPgpgout            uint64 `json:"total_pgpgout"`
	TotalRss                uint64 `json:"total_rss"`
	TotalUnevictable        uint64 `json:"total_unevictable"`
	Unevictable             uint64 `json:"unevictable"`
	Swap                    uint64 `json:"swap"`
	HierarchicalMemswLimit  uint64 `json:"hierarchical_memsw_limit"`
	TotalSwap               uint64 `json:"total_swap"`
	// A memory usage total which reports memory usage in the same way that limits are enforced.
	// This value includes memory consumed by nested containers.
	TotalUsageTowardLimit uint64
}

type ContainerCPUStat struct {
	Usage  uint64
	User   uint64
	System uint64
}

type ContainerDiskStat struct {
	TotalBytesUsed      uint64
	TotalInodesUsed     uint64
	ExclusiveBytesUsed  uint64
	ExclusiveInodesUsed uint64
}

type ContainerBandwidthStat struct {
	InRate   uint64
	InBurst  uint64
	OutRate  uint64
	OutBurst uint64
}

type ContainerNetworkStat struct {
	RxBytes uint64
	TxBytes uint64
}

type BandwidthLimits struct {
	RateInBytesPerSecond      uint64 `json:"rate,omitempty"`
	BurstRateInBytesPerSecond uint64 `json:"burst,omitempty"`
}

type ProcessLimits struct {
	CPU    CPULimits    `json:"cpu_limits,omitempty"`
	Memory MemoryLimits `json:"memory_limits,omitempty"`
}

type DiskLimits struct {
	InodeSoft uint64 `json:"inode_soft,omitempty"`
	InodeHard uint64 `json:"inode_hard,omitempty"`

	ByteSoft uint64 `json:"byte_soft,omitempty"`
	ByteHard uint64 `json:"byte_hard,omitempty"`

	Scope DiskLimitScope `json:"scope,omitempty"`
}

type MemoryLimits struct {
	//	Memory usage limit in bytes.
	LimitInBytes uint64 `json:"limit_in_bytes,omitempty"`
}

type CPULimits struct {
	LimitInShares uint64 `json:"limit_in_shares,omitempty"`
}

type PidLimits struct {
	// Limits the number of pids a container may create before new forks or clones are disallowed to processes in the container.
	// Note: this may only be enforced when a process attempts to fork, so it does not guarantee that a new container.Run(ProcessSpec)
	// will not succeed even if the limit has been exceeded, but the process will not be able to spawn further processes or threads.
	Max uint64 `json:"max,omitempty"`
}

// Resource limits.
//
// Please refer to the manual page of getrlimit for a description of the individual fields:
// http://www.kernel.org/doc/man-pages/online/pages/man2/getrlimit.2.html
type ResourceLimits struct {
	As         *uint64 `json:"as,omitempty"`
	Core       *uint64 `json:"core,omitempty"`
	Cpu        *uint64 `json:"cpu,omitempty"`
	Data       *uint64 `json:"data,omitempty"`
	Fsize      *uint64 `json:"fsize,omitempty"`
	Locks      *uint64 `json:"locks,omitempty"`
	Memlock    *uint64 `json:"memlock,omitempty"`
	Msgqueue   *uint64 `json:"msgqueue,omitempty"`
	Nice       *uint64 `json:"nice,omitempty"`
	Nofile     *uint64 `json:"nofile,omitempty"`
	Nproc      *uint64 `json:"nproc,omitempty"`
	Rss        *uint64 `json:"rss,omitempty"`
	Rtprio     *uint64 `json:"rtprio,omitempty"`
	Sigpending *uint64 `json:"sigpending,omitempty"`
	Stack      *uint64 `json:"stack,omitempty"`
}

type DiskLimitScope uint8

const DiskLimitScopeTotal DiskLimitScope = 0
const DiskLimitScopeExclusive DiskLimitScope = 1
<|endoftext|>


<commit_before>package es

import (
	"testing"
	"time"

	"github.com/dynport/dgtk/tskip/tskip"
)

type TestLog struct {
	Id      int
	Tag     string
	Host    string
	Created string
	Raw     string
}

var sleepFor = 2 * time.Second

var index = &Index{
	Host:  "127.0.0.1",
	Port:  9200,
	Index: "test",
	Type:  "logs",
}

func validateEsRunning(t *testing.T) {
	_, e := index.Status()
	if e != nil {
		tskip.Errorf(t, 1, "ElasticSearch is not running on %s. %s", index.BaseUrl(), e.Error())
		t.FailNow()
	}
}

func setupIndex() error {
	index.DeleteIndex()
	index.EnqueueBulkIndex("1", &TestLog{Id: 1, Tag: "nginx", Created: "2013-12-03"})
	index.EnqueueBulkIndex("2", &TestLog{Id: 1, Tag: "unicorn", Created: "2013-12-03"})

	index.EnqueueBulkIndex("3", &TestLog{Id: 1, Tag: "nginx", Created: "2013-12-02"})
	index.EnqueueBulkIndex("4", &TestLog{Id: 1, Tag: "unicorn", Created: "2013-12-02"})

	index.EnqueueBulkIndex("5", &TestLog{Id: 1, Tag: "nginx", Created: "2013-12-01"})
	index.EnqueueBulkIndex("6", &TestLog{Id: 1, Tag: "unicorn", Created: "2013-12-01"})
	return index.RunBatchIndex()
}

func TestDeleteFromImage(t *testing.T) {
	if !runIntegrationTests {
		t.SkipNow()
	}
	validateEsRunning(t)
	failIfError(t, setupIndex())
	failIfError(t, index.Refresh())
	req := &Request{
		Size: 10,
	}
	rsp, err := index.Search(req)
	failIfError(t, err)
	assertEqual(t, rsp.Hits.Total, 6)

	_, err = index.DeleteByQuery("Tag:nginx")
	failIfError(t, err)
	failIfError(t, index.Refresh())

	rsp, err = index.Search(req)
	failIfError(t, err)
	assertEqual(t, rsp.Hits.Total, 3)

	tags := map[string]int{}
	for _, hit := range rsp.Hits.Hits {
		tag := hit.Source["Tag"]
		switch tag := tag.(type) {
		case string:
			tags[tag]++
		}
	}
	assertEqual(t, tags["unicorn"], 3)

	_, err = index.DeleteByQuery("Tag:unicorn AND Created:[2013-12-01 TO 2013-12-02]")
	failIfError(t, err)
	failIfError(t, index.Refresh())

	rsp, err = index.Search(req)
	failIfError(t, err)
	assertEqual(t, rsp.Hits.Total, 1)
	assertEqual(t, rsp.Hits.Hits[0].Source["Created"], "2013-12-03")
}
<commit_msg>fix tests<commit_after>package es

import (
	"testing"
	"time"

	"github.com/dynport/dgtk/tskip/tskip"
)

type TestLog struct {
	Id      int
	Tag     string
	Host    string
	Created string
	Raw     string
}

var sleepFor = 2 * time.Second

var index = &Index{
	Address: "http://127.0.0.1:9200",
	Index:   "test",
	Type:    "logs",
}

func validateEsRunning(t *testing.T) {
	_, e := index.Status()
	if e != nil {
		tskip.Errorf(t, 1, "ElasticSearch is not running on %s. %s", index.BaseUrl(), e.Error())
		t.FailNow()
	}
}

func setupIndex() error {
	index.DeleteIndex()
	index.EnqueueBulkIndex("1", &TestLog{Id: 1, Tag: "nginx", Created: "2013-12-03"})
	index.EnqueueBulkIndex("2", &TestLog{Id: 1, Tag: "unicorn", Created: "2013-12-03"})

	index.EnqueueBulkIndex("3", &TestLog{Id: 1, Tag: "nginx", Created: "2013-12-02"})
	index.EnqueueBulkIndex("4", &TestLog{Id: 1, Tag: "unicorn", Created: "2013-12-02"})

	index.EnqueueBulkIndex("5", &TestLog{Id: 1, Tag: "nginx", Created: "2013-12-01"})
	index.EnqueueBulkIndex("6", &TestLog{Id: 1, Tag: "unicorn", Created: "2013-12-01"})
	return index.RunBatchIndex()
}

func TestDeleteFromImage(t *testing.T) {
	if !runIntegrationTests {
		t.SkipNow()
	}
	validateEsRunning(t)
	failIfError(t, setupIndex())
	failIfError(t, index.Refresh())
	req := &Request{
		Size: 10,
	}
	rsp, err := index.Search(req)
	failIfError(t, err)
	assertEqual(t, rsp.Hits.Total, 6)

	_, err = index.DeleteByQuery("Tag:nginx")
	failIfError(t, err)
	failIfError(t, index.Refresh())

	rsp, err = index.Search(req)
	failIfError(t, err)
	assertEqual(t, rsp.Hits.Total, 3)

	tags := map[string]int{}
	for _, hit := range rsp.Hits.Hits {
		tag := hit.Source["Tag"]
		switch tag := tag.(type) {
		case string:
			tags[tag]++
		}
	}
	assertEqual(t, tags["unicorn"], 3)

	_, err = index.DeleteByQuery("Tag:unicorn AND Created:[2013-12-01 TO 2013-12-02]")
	failIfError(t, err)
	failIfError(t, index.Refresh())

	rsp, err = index.Search(req)
	failIfError(t, err)
	assertEqual(t, rsp.Hits.Total, 1)
	assertEqual(t, rsp.Hits.Hits[0].Source["Created"], "2013-12-03")
}
<|endoftext|>


<commit_before>// Copyright 2018 The Go Cloud 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
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pubsub_test

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"

	"github.com/google/go-cloud/internal/pubsub"
	"github.com/google/go-cloud/internal/pubsub/driver"
)

type driverTopic struct {
	subs []*driverSub
}

func (t *driverTopic) SendBatch(ctx context.Context, ms []*driver.Message) error {
	for _, s := range t.subs {
		select {
		case <-s.sem:
			s.q = append(s.q, ms...)
			s.sem <- struct{}{}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	return nil
}

func (t *driverTopic) Close() error {
	return nil
}

func (s *driverTopic) IsRetryable(error) bool { return false }

type driverSub struct {
	sem chan struct{}
	// Normally this queue would live on a separate server in the cloud.
	q []*driver.Message
}

func NewDriverSub() *driverSub {
	ds := &driverSub{
		sem: make(chan struct{}, 1),
	}
	ds.sem <- struct{}{}
	return ds
}

func (s *driverSub) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	for {
		select {
		case <-s.sem:
			ms := s.grabQueue(maxMessages)
			if len(ms) != 0 {
				return ms, nil
			}
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}
	}
}

func (s *driverSub) grabQueue(maxMessages int) []*driver.Message {
	defer func() { s.sem <- struct{}{} }()
	if len(s.q) > 0 {
		if len(s.q) <= maxMessages {
			ms := s.q
			s.q = nil
			return ms
		}
		ms := s.q[:maxMessages]
		s.q = s.q[maxMessages:]
		return ms
	}
	return nil
}

func (s *driverSub) SendAcks(ctx context.Context, ackIDs []driver.AckID) error {
	return nil
}

func (s *driverSub) Close() error {
	return nil
}

func (s *driverSub) IsRetryable(error) bool { return false }

func TestSendReceive(t *testing.T) {
	ctx := context.Background()
	ds := NewDriverSub()
	dt := &driverTopic{
		subs: []*driverSub{ds},
	}
	topic := pubsub.NewTopic(dt)
	defer topic.Close()
	m := &pubsub.Message{Body: []byte("user signed up")}
	if err := topic.Send(ctx, m); err != nil {
		t.Fatal(err)
	}

	sub := pubsub.NewSubscription(ds)
	defer sub.Close()
	m2, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if string(m2.Body) != string(m.Body) {
		t.Fatalf("received message has body %q, want %q", m2.Body, m.Body)
	}
}

func TestConcurrentReceivesGetAllTheMessages(t *testing.T) {
	howManyToSend := int(1e3)
	ctx, cancel := context.WithCancel(context.Background())
	dt := &driverTopic{}

	// wg is used to wait until all messages are received.
	var wg sync.WaitGroup
	wg.Add(howManyToSend)

	// Make a subscription.
	ds := NewDriverSub()
	dt.subs = append(dt.subs, ds)
	s := pubsub.NewSubscription(ds)
	defer s.Close()

	// Start 10 goroutines to receive from it.
	var mu sync.Mutex
	receivedMsgs := make(map[string]int)
	for i := 0; i < 10; i++ {
		go func() {
			for {
				m, err := s.Receive(ctx)
				if err != nil {
					// Permanent error; ctx cancelled or subscription closed is
					// expected once we've received all the messages.
					mu.Lock()
					n := len(receivedMsgs)
					mu.Unlock()
					if n != howManyToSend {
						t.Errorf("Worker's Receive failed before all messages were received (%d)", n)
					}
					return
				}
				mu.Lock()
				receivedMsgs[string(m.Body)]++
				mu.Unlock()
				wg.Done()
			}
		}()
	}

	// Send messages.
	topic := pubsub.NewTopic(dt)
	defer topic.Close()
	sentMsgs := make(map[string]int)
	for i := 0; i < howManyToSend; i++ {
		bod := fmt.Sprintf("%d", i)
		m := &pubsub.Message{Body: []byte(bod)}
		sentMsgs[string(m.Body)]++
		if err := topic.Send(ctx, m); err != nil {
			t.Fatal(err)
		}
	}

	// Wait for the goroutines to receive all of the messages, then cancel the
	// ctx so they all exit.
	wg.Wait()
	defer cancel()

	// Check that all the messages were received.
	sum := 0
	for _, n := range receivedMsgs {
		sum += n
	}
	if sum != howManyToSend {
		t.Errorf("received %d messages, want %d", sum, howManyToSend)
	}
	for k, v := range sentMsgs {
		v2 := receivedMsgs[k]
		if v2 != v {
			t.Errorf("got %d for %q, want %d", v2, k, v)
		}
	}
}

func TestCancelSend(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	ds := NewDriverSub()
	dt := &driverTopic{
		subs: []*driverSub{ds},
	}
	topic := pubsub.NewTopic(dt)
	defer topic.Close()
	m := &pubsub.Message{}

	// Intentionally break the driver subscription by acquiring its semaphore.
	// Now topic.Send will have to wait for cancellation.
	<-ds.sem

	cancel()
	if err := topic.Send(ctx, m); err == nil {
		t.Error("got nil, want cancellation error")
	}
}

func TestCancelReceive(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	ds := NewDriverSub()
	s := pubsub.NewSubscription(ds)
	defer s.Close()
	cancel()
	// Without cancellation, this Receive would hang.
	if _, err := s.Receive(ctx); err == nil {
		t.Error("got nil, want cancellation error")
	}
}

func TestRetryTopic(t *testing.T) {
	// Test that Send is retried if the driver returns a retryable error.
	ft := &failTopic{}
	top := pubsub.NewTopic(ft)
	err := top.Send(context.Background(), &pubsub.Message{})
	if err != nil {
		t.Errorf("Send: got %v, want nil", err)
	}
	if got, want := ft.calls, nRetryCalls+1; got != want {
		t.Errorf("calls: got %d, want %d", got, want)
	}
}

var errRetry = errors.New("retry")

func isRetryable(err error) bool {
	return err == errRetry
}

const nRetryCalls = 2

type failTopic struct {
	driver.Topic
	calls int
}

func (t *failTopic) SendBatch(ctx context.Context, ms []*driver.Message) error {
	t.calls++
	if t.calls <= nRetryCalls {
		return errRetry
	}
	return nil
}

func (t *failTopic) IsRetryable(err error) bool { return isRetryable(err) }

func TestRetryReceive(t *testing.T) {
	fs := &failSub{}
	sub := pubsub.NewSubscription(fs)
	_, err := sub.Receive(context.Background())
	if err != nil {
		t.Errorf("Receive: got %v, want nil", err)
	}
	if got, want := fs.calls, nRetryCalls+1; got != want {
		t.Errorf("calls: got %d, want %d", got, want)
	}
}

type failSub struct {
	driver.Subscription
	calls int
}

func (t *failSub) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	t.calls++
	if t.calls <= nRetryCalls {
		return nil, errRetry
	}
	return []*driver.Message{{Body: []byte("")}}, nil
}

func (t *failSub) IsRetryable(err error) bool { return isRetryable(err) }

// TODO(jba): add a test for retry of SendAcks.
<commit_msg>pubsub/tests: simplify TestConcurrentReceivesGetAllTheMessages (#868)<commit_after>// Copyright 2018 The Go Cloud 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
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package pubsub_test

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"

	"github.com/google/go-cloud/internal/pubsub"
	"github.com/google/go-cloud/internal/pubsub/driver"
)

type driverTopic struct {
	subs []*driverSub
}

func (t *driverTopic) SendBatch(ctx context.Context, ms []*driver.Message) error {
	for _, s := range t.subs {
		select {
		case <-s.sem:
			s.q = append(s.q, ms...)
			s.sem <- struct{}{}
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	return nil
}

func (t *driverTopic) Close() error {
	return nil
}

func (s *driverTopic) IsRetryable(error) bool { return false }

type driverSub struct {
	sem chan struct{}
	// Normally this queue would live on a separate server in the cloud.
	q []*driver.Message
}

func NewDriverSub() *driverSub {
	ds := &driverSub{
		sem: make(chan struct{}, 1),
	}
	ds.sem <- struct{}{}
	return ds
}

func (s *driverSub) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	for {
		select {
		case <-s.sem:
			ms := s.grabQueue(maxMessages)
			if len(ms) != 0 {
				return ms, nil
			}
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}
	}
}

func (s *driverSub) grabQueue(maxMessages int) []*driver.Message {
	defer func() { s.sem <- struct{}{} }()
	if len(s.q) > 0 {
		if len(s.q) <= maxMessages {
			ms := s.q
			s.q = nil
			return ms
		}
		ms := s.q[:maxMessages]
		s.q = s.q[maxMessages:]
		return ms
	}
	return nil
}

func (s *driverSub) SendAcks(ctx context.Context, ackIDs []driver.AckID) error {
	return nil
}

func (s *driverSub) Close() error {
	return nil
}

func (s *driverSub) IsRetryable(error) bool { return false }

func TestSendReceive(t *testing.T) {
	ctx := context.Background()
	ds := NewDriverSub()
	dt := &driverTopic{
		subs: []*driverSub{ds},
	}
	topic := pubsub.NewTopic(dt)
	defer topic.Close()
	m := &pubsub.Message{Body: []byte("user signed up")}
	if err := topic.Send(ctx, m); err != nil {
		t.Fatal(err)
	}

	sub := pubsub.NewSubscription(ds)
	defer sub.Close()
	m2, err := sub.Receive(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if string(m2.Body) != string(m.Body) {
		t.Fatalf("received message has body %q, want %q", m2.Body, m.Body)
	}
}

func TestConcurrentReceivesGetAllTheMessages(t *testing.T) {
	howManyToSend := int(1e3)
	ctx, cancel := context.WithCancel(context.Background())
	dt := &driverTopic{}

	// wg is used to wait until all messages are received.
	var wg sync.WaitGroup
	wg.Add(howManyToSend)

	// Make a subscription.
	ds := NewDriverSub()
	dt.subs = append(dt.subs, ds)
	s := pubsub.NewSubscription(ds)
	defer s.Close()

	// Start 10 goroutines to receive from it.
	var mu sync.Mutex
	receivedMsgs := make(map[string]bool)
	for i := 0; i < 10; i++ {
		go func() {
			for {
				m, err := s.Receive(ctx)
				if err != nil {
					// Permanent error; ctx cancelled or subscription closed is
					// expected once we've received all the messages.
					mu.Lock()
					n := len(receivedMsgs)
					mu.Unlock()
					if n != howManyToSend {
						t.Errorf("Worker's Receive failed before all messages were received (%d)", n)
					}
					return
				}
				mu.Lock()
				receivedMsgs[string(m.Body)] = true
				mu.Unlock()
				wg.Done()
			}
		}()
	}

	// Send messages. Each message has a unique body used as a key to receivedMsgs.
	topic := pubsub.NewTopic(dt)
	defer topic.Close()
	for i := 0; i < howManyToSend; i++ {
		key := fmt.Sprintf("message #%d", i)
		m := &pubsub.Message{Body: []byte(key)}
		if err := topic.Send(ctx, m); err != nil {
			t.Fatal(err)
		}
	}

	// Wait for the goroutines to receive all of the messages, then cancel the
	// ctx so they all exit.
	wg.Wait()
	defer cancel()

	// Check that all the messages were received.
	for i := 0; i < howManyToSend; i++ {
		key := fmt.Sprintf("message #%d", i)
		if !receivedMsgs[key] {
			t.Errorf("message %q was not received", key)
		}
	}
}

func TestCancelSend(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	ds := NewDriverSub()
	dt := &driverTopic{
		subs: []*driverSub{ds},
	}
	topic := pubsub.NewTopic(dt)
	defer topic.Close()
	m := &pubsub.Message{}

	// Intentionally break the driver subscription by acquiring its semaphore.
	// Now topic.Send will have to wait for cancellation.
	<-ds.sem

	cancel()
	if err := topic.Send(ctx, m); err == nil {
		t.Error("got nil, want cancellation error")
	}
}

func TestCancelReceive(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	ds := NewDriverSub()
	s := pubsub.NewSubscription(ds)
	defer s.Close()
	cancel()
	// Without cancellation, this Receive would hang.
	if _, err := s.Receive(ctx); err == nil {
		t.Error("got nil, want cancellation error")
	}
}

func TestRetryTopic(t *testing.T) {
	// Test that Send is retried if the driver returns a retryable error.
	ft := &failTopic{}
	top := pubsub.NewTopic(ft)
	err := top.Send(context.Background(), &pubsub.Message{})
	if err != nil {
		t.Errorf("Send: got %v, want nil", err)
	}
	if got, want := ft.calls, nRetryCalls+1; got != want {
		t.Errorf("calls: got %d, want %d", got, want)
	}
}

var errRetry = errors.New("retry")

func isRetryable(err error) bool {
	return err == errRetry
}

const nRetryCalls = 2

type failTopic struct {
	driver.Topic
	calls int
}

func (t *failTopic) SendBatch(ctx context.Context, ms []*driver.Message) error {
	t.calls++
	if t.calls <= nRetryCalls {
		return errRetry
	}
	return nil
}

func (t *failTopic) IsRetryable(err error) bool { return isRetryable(err) }

func TestRetryReceive(t *testing.T) {
	fs := &failSub{}
	sub := pubsub.NewSubscription(fs)
	_, err := sub.Receive(context.Background())
	if err != nil {
		t.Errorf("Receive: got %v, want nil", err)
	}
	if got, want := fs.calls, nRetryCalls+1; got != want {
		t.Errorf("calls: got %d, want %d", got, want)
	}
}

type failSub struct {
	driver.Subscription
	calls int
}

func (t *failSub) ReceiveBatch(ctx context.Context, maxMessages int) ([]*driver.Message, error) {
	t.calls++
	if t.calls <= nRetryCalls {
		return nil, errRetry
	}
	return []*driver.Message{{Body: []byte("")}}, nil
}

func (t *failSub) IsRetryable(err error) bool { return isRetryable(err) }

// TODO(jba): add a test for retry of SendAcks.
<|endoftext|>


<commit_before>package govalidator

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
)

// ToString convert the input to a string.
func ToString(obj interface{}) string {
	res := fmt.Sprintf("%v", obj)
	return string(res)
}

// ToJSON convert the input to a valid JSON string
func ToJSON(obj interface{}) (string, error) {
	res, err := json.Marshal(obj)
	if err != nil {
		res = []byte("")
	}
	return string(res), err
}

// ToFloat convert the input string to a float, or 0.0 if the input is not a float.
func ToFloat(str string) (float64, error) {
	res, err := strconv.ParseFloat(str, 64)
	if err != nil {
		res = 0.0
	}
	return res, err
}

func ToInt(value interface{}) (res int64, err error) {
	val := reflect.ValueOf(value)

	switch value.(type) {
	case int, int8, int16, int32, int64:
		res = val.Int()
	case uint, uint8, uint16, uint32, uint64:
		res = int64(val.Uint())
	case string:
		if IsInt(val.String()) {
			res, err = strconv.ParseInt(val.String(), 0, 64)
			if err != nil {
				res = 0
			}
		} else {
			err = fmt.Errorf("math: square root of negative number %g", value)
			res = 0
		}
	default:
		err = fmt.Errorf("math: square root of negative number %g", value)
		res = 0
	}

	return
}

// ToBoolean convert the input string to a boolean.
func ToBoolean(str string) (bool, error) {
	return strconv.ParseBool(str)
}
<commit_msg>Update ToInt comment<commit_after>package govalidator

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
)

// ToString convert the input to a string.
func ToString(obj interface{}) string {
	res := fmt.Sprintf("%v", obj)
	return string(res)
}

// ToJSON convert the input to a valid JSON string
func ToJSON(obj interface{}) (string, error) {
	res, err := json.Marshal(obj)
	if err != nil {
		res = []byte("")
	}
	return string(res), err
}

// ToFloat convert the input string to a float, or 0.0 if the input is not a float.
func ToFloat(str string) (float64, error) {
	res, err := strconv.ParseFloat(str, 64)
	if err != nil {
		res = 0.0
	}
	return res, err
}

// ToInt convert the input string or any int type to an integer type 64, or 0 if the input is not an integer.
func ToInt(value interface{}) (res int64, err error) {
	val := reflect.ValueOf(value)

	switch value.(type) {
	case int, int8, int16, int32, int64:
		res = val.Int()
	case uint, uint8, uint16, uint32, uint64:
		res = int64(val.Uint())
	case string:
		if IsInt(val.String()) {
			res, err = strconv.ParseInt(val.String(), 0, 64)
			if err != nil {
				res = 0
			}
		} else {
			err = fmt.Errorf("math: square root of negative number %g", value)
			res = 0
		}
	default:
		err = fmt.Errorf("math: square root of negative number %g", value)
		res = 0
	}

	return
}

// ToBoolean convert the input string to a boolean.
func ToBoolean(str string) (bool, error) {
	return strconv.ParseBool(str)
}
<|endoftext|>