<commit_before>// Copyright 2016 The etcd 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.

// Based on github.com/grpc-ecosystem/go-grpc-middleware/retry, but modified to support the more
// fine grained error checking required by write-at-most-once retry semantics of etcd.

package clientv3

import (
	"context"
	"io"
	"sync"
	"time"

	"go.etcd.io/etcd/etcdserver/api/v3rpc/rpctypes"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
)

// unaryClientInterceptor returns a new retrying unary client interceptor.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
func (c *Client) unaryClientInterceptor(logger *zap.Logger, optFuncs ...retryOption) grpc.UnaryClientInterceptor {
	intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		grpcOpts, retryOpts := filterCallOptions(opts)
		callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
		// short circuit for simplicity, and avoiding allocations.
		if callOpts.max == 0 {
			return invoker(ctx, method, req, reply, cc, grpcOpts...)
		}
		var lastErr error
		for attempt := uint(0); attempt < callOpts.max; attempt++ {
			if err := waitRetryBackoff(ctx, attempt, callOpts); err != nil {
				return err
			}
			lastErr = invoker(ctx, method, req, reply, cc, grpcOpts...)
			logger.Info("retry unary intercept", zap.Uint("attempt", attempt), zap.Error(lastErr))
			if lastErr == nil {
				return nil
			}
			if isContextError(lastErr) {
				if ctx.Err() != nil {
					// its the context deadline or cancellation.
					return lastErr
				}
				// its the callCtx deadline or cancellation, in which case try again.
				continue
			}
			if callOpts.retryAuth && rpctypes.Error(lastErr) == rpctypes.ErrInvalidAuthToken {
				gterr := c.getToken(ctx)
				if gterr != nil {
					logger.Info("retry failed to fetch new auth token", zap.Error(gterr))
					return lastErr // return the original error for simplicity
				}
				continue
			}
			if !isSafeRetry(c.lg, lastErr, callOpts) {
				return lastErr
			}
		}
		return lastErr
	}
}

// streamClientInterceptor returns a new retrying stream client interceptor for server side streaming calls.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
//
// Retry logic is available *only for ServerStreams*, i.e. 1:n streams, as the internal logic needs
// to buffer the messages sent by the client. If retry is enabled on any other streams (ClientStreams,
// BidiStreams), the retry interceptor will fail the call.
func (c *Client) streamClientInterceptor(logger *zap.Logger, optFuncs ...retryOption) grpc.StreamClientInterceptor {
	intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
		grpcOpts, retryOpts := filterCallOptions(opts)
		callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
		// short circuit for simplicity, and avoiding allocations.
		if callOpts.max == 0 {
			return streamer(ctx, desc, cc, method, grpcOpts...)
		}
		if desc.ClientStreams {
			return nil, grpc.Errorf(codes.Unimplemented, "clientv3/retry_interceptor: cannot retry on ClientStreams, set Disable()")
		}
		newStreamer, err := streamer(ctx, desc, cc, method, grpcOpts...)
		logger.Info("retry stream intercept", zap.Error(err))
		if err != nil {
			// TODO(mwitkow): Maybe dial and transport errors should be retriable?
			return nil, err
		}
		retryingStreamer := &serverStreamingRetryingStream{
			client:       c,
			ClientStream: newStreamer,
			callOpts:     callOpts,
			ctx:          ctx,
			streamerCall: func(ctx context.Context) (grpc.ClientStream, error) {
				return streamer(ctx, desc, cc, method, grpcOpts...)
			},
		}
		return retryingStreamer, nil
	}
}

// type serverStreamingRetryingStream is the implementation of grpc.ClientStream that acts as a
// proxy to the underlying call. If any of the RecvMsg() calls fail, it will try to reestablish
// a new ClientStream according to the retry policy.
type serverStreamingRetryingStream struct {
	grpc.ClientStream
	client        *Client
	bufferedSends []interface{} // single message that the client can sen
	receivedGood  bool          // indicates whether any prior receives were successful
	wasClosedSend bool          // indicates that CloseSend was closed
	ctx           context.Context
	callOpts      *options
	streamerCall  func(ctx context.Context) (grpc.ClientStream, error)
	mu            sync.RWMutex
}

func (s *serverStreamingRetryingStream) setStream(clientStream grpc.ClientStream) {
	s.mu.Lock()
	s.ClientStream = clientStream
	s.mu.Unlock()
}

func (s *serverStreamingRetryingStream) getStream() grpc.ClientStream {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.ClientStream
}

func (s *serverStreamingRetryingStream) SendMsg(m interface{}) error {
	s.mu.Lock()
	s.bufferedSends = append(s.bufferedSends, m)
	s.mu.Unlock()
	return s.getStream().SendMsg(m)
}

func (s *serverStreamingRetryingStream) CloseSend() error {
	s.mu.Lock()
	s.wasClosedSend = true
	s.mu.Unlock()
	return s.getStream().CloseSend()
}

func (s *serverStreamingRetryingStream) Header() (metadata.MD, error) {
	return s.getStream().Header()
}

func (s *serverStreamingRetryingStream) Trailer() metadata.MD {
	return s.getStream().Trailer()
}

func (s *serverStreamingRetryingStream) RecvMsg(m interface{}) error {
	attemptRetry, lastErr := s.receiveMsgAndIndicateRetry(m)
	if !attemptRetry {
		return lastErr // success or hard failure
	}
	// We start off from attempt 1, because zeroth was already made on normal SendMsg().
	for attempt := uint(1); attempt < s.callOpts.max; attempt++ {
		if err := waitRetryBackoff(s.ctx, attempt, s.callOpts); err != nil {
			return err
		}
		newStream, err := s.reestablishStreamAndResendBuffer(s.ctx)
		if err != nil {
			// TODO(mwitkow): Maybe dial and transport errors should be retriable?
			return err
		}
		s.setStream(newStream)
		attemptRetry, lastErr = s.receiveMsgAndIndicateRetry(m)
		//fmt.Printf("Received message and indicate: %v  %v\n", attemptRetry, lastErr)
		if !attemptRetry {
			return lastErr
		}
	}
	return lastErr
}

func (s *serverStreamingRetryingStream) receiveMsgAndIndicateRetry(m interface{}) (bool, error) {
	s.mu.RLock()
	wasGood := s.receivedGood
	s.mu.RUnlock()
	err := s.getStream().RecvMsg(m)
	if err == nil || err == io.EOF {
		s.mu.Lock()
		s.receivedGood = true
		s.mu.Unlock()
		return false, err
	} else if wasGood {
		// previous RecvMsg in the stream succeeded, no retry logic should interfere
		return false, err
	}
	if isContextError(err) {
		if s.ctx.Err() != nil {
			return false, err
		}
		// its the callCtx deadline or cancellation, in which case try again.
		return true, err
	}
	if s.callOpts.retryAuth && rpctypes.Error(err) == rpctypes.ErrInvalidAuthToken {
		gterr := s.client.getToken(s.ctx)
		if gterr != nil {
			s.client.lg.Info("retry failed to fetch new auth token", zap.Error(gterr))
			return false, err // return the original error for simplicity
		}
		return true, err

	}
	return isSafeRetry(s.client.lg, err, s.callOpts), err
}

func (s *serverStreamingRetryingStream) reestablishStreamAndResendBuffer(callCtx context.Context) (grpc.ClientStream, error) {
	s.mu.RLock()
	bufferedSends := s.bufferedSends
	s.mu.RUnlock()
	newStream, err := s.streamerCall(callCtx)
	if err != nil {
		return nil, err
	}
	for _, msg := range bufferedSends {
		if err := newStream.SendMsg(msg); err != nil {
			return nil, err
		}
	}
	if err := newStream.CloseSend(); err != nil {
		return nil, err
	}
	return newStream, nil
}

func waitRetryBackoff(ctx context.Context, attempt uint, callOpts *options) error {
	waitTime := time.Duration(0)
	if attempt > 0 {
		waitTime = callOpts.backoffFunc(attempt)
	}
	if waitTime > 0 {
		timer := time.NewTimer(waitTime)
		select {
		case <-ctx.Done():
			timer.Stop()
			return contextErrToGrpcErr(ctx.Err())
		case <-timer.C:
		}
	}
	return nil
}

// isSafeRetry returns "true", if request is safe for retry with the given error.
func isSafeRetry(lg *zap.Logger, err error, callOpts *options) bool {
	if isContextError(err) {
		return false
	}
	switch callOpts.retryPolicy {
	case repeatable:
		return isSafeRetryImmutableRPC(err)
	case nonRepeatable:
		return isSafeRetryMutableRPC(err)
	default:
		lg.Warn("unrecognized retry policy", zap.String("retryPolicy", callOpts.retryPolicy.String()))
		return false
	}
}

func isContextError(err error) bool {
	return grpc.Code(err) == codes.DeadlineExceeded || grpc.Code(err) == codes.Canceled
}

func contextErrToGrpcErr(err error) error {
	switch err {
	case context.DeadlineExceeded:
		return grpc.Errorf(codes.DeadlineExceeded, err.Error())
	case context.Canceled:
		return grpc.Errorf(codes.Canceled, err.Error())
	default:
		return grpc.Errorf(codes.Unknown, err.Error())
	}
}

var (
	defaultOptions = &options{
		retryPolicy: nonRepeatable,
		max:         0, // disable
		backoffFunc: backoffLinearWithJitter(50*time.Millisecond /*jitter*/, 0.10),
		retryAuth:   true,
	}
)

// backoffFunc denotes a family of functions that control the backoff duration between call retries.
//
// They are called with an identifier of the attempt, and should return a time the system client should
// hold off for. If the time returned is longer than the `context.Context.Deadline` of the request
// the deadline of the request takes precedence and the wait will be interrupted before proceeding
// with the next iteration.
type backoffFunc func(attempt uint) time.Duration

// withRetryPolicy sets the retry policy of this call.
func withRetryPolicy(rp retryPolicy) retryOption {
	return retryOption{applyFunc: func(o *options) {
		o.retryPolicy = rp
	}}
}

// withAuthRetry sets enables authentication retries.
func withAuthRetry(retryAuth bool) retryOption {
	return retryOption{applyFunc: func(o *options) {
		o.retryAuth = retryAuth
	}}
}

// withMax sets the maximum number of retries on this call, or this interceptor.
func withMax(maxRetries uint) retryOption {
	return retryOption{applyFunc: func(o *options) {
		o.max = maxRetries
	}}
}

// WithBackoff sets the `BackoffFunc `used to control time between retries.
func withBackoff(bf backoffFunc) retryOption {
	return retryOption{applyFunc: func(o *options) {
		o.backoffFunc = bf
	}}
}

type options struct {
	retryPolicy retryPolicy
	max         uint
	backoffFunc backoffFunc
	retryAuth   bool
}

// retryOption is a grpc.CallOption that is local to clientv3's retry interceptor.
type retryOption struct {
	grpc.EmptyCallOption // make sure we implement private after() and before() fields so we don't panic.
	applyFunc            func(opt *options)
}

func reuseOrNewWithCallOptions(opt *options, retryOptions []retryOption) *options {
	if len(retryOptions) == 0 {
		return opt
	}
	optCopy := &options{}
	*optCopy = *opt
	for _, f := range retryOptions {
		f.applyFunc(optCopy)
	}
	return optCopy
}

func filterCallOptions(callOptions []grpc.CallOption) (grpcOptions []grpc.CallOption, retryOptions []retryOption) {
	for _, opt := range callOptions {
		if co, ok := opt.(retryOption); ok {
			retryOptions = append(retryOptions, co)
		} else {
			grpcOptions = append(grpcOptions, opt)
		}
	}
	return grpcOptions, retryOptions
}

// BackoffLinearWithJitter waits a set period of time, allowing for jitter (fractional adjustment).
//
// For example waitBetween=1s and jitter=0.10 can generate waits between 900ms and 1100ms.
func backoffLinearWithJitter(waitBetween time.Duration, jitterFraction float64) backoffFunc {
	return func(attempt uint) time.Duration {
		return jitterUp(waitBetween, jitterFraction)
	}
}
<commit_msg>clientv3: clarify retry interceptor logging<commit_after>// Copyright 2016 The etcd 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.

// Based on github.com/grpc-ecosystem/go-grpc-middleware/retry, but modified to support the more
// fine grained error checking required by write-at-most-once retry semantics of etcd.

package clientv3

import (
	"context"
	"io"
	"sync"
	"time"

	"go.etcd.io/etcd/etcdserver/api/v3rpc/rpctypes"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
)

// unaryClientInterceptor returns a new retrying unary client interceptor.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
func (c *Client) unaryClientInterceptor(logger *zap.Logger, optFuncs ...retryOption) grpc.UnaryClientInterceptor {
	intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		grpcOpts, retryOpts := filterCallOptions(opts)
		callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
		// short circuit for simplicity, and avoiding allocations.
		if callOpts.max == 0 {
			return invoker(ctx, method, req, reply, cc, grpcOpts...)
		}
		var lastErr error
		for attempt := uint(0); attempt < callOpts.max; attempt++ {
			if err := waitRetryBackoff(ctx, attempt, callOpts); err != nil {
				return err
			}
			logger.Info(
				"retrying of unary invoker",
				zap.String("target", cc.Target()),
				zap.Uint("attempt", attempt),
			)
			lastErr = invoker(ctx, method, req, reply, cc, grpcOpts...)
			if lastErr == nil {
				return nil
			}
			logger.Warn(
				"retrying of unary invoker failed",
				zap.String("target", cc.Target()),
				zap.Uint("attempt", attempt),
				zap.Error(lastErr),
			)
			if isContextError(lastErr) {
				if ctx.Err() != nil {
					// its the context deadline or cancellation.
					return lastErr
				}
				// its the callCtx deadline or cancellation, in which case try again.
				continue
			}
			if callOpts.retryAuth && rpctypes.Error(lastErr) == rpctypes.ErrInvalidAuthToken {
				gterr := c.getToken(ctx)
				if gterr != nil {
					logger.Warn(
						"retrying of unary invoker failed to fetch new auth token",
						zap.String("target", cc.Target()),
						zap.Error(gterr),
					)
					return lastErr // return the original error for simplicity
				}
				continue
			}
			if !isSafeRetry(c.lg, lastErr, callOpts) {
				return lastErr
			}
		}
		return lastErr
	}
}

// streamClientInterceptor returns a new retrying stream client interceptor for server side streaming calls.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
//
// Retry logic is available *only for ServerStreams*, i.e. 1:n streams, as the internal logic needs
// to buffer the messages sent by the client. If retry is enabled on any other streams (ClientStreams,
// BidiStreams), the retry interceptor will fail the call.
func (c *Client) streamClientInterceptor(logger *zap.Logger, optFuncs ...retryOption) grpc.StreamClientInterceptor {
	intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
		grpcOpts, retryOpts := filterCallOptions(opts)
		callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
		// short circuit for simplicity, and avoiding allocations.
		if callOpts.max == 0 {
			return streamer(ctx, desc, cc, method, grpcOpts...)
		}
		if desc.ClientStreams {
			return nil, grpc.Errorf(codes.Unimplemented, "clientv3/retry_interceptor: cannot retry on ClientStreams, set Disable()")
		}
		newStreamer, err := streamer(ctx, desc, cc, method, grpcOpts...)
		logger.Info("retry stream intercept", zap.Error(err))
		if err != nil {
			// TODO(mwitkow): Maybe dial and transport errors should be retriable?
			return nil, err
		}
		retryingStreamer := &serverStreamingRetryingStream{
			client:       c,
			ClientStream: newStreamer,
			callOpts:     callOpts,
			ctx:          ctx,
			streamerCall: func(ctx context.Context) (grpc.ClientStream, error) {
				return streamer(ctx, desc, cc, method, grpcOpts...)
			},
		}
		return retryingStreamer, nil
	}
}

// type serverStreamingRetryingStream is the implementation of grpc.ClientStream that acts as a
// proxy to the underlying call. If any of the RecvMsg() calls fail, it will try to reestablish
// a new ClientStream according to the retry policy.
type serverStreamingRetryingStream struct {
	grpc.ClientStream
	client        *Client
	bufferedSends []interface{} // single message that the client can sen
	receivedGood  bool          // indicates whether any prior receives were successful
	wasClosedSend bool          // indicates that CloseSend was closed
	ctx           context.Context
	callOpts      *options
	streamerCall  func(ctx context.Context) (grpc.ClientStream, error)
	mu            sync.RWMutex
}

func (s *serverStreamingRetryingStream) setStream(clientStream grpc.ClientStream) {
	s.mu.Lock()
	s.ClientStream = clientStream
	s.mu.Unlock()
}

func (s *serverStreamingRetryingStream) getStream() grpc.ClientStream {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.ClientStream
}

func (s *serverStreamingRetryingStream) SendMsg(m interface{}) error {
	s.mu.Lock()
	s.bufferedSends = append(s.bufferedSends, m)
	s.mu.Unlock()
	return s.getStream().SendMsg(m)
}

func (s *serverStreamingRetryingStream) CloseSend() error {
	s.mu.Lock()
	s.wasClosedSend = true
	s.mu.Unlock()
	return s.getStream().CloseSend()
}

func (s *serverStreamingRetryingStream) Header() (metadata.MD, error) {
	return s.getStream().Header()
}

func (s *serverStreamingRetryingStream) Trailer() metadata.MD {
	return s.getStream().Trailer()
}

func (s *serverStreamingRetryingStream) RecvMsg(m interface{}) error {
	attemptRetry, lastErr := s.receiveMsgAndIndicateRetry(m)
	if !attemptRetry {
		return lastErr // success or hard failure
	}
	// We start off from attempt 1, because zeroth was already made on normal SendMsg().
	for attempt := uint(1); attempt < s.callOpts.max; attempt++ {
		if err := waitRetryBackoff(s.ctx, attempt, s.callOpts); err != nil {
			return err
		}
		newStream, err := s.reestablishStreamAndResendBuffer(s.ctx)
		if err != nil {
			// TODO(mwitkow): Maybe dial and transport errors should be retriable?
			return err
		}
		s.setStream(newStream)
		attemptRetry, lastErr = s.receiveMsgAndIndicateRetry(m)
		//fmt.Printf("Received message and indicate: %v  %v\n", attemptRetry, lastErr)
		if !attemptRetry {
			return lastErr
		}
	}
	return lastErr
}

func (s *serverStreamingRetryingStream) receiveMsgAndIndicateRetry(m interface{}) (bool, error) {
	s.mu.RLock()
	wasGood := s.receivedGood
	s.mu.RUnlock()
	err := s.getStream().RecvMsg(m)
	if err == nil || err == io.EOF {
		s.mu.Lock()
		s.receivedGood = true
		s.mu.Unlock()
		return false, err
	} else if wasGood {
		// previous RecvMsg in the stream succeeded, no retry logic should interfere
		return false, err
	}
	if isContextError(err) {
		if s.ctx.Err() != nil {
			return false, err
		}
		// its the callCtx deadline or cancellation, in which case try again.
		return true, err
	}
	if s.callOpts.retryAuth && rpctypes.Error(err) == rpctypes.ErrInvalidAuthToken {
		gterr := s.client.getToken(s.ctx)
		if gterr != nil {
			s.client.lg.Info("retry failed to fetch new auth token", zap.Error(gterr))
			return false, err // return the original error for simplicity
		}
		return true, err

	}
	return isSafeRetry(s.client.lg, err, s.callOpts), err
}

func (s *serverStreamingRetryingStream) reestablishStreamAndResendBuffer(callCtx context.Context) (grpc.ClientStream, error) {
	s.mu.RLock()
	bufferedSends := s.bufferedSends
	s.mu.RUnlock()
	newStream, err := s.streamerCall(callCtx)
	if err != nil {
		return nil, err
	}
	for _, msg := range bufferedSends {
		if err := newStream.SendMsg(msg); err != nil {
			return nil, err
		}
	}
	if err := newStream.CloseSend(); err != nil {
		return nil, err
	}
	return newStream, nil
}

func waitRetryBackoff(ctx context.Context, attempt uint, callOpts *options) error {
	waitTime := time.Duration(0)
	if attempt > 0 {
		waitTime = callOpts.backoffFunc(attempt)
	}
	if waitTime > 0 {
		timer := time.NewTimer(waitTime)
		select {
		case <-ctx.Done():
			timer.Stop()
			return contextErrToGrpcErr(ctx.Err())
		case <-timer.C:
		}
	}
	return nil
}

// isSafeRetry returns "true", if request is safe for retry with the given error.
func isSafeRetry(lg *zap.Logger, err error, callOpts *options) bool {
	if isContextError(err) {
		return false
	}
	switch callOpts.retryPolicy {
	case repeatable:
		return isSafeRetryImmutableRPC(err)
	case nonRepeatable:
		return isSafeRetryMutableRPC(err)
	default:
		lg.Warn("unrecognized retry policy", zap.String("retryPolicy", callOpts.retryPolicy.String()))
		return false
	}
}

func isContextError(err error) bool {
	return grpc.Code(err) == codes.DeadlineExceeded || grpc.Code(err) == codes.Canceled
}

func contextErrToGrpcErr(err error) error {
	switch err {
	case context.DeadlineExceeded:
		return grpc.Errorf(codes.DeadlineExceeded, err.Error())
	case context.Canceled:
		return grpc.Errorf(codes.Canceled, err.Error())
	default:
		return grpc.Errorf(codes.Unknown, err.Error())
	}
}

var (
	defaultOptions = &options{
		retryPolicy: nonRepeatable,
		max:         0, // disable
		backoffFunc: backoffLinearWithJitter(50*time.Millisecond /*jitter*/, 0.10),
		retryAuth:   true,
	}
)

// backoffFunc denotes a family of functions that control the backoff duration between call retries.
//
// They are called with an identifier of the attempt, and should return a time the system client should
// hold off for. If the time returned is longer than the `context.Context.Deadline` of the request
// the deadline of the request takes precedence and the wait will be interrupted before proceeding
// with the next iteration.
type backoffFunc func(attempt uint) time.Duration

// withRetryPolicy sets the retry policy of this call.
func withRetryPolicy(rp retryPolicy) retryOption {
	return retryOption{applyFunc: func(o *options) {
		o.retryPolicy = rp
	}}
}

// withAuthRetry sets enables authentication retries.
func withAuthRetry(retryAuth bool) retryOption {
	return retryOption{applyFunc: func(o *options) {
		o.retryAuth = retryAuth
	}}
}

// withMax sets the maximum number of retries on this call, or this interceptor.
func withMax(maxRetries uint) retryOption {
	return retryOption{applyFunc: func(o *options) {
		o.max = maxRetries
	}}
}

// WithBackoff sets the `BackoffFunc `used to control time between retries.
func withBackoff(bf backoffFunc) retryOption {
	return retryOption{applyFunc: func(o *options) {
		o.backoffFunc = bf
	}}
}

type options struct {
	retryPolicy retryPolicy
	max         uint
	backoffFunc backoffFunc
	retryAuth   bool
}

// retryOption is a grpc.CallOption that is local to clientv3's retry interceptor.
type retryOption struct {
	grpc.EmptyCallOption // make sure we implement private after() and before() fields so we don't panic.
	applyFunc            func(opt *options)
}

func reuseOrNewWithCallOptions(opt *options, retryOptions []retryOption) *options {
	if len(retryOptions) == 0 {
		return opt
	}
	optCopy := &options{}
	*optCopy = *opt
	for _, f := range retryOptions {
		f.applyFunc(optCopy)
	}
	return optCopy
}

func filterCallOptions(callOptions []grpc.CallOption) (grpcOptions []grpc.CallOption, retryOptions []retryOption) {
	for _, opt := range callOptions {
		if co, ok := opt.(retryOption); ok {
			retryOptions = append(retryOptions, co)
		} else {
			grpcOptions = append(grpcOptions, opt)
		}
	}
	return grpcOptions, retryOptions
}

// BackoffLinearWithJitter waits a set period of time, allowing for jitter (fractional adjustment).
//
// For example waitBetween=1s and jitter=0.10 can generate waits between 900ms and 1100ms.
func backoffLinearWithJitter(waitBetween time.Duration, jitterFraction float64) backoffFunc {
	return func(attempt uint) time.Duration {
		return jitterUp(waitBetween, jitterFraction)
	}
}
<|endoftext|>


<commit_before>package fs

import (
	"fmt"
	"io"
	"os"
	"sync"
	"time"

	"github.com/chzyer/flow"
	"github.com/chzyer/logex"
)

type File struct {
	ref      int32
	refGuard sync.Mutex

	flow      *flow.Flow
	cfg       *FileConfig
	ino       int32
	flags     int
	delegate  FileDelegater
	inodePool *InodePool
	flusher   FileFlusher
	cobuf     *Cobuffer

	flushWaiter sync.WaitGroup
	flushChan   chan struct{}
}

type FileDelegater interface {
	InodePoolDelegate
	ReadData(offset ShortAddr, n int) ([]byte, error)
}

type FileFlusher interface {
	WriteByInode(*InodePool, []byte, chan int)
	Flush()
}

type FileConfig struct {
	Ino           int32
	Name          string
	Flags         int
	Delegate      FileDelegater
	FlushInterval time.Duration
	Flusher       FileFlusher
	FlushSize     int
}

func IsFileCreate(flags int) bool {
	return flags&os.O_CREATE > 0
}

func NewFile(f *flow.Flow, cfg *FileConfig) (*File, error) {
	inodePool := NewInodePool(cfg.Ino, cfg.Delegate)

	if _, err := inodePool.GetLastest(); err != nil {
		if IsFileCreate(cfg.Flags) {
			inodePool.InitInode()
			err = nil
		} else {
			return nil, logex.Trace(err)
		}
	}

	file := &File{
		cfg:       cfg,
		flow:      f.Fork(1),
		ino:       cfg.Ino,
		flags:     cfg.Flags,
		delegate:  cfg.Delegate,
		inodePool: inodePool,
		flusher:   cfg.Flusher,
		cobuf:     NewCobuffer(1<<10, cfg.FlushSize),

		flushChan: make(chan struct{}, 1),
	}
	f.SetOnClose(func() {
		file.Close()
	})
	go file.writeLoop()

	return file, nil
}

func (f *File) Name() string {
	return f.cfg.Name
}

func (f *File) Ino() int32 {
	return f.ino
}

func (f *File) writeLoop() {
	defer f.flow.DoneAndClose()
	var (
		wantFlush bool
		wantClose bool
		timer     <-chan time.Time

		flushReply = make(chan int, 100)
		flushStart = time.Now()
		buffer     []byte
		bufferOps  int
	)

	for {
		select {
		case <-f.cobuf.IsWritten():
			timer = time.NewTimer(f.cfg.FlushInterval).C
			flushStart = time.Now()
			continue
		case <-f.cobuf.IsFlush():
		case <-timer:
		case <-f.flushChan:
			wantFlush = true
		case err := <-flushReply:
			bufferOps -= err
			//if err != nil {
			//	logex.Error("write error:", err)
			//}
			continue
		case <-f.flow.IsClose():
			// println("want close")
			wantClose = true
		}

		Stat.File.FlushDuration.AddNow(flushStart)
		n := f.cobuf.GetData(buffer)
		for n > len(buffer) {
			Stat.File.RegenBuffer.Hit()
			buffer = make([]byte, n)
			n = f.cobuf.GetData(buffer)
		}

		if n > 0 {
			Stat.File.FlushSize.AddBuf(buffer[:n])
			f.flusher.WriteByInode(f.inodePool, buffer[:n], flushReply)
			bufferOps++
		}
		if wantFlush {
			now := time.Now()
			Stat.File.Flush.WaitSize.HitN(bufferOps)
			for bufferOps > 0 {
				bufferOps -= <-flushReply
			}
			Stat.File.Flush.WaitReply.AddNow(now)
			f.flushWaiter.Done()
			wantFlush = false
		}
		if wantClose {
			// println("file: remain ops:", bufferOps)
			// now := time.Now()
			for bufferOps > 0 {
				bufferOps -= <-flushReply
			}
			// println("file wait time:", time.Now().Sub(now).String())
			break
		}
	}
}

func (f *File) Stat() (*Inode, error) {
	return f.inodePool.GetLastest()
}

func (f *File) Size() int64 {
	ino, err := f.inodePool.GetLastest()
	if err != nil {
		panic(err)
	}
	return int64(ino.Start)*BlockSize + int64(ino.Size)
}

func (f *File) ReadAt(b []byte, off int64) (readBytes int, err error) {
	inode, err := f.inodePool.SeekPrev(off)
	if err != nil {
		return 0, err
	}

	var nextInode *Inode

getOffset:
	idx, ok := inode.SeekIdx(off)
	if !ok {
		nextInode, err = f.inodePool.SeekNext(inode)
		if err != nil {
			return
		}
		if inode == nextInode {
			panic("SeekNext() is not working")
		}
		inode = nextInode
		goto getOffset
	}

	remainBytes := inode.GetRemainInBlock(off)
	if remainBytes < len(b)-readBytes {
		remainBytes = len(b) - readBytes
	}
	readAddr := inode.Offsets[idx] + ShortAddr(off&(BlockSize-1))
	data, err := f.delegate.ReadData(readAddr, remainBytes)
	if err != nil {
		return
	}

	if len(data) != remainBytes {
		err = io.EOF
		return
	}

	readBytes += copy(b[readBytes:], data)
	if readBytes == len(b) {
		return
	}
	off += int64(readBytes)

	goto getOffset
}

func (f *File) Write(b []byte) (int, error) {
	if f.flow.IsClosed() {
		return 0, fmt.Errorf("closed")
	}
	f.cobuf.WriteData(b)
	return len(b), nil
}

func (f *File) Sync() {
	f.flushWaiter.Add(1)
	select {
	case f.flushChan <- struct{}{}:
		f.flushWaiter.Wait()
	default:
		f.flushWaiter.Done()
	}

	f.flusher.Flush()
}

func (f *File) AddRef() bool {
	succ := false
	f.refGuard.Lock()
	if f.ref >= 0 {
		f.ref++
		succ = true
	}
	f.refGuard.Unlock()
	return succ
}

func (f *File) Close() error {
	f.refGuard.Lock()
	if f.ref >= 1 {
		f.ref--
	}
	isClose := f.ref == 0
	if isClose {
		f.ref = -1
	}
	f.refGuard.Unlock()

	if !isClose {
		return nil
	}

	if !f.flow.MarkExit() {
		return nil
	}

	now := time.Now()

	f.flow.Close()
	f.cobuf.Close()
	Stat.File.CloseTime.AddNow(now)
	return nil
}
<commit_msg>[fs] file: fix flush logic<commit_after>package fs

import (
	"fmt"
	"io"
	"os"
	"sync"
	"time"

	"github.com/chzyer/flow"
	"github.com/chzyer/logex"
)

type File struct {
	ref      int32
	refGuard sync.Mutex

	flow      *flow.Flow
	cfg       *FileConfig
	ino       int32
	flags     int
	delegate  FileDelegater
	inodePool *InodePool
	flusher   FileFlusher
	cobuf     *Cobuffer

	flushWaiter sync.WaitGroup
	flushChan   chan struct{}
}

type FileDelegater interface {
	InodePoolDelegate
	ReadData(offset ShortAddr, n int) ([]byte, error)
}

type FileFlusher interface {
	WriteByInode(*InodePool, []byte, chan int)
	Flush()
}

type FileConfig struct {
	Ino           int32
	Name          string
	Flags         int
	Delegate      FileDelegater
	FlushInterval time.Duration
	Flusher       FileFlusher
	FlushSize     int
}

func IsFileCreate(flags int) bool {
	return flags&os.O_CREATE > 0
}

func NewFile(f *flow.Flow, cfg *FileConfig) (*File, error) {
	inodePool := NewInodePool(cfg.Ino, cfg.Delegate)

	if _, err := inodePool.GetLastest(); err != nil {
		if IsFileCreate(cfg.Flags) {
			inodePool.InitInode()
			err = nil
		} else {
			return nil, logex.Trace(err)
		}
	}

	file := &File{
		cfg:       cfg,
		flow:      f.Fork(1),
		ino:       cfg.Ino,
		flags:     cfg.Flags,
		delegate:  cfg.Delegate,
		inodePool: inodePool,
		flusher:   cfg.Flusher,
		cobuf:     NewCobuffer(1<<10, cfg.FlushSize),

		flushChan: make(chan struct{}, 1),
	}
	f.SetOnClose(func() {
		file.Close()
	})
	go file.writeLoop()

	return file, nil
}

func (f *File) Name() string {
	return f.cfg.Name
}

func (f *File) Ino() int32 {
	return f.ino
}

func (f *File) writeLoop() {
	defer f.flow.DoneAndClose()
	var (
		wantFlush bool
		wantClose bool
		timer     <-chan time.Time

		flushReply = make(chan int, 100)
		flushStart = time.Now()
		buffer     []byte
		bufferOps  int
	)

	for {
		select {
		case <-f.cobuf.IsWritten():
			timer = time.NewTimer(f.cfg.FlushInterval).C
			flushStart = time.Now()
			continue
		case <-f.cobuf.IsFlush():
		case <-timer:
		case <-f.flushChan:
			wantFlush = true
			f.cobuf.Flush()
			continue
		case err := <-flushReply:
			bufferOps -= err
			//if err != nil {
			//	logex.Error("write error:", err)
			//}
			continue
		case <-f.flow.IsClose():
			// println("want close")
			wantClose = true
		}

		Stat.File.Loop.BufferDuration.AddNow(flushStart)
		n := f.cobuf.GetData(buffer)
		for n > len(buffer) {
			Stat.File.RegenBuffer.Hit()
			buffer = make([]byte, n)
			n = f.cobuf.GetData(buffer)
		}

		if n > 0 {
			Stat.File.FlushSize.AddBuf(buffer[:n])
			f.flusher.WriteByInode(f.inodePool, buffer[:n], flushReply)
			bufferOps++
		}
		if wantFlush {
			now := time.Now()
			Stat.File.Flush.WaitSize.HitN(bufferOps)
			if bufferOps > 0 {
				// println("file: flush them")
				f.flusher.Flush()
				// println("file: done with flush")
			}
			for bufferOps > 0 {
				bufferOps -= <-flushReply
			}
			Stat.File.Flush.WaitReply.AddNow(now)
			f.flushWaiter.Done()
			wantFlush = false
		}
		if wantClose {
			// println("file: remain ops:", bufferOps)
			// now := time.Now()
			for bufferOps > 0 {
				bufferOps -= <-flushReply
			}
			// println("file wait time:", time.Now().Sub(now).String())
			break
		}
	}
}

func (f *File) Stat() (*Inode, error) {
	return f.inodePool.GetLastest()
}

func (f *File) Size() int64 {
	ino, err := f.inodePool.GetLastest()
	if err != nil {
		panic(err)
	}
	return int64(ino.Start)*BlockSize + int64(ino.Size)
}

func (f *File) ReadAt(b []byte, off int64) (readBytes int, err error) {
	inode, err := f.inodePool.SeekPrev(off)
	if err != nil {
		return 0, err
	}

	var nextInode *Inode

getOffset:
	idx, ok := inode.SeekIdx(off)
	if !ok {
		nextInode, err = f.inodePool.SeekNext(inode)
		if err != nil {
			return
		}
		if inode == nextInode {
			panic("SeekNext() is not working")
		}
		inode = nextInode
		goto getOffset
	}

	remainBytes := inode.GetRemainInBlock(off)
	if remainBytes < len(b)-readBytes {
		remainBytes = len(b) - readBytes
	}
	readAddr := inode.Offsets[idx] + ShortAddr(off&(BlockSize-1))
	data, err := f.delegate.ReadData(readAddr, remainBytes)
	if err != nil {
		return
	}

	if len(data) != remainBytes {
		err = io.EOF
		return
	}

	readBytes += copy(b[readBytes:], data)
	if readBytes == len(b) {
		return
	}
	off += int64(readBytes)

	goto getOffset
}

func (f *File) Write(b []byte) (int, error) {
	if f.flow.IsClosed() {
		return 0, fmt.Errorf("closed")
	}
	f.cobuf.WriteData(b)
	return len(b), nil
}

func (f *File) Sync() {
	// println("file: want sync")
	f.flushWaiter.Add(1)
	select {
	case f.flushChan <- struct{}{}:
		f.flushWaiter.Wait()
	default:
		f.flushWaiter.Done()
	}
}

func (f *File) AddRef() bool {
	succ := false
	f.refGuard.Lock()
	if f.ref >= 0 {
		f.ref++
		succ = true
	}
	f.refGuard.Unlock()
	return succ
}

func (f *File) Close() error {
	f.refGuard.Lock()
	if f.ref >= 1 {
		f.ref--
	}
	isClose := f.ref == 0
	if isClose {
		f.ref = -1
	}
	f.refGuard.Unlock()

	if !isClose {
		return nil
	}

	if !f.flow.MarkExit() {
		return nil
	}

	now := time.Now()

	f.flow.Close()
	f.cobuf.Close()
	Stat.File.CloseTime.AddNow(now)
	return nil
}
<|endoftext|>


<commit_before>package godux

import "testing"

func TestStoreStateCount(t *testing.T) {
	store := NewStore()
	store.Setstate("count", 1)
	if store.GetState("count") != 1 {
		t.Error("Expected 1 for Count State value.")
	}

}
func TestActionType(t *testing.T) {
	store := NewStore()
	store.Setstate("count", 1)

	increment := func(number int) Action {
		return Action{
			Type:  "INCREMENT",
			Value: number,
		}
	}

	reducer := func(action Action) interface{} {
		switch action.Type {
		case "INCREMENT":
			return store.GetState("count").(int) + action.Value.(int)
		default:
			return store.GetAllState()
		}
	}
	store.Reducer(reducer)
	value := store.Dispatch(increment(1))
	if value != 2 {
		t.Error("The value are different. Action not correct.")
	}
}
<commit_msg>Fix for issue https://github.com/luisvinicius167/godux/issues/2<commit_after>package godux

import "testing"

func TestStoreStateCount(t *testing.T) {
	store := NewStore()
	store.SetState("count", 1)
	if store.GetState("count") != 1 {
		t.Error("Expected 1 for Count State value.")
	}

}
func TestActionType(t *testing.T) {
	store := NewStore()
	store.SetState("count", 1)

	increment := func(number int) Action {
		return Action{
			Type:  "INCREMENT",
			Value: number,
		}
	}

	reducer := func(action Action) interface{} {
		switch action.Type {
		case "INCREMENT":
			return store.GetState("count").(int) + action.Value.(int)
		default:
			return store.GetAllState()
		}
	}
	store.Reducer(reducer)
	value := store.Dispatch(increment(1))
	if value != 2 {
		t.Error("The value are different. Action not correct.")
	}
}
<|endoftext|>


<commit_before>package godbg

import (
	"testing"
	. "github.com/smartystreets/goconvey/convey"
)

func TestProject(t *testing.T) {
	Convey("Test buffers", t, func() {

	})
}
<commit_msg>Test default buffers<commit_after>package godbg

import (
	"os"
	"testing"
	. "github.com/smartystreets/goconvey/convey"
)

func TestProject(t *testing.T) {
	Convey("Test buffers", t, func() {

		Convey("By Default, equals to std", func() {
			So(Out(), ShouldEqual, os.Stdout)
			So(Err(), ShouldEqual, os.Stderr)
		})
	})
}
<|endoftext|>


<commit_before>package monitor

import (
	"flag"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/docker/leadership"
	"github.com/docker/libkv/store"
	"github.com/docker/libkv/store/zookeeper"
	"github.com/funkygao/gafka"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/reporter"
	"github.com/funkygao/gafka/reporter/influxdb"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/golib/signal"
	log "github.com/funkygao/log4go"
	"github.com/julienschmidt/httprouter"
)

// Monitor is the engine that will start/stop plugin watchers.
// It itself is an implementation of Context.
type Monitor struct {
	influxdbAddr   string
	influxdbDbName string
	apiAddr        string

	router *httprouter.Router
	zkzone *zk.ZkZone

	candidate *leadership.Candidate

	inflight *sync.WaitGroup
	stop     chan struct{} // broadcast to all watchers to stop, but might restart again
	quit     chan struct{}
	leader   bool
}

func (this *Monitor) Init() {
	var logFile, zone string
	flag.StringVar(&logFile, "log", "stdout", "log filename")
	flag.StringVar(&zone, "z", "", "zone, required")
	flag.StringVar(&this.apiAddr, "http", ":10025", "api http server addr")
	flag.StringVar(&this.influxdbAddr, "influxAddr", "", "influxdb addr, required")
	flag.StringVar(&this.influxdbDbName, "db", "", "influxdb db name, required")
	flag.Parse()

	if zone == "" || this.influxdbDbName == "" || this.influxdbAddr == "" {
		panic("run help ")
	}

	ctx.LoadFromHome()
	this.zkzone = zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	this.quit = make(chan struct{})

	// export RESTful api
	this.router = httprouter.New()
	this.router.GET("/metrics", this.metricsHandler)

	if logFile == "stdout" {
		log.AddFilter("stdout", log.INFO, log.NewConsoleLogWriter())
	} else {
		log.DeleteFilter("stdout")

		filer := log.NewFileLogWriter(logFile, true, false, 0644)
		filer.SetFormat("[%d %T] [%L] (%S) %M")
		log.AddFilter("file", log.INFO, filer)
	}

	reporterConfig, err := influxdb.NewConfig(this.influxdbAddr, this.influxdbDbName, "", "", time.Minute)
	if err != nil {
		panic(err)
	}
	reporter.Default = influxdb.New(metrics.DefaultRegistry, reporterConfig)
}

func (this *Monitor) Stop() {
	if this.leader {
		this.leader = false

		log.Info("stopping all watchers ...")
		close(this.stop)

		log.Info("stopping reporter...")
		reporter.Default.Stop()
	}
}

func (this *Monitor) Start() {
	this.leader = true

	this.stop = make(chan struct{})

	go func() {
		if err := reporter.Default.Start(); err != nil {
			log.Error("reporter: %v", err)
		}
	}()

	this.inflight = new(sync.WaitGroup)
	for name, watcherFactory := range registeredWatchers {
		watcher := watcherFactory()
		watcher.Init(this)

		log.Info("created and starting watcher: %s", name)

		this.inflight.Add(1)
		go watcher.Run()
	}

	log.Info("all watchers ready!")

	<-this.stop
	this.inflight.Wait()

	log.Info("all watchers stopped")
}

func (this *Monitor) ServeForever() {
	defer this.zkzone.Close()

	log.Info("kguard[%s] starting...", gafka.BuildId)

	signal.RegisterSignalsHandler(func(sig os.Signal) {
		log.Info("kguard[%s] received signal: %s", gafka.BuildId, strings.ToUpper(sig.String()))

		if this.leader {
			//this.candidate.Resign()
		}
		this.candidate.Stop()
		log.Info("election stopped, stopping watchers...")
		this.Stop()
		close(this.quit)
	}, syscall.SIGINT, syscall.SIGTERM)

	// start the api server
	apiServer := &http.Server{
		Addr:    this.apiAddr,
		Handler: this.router,
	}
	apiListener, err := net.Listen("tcp", this.apiAddr)
	if err == nil {
		go apiServer.Serve(apiListener)
	} else {
		log.Error("api http server: %v", err)
	}

	backend, err := zookeeper.New(this.zkzone.ZkAddrList(), &store.Config{})
	if err != nil {
		panic(err)
	}

	ip, _ := ctx.LocalIP()
	this.candidate = leadership.NewCandidate(backend, zk.KguardLeaderPath, ip.String(), 15*time.Second)
	electedCh, errCh := this.candidate.RunForElection()
	if err != nil {
		panic("Cannot run for election, store is probably down")
	}

	for {
		select {
		case isElected := <-electedCh:
			if isElected {
				log.Info("Won the election, starting all watchers")

				this.Start()
			} else {
				log.Warn("Fails the election, watching election events...")
				this.Stop()
			}

		case err := <-errCh:
			log.Error("Error during election: %v", err)

		case <-this.quit:
			log.Info("kguard[%s] bye!", gafka.BuildId)
			log.Close()
			return
		}
	}

}

func (this *Monitor) ZkZone() *zk.ZkZone {
	return this.zkzone
}

func (this *Monitor) StopChan() <-chan struct{} {
	return this.stop
}

func (this *Monitor) Inflight() *sync.WaitGroup {
	return this.inflight
}

func (this *Monitor) InfluxAddr() string {
	return this.influxdbAddr
}

func (this *Monitor) InfluxDB() string {
	return this.influxdbDbName
}
<commit_msg>add build at footprint<commit_after>package monitor

import (
	"flag"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/docker/leadership"
	"github.com/docker/libkv/store"
	"github.com/docker/libkv/store/zookeeper"
	"github.com/funkygao/gafka"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/reporter"
	"github.com/funkygao/gafka/reporter/influxdb"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/golib/signal"
	log "github.com/funkygao/log4go"
	"github.com/julienschmidt/httprouter"
)

// Monitor is the engine that will start/stop plugin watchers.
// It itself is an implementation of Context.
type Monitor struct {
	influxdbAddr   string
	influxdbDbName string
	apiAddr        string

	router *httprouter.Router
	zkzone *zk.ZkZone

	candidate *leadership.Candidate

	inflight *sync.WaitGroup
	stop     chan struct{} // broadcast to all watchers to stop, but might restart again
	quit     chan struct{}
	leader   bool
}

func (this *Monitor) Init() {
	var logFile, zone string
	flag.StringVar(&logFile, "log", "stdout", "log filename")
	flag.StringVar(&zone, "z", "", "zone, required")
	flag.StringVar(&this.apiAddr, "http", ":10025", "api http server addr")
	flag.StringVar(&this.influxdbAddr, "influxAddr", "", "influxdb addr, required")
	flag.StringVar(&this.influxdbDbName, "db", "", "influxdb db name, required")
	flag.Parse()

	if zone == "" || this.influxdbDbName == "" || this.influxdbAddr == "" {
		panic("run help ")
	}

	ctx.LoadFromHome()
	this.zkzone = zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	this.quit = make(chan struct{})

	// export RESTful api
	this.router = httprouter.New()
	this.router.GET("/metrics", this.metricsHandler)

	if logFile == "stdout" {
		log.AddFilter("stdout", log.INFO, log.NewConsoleLogWriter())
	} else {
		log.DeleteFilter("stdout")

		filer := log.NewFileLogWriter(logFile, true, false, 0644)
		filer.SetFormat("[%d %T] [%L] (%S) %M")
		log.AddFilter("file", log.INFO, filer)
	}

	reporterConfig, err := influxdb.NewConfig(this.influxdbAddr, this.influxdbDbName, "", "", time.Minute)
	if err != nil {
		panic(err)
	}
	reporter.Default = influxdb.New(metrics.DefaultRegistry, reporterConfig)
}

func (this *Monitor) Stop() {
	if this.leader {
		this.leader = false

		log.Info("stopping all watchers ...")
		close(this.stop)

		log.Info("stopping reporter...")
		reporter.Default.Stop()
	}
}

func (this *Monitor) Start() {
	this.leader = true

	this.stop = make(chan struct{})

	go func() {
		if err := reporter.Default.Start(); err != nil {
			log.Error("reporter: %v", err)
		}
	}()

	this.inflight = new(sync.WaitGroup)
	for name, watcherFactory := range registeredWatchers {
		watcher := watcherFactory()
		watcher.Init(this)

		log.Info("created and starting watcher: %s", name)

		this.inflight.Add(1)
		go watcher.Run()
	}

	log.Info("all watchers ready!")

	<-this.stop
	this.inflight.Wait()

	log.Info("all watchers stopped")
}

func (this *Monitor) ServeForever() {
	defer this.zkzone.Close()

	log.Info("kguard[%s@%s] starting...", gafka.BuildId, gafka.BuiltAt)

	signal.RegisterSignalsHandler(func(sig os.Signal) {
		log.Info("kguard[%s@%s] received signal: %s", gafka.BuildId, gafka.BuiltAt, strings.ToUpper(sig.String()))

		if this.leader {
			//this.candidate.Resign()
		}
		this.candidate.Stop()
		log.Info("election stopped, stopping watchers...")
		this.Stop()
		close(this.quit)
	}, syscall.SIGINT, syscall.SIGTERM)

	// start the api server
	apiServer := &http.Server{
		Addr:    this.apiAddr,
		Handler: this.router,
	}
	apiListener, err := net.Listen("tcp", this.apiAddr)
	if err == nil {
		go apiServer.Serve(apiListener)
	} else {
		log.Error("api http server: %v", err)
	}

	backend, err := zookeeper.New(this.zkzone.ZkAddrList(), &store.Config{})
	if err != nil {
		panic(err)
	}

	ip, _ := ctx.LocalIP()
	this.candidate = leadership.NewCandidate(backend, zk.KguardLeaderPath, ip.String(), 15*time.Second)
	electedCh, errCh := this.candidate.RunForElection()
	if err != nil {
		panic("Cannot run for election, store is probably down")
	}

	for {
		select {
		case isElected := <-electedCh:
			if isElected {
				log.Info("Won the election, starting all watchers")

				this.Start()
			} else {
				log.Warn("Fails the election, watching election events...")
				this.Stop()
			}

		case err := <-errCh:
			log.Error("Error during election: %v", err)

		case <-this.quit:
			log.Info("kguard[%s@%s] bye!", gafka.BuildId, gafka.BuiltAt)
			log.Close()
			return
		}
	}

}

func (this *Monitor) ZkZone() *zk.ZkZone {
	return this.zkzone
}

func (this *Monitor) StopChan() <-chan struct{} {
	return this.stop
}

func (this *Monitor) Inflight() *sync.WaitGroup {
	return this.inflight
}

func (this *Monitor) InfluxAddr() string {
	return this.influxdbAddr
}

func (this *Monitor) InfluxDB() string {
	return this.influxdbDbName
}
<|endoftext|>


<commit_before>package monitor

import (
	"flag"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/docker/leadership"
	"github.com/docker/libkv/store"
	"github.com/docker/libkv/store/zookeeper"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/reporter"
	"github.com/funkygao/gafka/reporter/influxdb"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/golib/signal"
	log "github.com/funkygao/log4go"
	"github.com/julienschmidt/httprouter"
)

// Monitor is the engine that will start/stop plugin watchers.
// It itself is an implementation of Context.
type Monitor struct {
	influxdbAddr   string
	influxdbDbName string
	apiAddr        string

	router *httprouter.Router
	zkzone *zk.ZkZone

	candidate *leadership.Candidate

	wg     *sync.WaitGroup
	stop   chan struct{} // broadcast to all watchers to stop, but might restart again
	quit   chan struct{}
	leader bool
}

func (this *Monitor) Init() {
	var logFile, zone string
	flag.StringVar(&logFile, "log", "stdout", "log filename")
	flag.StringVar(&zone, "z", "", "zone, required")
	flag.StringVar(&this.apiAddr, "http", ":10025", "api http server addr")
	flag.StringVar(&this.influxdbAddr, "influxAddr", "", "influxdb addr, required")
	flag.StringVar(&this.influxdbDbName, "db", "", "influxdb db name, required")
	flag.Parse()

	if zone == "" || this.influxdbDbName == "" || this.influxdbAddr == "" {
		panic("run help ")
	}

	ctx.LoadFromHome()
	this.zkzone = zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	this.quit = make(chan struct{})

	// export RESTful api
	this.router = httprouter.New()
	this.router.GET("/metrics", this.metricsHandler)

	if logFile == "stdout" {
		log.AddFilter("stdout", log.INFO, log.NewConsoleLogWriter())
	} else {
		log.DeleteFilter("stdout")

		filer := log.NewFileLogWriter(logFile, true, false, 0644)
		filer.SetFormat("[%d %T] [%L] (%S) %M")
		log.AddFilter("file", log.INFO, filer)
	}

	reporterConfig, err := influxdb.NewConfig(this.influxdbAddr, this.influxdbDbName, "", "", time.Minute)
	if err != nil {
		panic(err)
	}
	reporter.Default = influxdb.New(metrics.DefaultRegistry, reporterConfig)
}

func (this *Monitor) Stop() {
	if this.leader {
		log.Info("stopping all watchers ...")
		close(this.stop)

		log.Info("stopping reporter...")
		reporter.Default.Stop()
	}
	this.leader = false
}

func (this *Monitor) Start() {
	this.leader = true

	this.stop = make(chan struct{})

	go func() {
		if err := reporter.Default.Start(); err != nil {
			log.Error("reporter: %v", err)
		}
	}()

	this.wg = new(sync.WaitGroup)
	for name, watcherFactory := range registeredWatchers {
		watcher := watcherFactory()
		watcher.Init(this)

		log.Info("created and starting watcher: %s", name)

		this.wg.Add(1)
		go watcher.Run()
	}

	log.Info("all watchers ready!")

	<-this.stop
	this.wg.Wait()

	log.Info("all watchers stopped")
}

func (this *Monitor) ServeForever() {
	defer this.zkzone.Close()

	log.Info("starting...")

	signal.RegisterSignalsHandler(func(sig os.Signal) {
		log.Info("received signal: %s", strings.ToUpper(sig.String()))

		if this.leader {
			//this.candidate.Resign()
		}
		this.candidate.Stop()
		log.Info("election stopped, stopping watchers...")
		this.Stop()
		close(this.quit)
	}, syscall.SIGINT, syscall.SIGTERM)

	// start the api server
	apiServer := &http.Server{
		Addr:    this.apiAddr,
		Handler: this.router,
	}
	apiListener, err := net.Listen("tcp", this.apiAddr)
	if err == nil {
		go apiServer.Serve(apiListener)
	} else {
		log.Error("api http server: %v", err)
	}

	backend, err := zookeeper.New(this.zkzone.ZkAddrList(), &store.Config{})
	if err != nil {
		panic(err)
	}

	ip, _ := ctx.LocalIP()
	this.candidate = leadership.NewCandidate(backend, zk.KguardLeaderPath, ip.String(), 15*time.Second)
	electedCh, errCh := this.candidate.RunForElection()
	if err != nil {
		panic("Cannot run for election, store is probably down")
	}

	for {
		select {
		case isElected := <-electedCh:
			if isElected {
				log.Info("Won the election, starting all watchers")

				this.Start()
			} else {
				log.Warn("Fails the election, watching election events...")
				this.Stop()
			}

		case err := <-errCh:
			log.Error("Error during election: %v", err)

		case <-this.quit:
			log.Info("kguard bye!")
			log.Close()
			return
		}
	}

}

func (this *Monitor) ZkZone() *zk.ZkZone {
	return this.zkzone
}

func (this *Monitor) StopChan() <-chan struct{} {
	return this.stop
}

func (this *Monitor) WaitGroup() *sync.WaitGroup {
	return this.wg
}
<commit_msg>fix 'close a closed channel' issue<commit_after>package monitor

import (
	"flag"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/docker/leadership"
	"github.com/docker/libkv/store"
	"github.com/docker/libkv/store/zookeeper"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/reporter"
	"github.com/funkygao/gafka/reporter/influxdb"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/golib/signal"
	log "github.com/funkygao/log4go"
	"github.com/julienschmidt/httprouter"
)

// Monitor is the engine that will start/stop plugin watchers.
// It itself is an implementation of Context.
type Monitor struct {
	influxdbAddr   string
	influxdbDbName string
	apiAddr        string

	router *httprouter.Router
	zkzone *zk.ZkZone

	candidate *leadership.Candidate

	wg     *sync.WaitGroup
	stop   chan struct{} // broadcast to all watchers to stop, but might restart again
	quit   chan struct{}
	leader bool
}

func (this *Monitor) Init() {
	var logFile, zone string
	flag.StringVar(&logFile, "log", "stdout", "log filename")
	flag.StringVar(&zone, "z", "", "zone, required")
	flag.StringVar(&this.apiAddr, "http", ":10025", "api http server addr")
	flag.StringVar(&this.influxdbAddr, "influxAddr", "", "influxdb addr, required")
	flag.StringVar(&this.influxdbDbName, "db", "", "influxdb db name, required")
	flag.Parse()

	if zone == "" || this.influxdbDbName == "" || this.influxdbAddr == "" {
		panic("run help ")
	}

	ctx.LoadFromHome()
	this.zkzone = zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	this.quit = make(chan struct{})

	// export RESTful api
	this.router = httprouter.New()
	this.router.GET("/metrics", this.metricsHandler)

	if logFile == "stdout" {
		log.AddFilter("stdout", log.INFO, log.NewConsoleLogWriter())
	} else {
		log.DeleteFilter("stdout")

		filer := log.NewFileLogWriter(logFile, true, false, 0644)
		filer.SetFormat("[%d %T] [%L] (%S) %M")
		log.AddFilter("file", log.INFO, filer)
	}

	reporterConfig, err := influxdb.NewConfig(this.influxdbAddr, this.influxdbDbName, "", "", time.Minute)
	if err != nil {
		panic(err)
	}
	reporter.Default = influxdb.New(metrics.DefaultRegistry, reporterConfig)
}

func (this *Monitor) Stop() {
	if this.leader {
		this.leader = false

		log.Info("stopping all watchers ...")
		close(this.stop)

		log.Info("stopping reporter...")
		reporter.Default.Stop()
	}
}

func (this *Monitor) Start() {
	this.leader = true

	this.stop = make(chan struct{})

	go func() {
		if err := reporter.Default.Start(); err != nil {
			log.Error("reporter: %v", err)
		}
	}()

	this.wg = new(sync.WaitGroup)
	for name, watcherFactory := range registeredWatchers {
		watcher := watcherFactory()
		watcher.Init(this)

		log.Info("created and starting watcher: %s", name)

		this.wg.Add(1)
		go watcher.Run()
	}

	log.Info("all watchers ready!")

	<-this.stop
	this.wg.Wait()

	log.Info("all watchers stopped")
}

func (this *Monitor) ServeForever() {
	defer this.zkzone.Close()

	log.Info("starting...")

	signal.RegisterSignalsHandler(func(sig os.Signal) {
		log.Info("received signal: %s", strings.ToUpper(sig.String()))

		if this.leader {
			//this.candidate.Resign()
		}
		this.candidate.Stop()
		log.Info("election stopped, stopping watchers...")
		this.Stop()
		close(this.quit)
	}, syscall.SIGINT, syscall.SIGTERM)

	// start the api server
	apiServer := &http.Server{
		Addr:    this.apiAddr,
		Handler: this.router,
	}
	apiListener, err := net.Listen("tcp", this.apiAddr)
	if err == nil {
		go apiServer.Serve(apiListener)
	} else {
		log.Error("api http server: %v", err)
	}

	backend, err := zookeeper.New(this.zkzone.ZkAddrList(), &store.Config{})
	if err != nil {
		panic(err)
	}

	ip, _ := ctx.LocalIP()
	this.candidate = leadership.NewCandidate(backend, zk.KguardLeaderPath, ip.String(), 15*time.Second)
	electedCh, errCh := this.candidate.RunForElection()
	if err != nil {
		panic("Cannot run for election, store is probably down")
	}

	for {
		select {
		case isElected := <-electedCh:
			if isElected {
				log.Info("Won the election, starting all watchers")

				this.Start()
			} else {
				log.Warn("Fails the election, watching election events...")
				this.Stop()
			}

		case err := <-errCh:
			log.Error("Error during election: %v", err)

		case <-this.quit:
			log.Info("kguard bye!")
			log.Close()
			return
		}
	}

}

func (this *Monitor) ZkZone() *zk.ZkZone {
	return this.zkzone
}

func (this *Monitor) StopChan() <-chan struct{} {
	return this.stop
}

func (this *Monitor) WaitGroup() *sync.WaitGroup {
	return this.wg
}
<|endoftext|>


<commit_before>package queue

import (
	"bytes"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/StackExchange/tcollector/opentsdb"
)

var l = log.New(os.Stdout, "", log.LstdFlags)

type Queue struct {
	sync.Mutex
	host  string
	queue opentsdb.MultiDataPoint
	c     chan *opentsdb.DataPoint
}

// Creates and starts a new Queue.
func New(host string, c chan *opentsdb.DataPoint) *Queue {
	q := Queue{
		host: host,
		c:    c,
	}

	go func() {
		for dp := range c {
			q.Lock()
			q.queue = append(q.queue, dp)
			q.Unlock()
		}
	}()
	go q.send()
	return &q
}

const MAX_PERSEC = 50

func (q *Queue) send() {
	for {
		if len(q.queue) > 0 {
			q.Lock()
			i := len(q.queue)
			if i > MAX_PERSEC {
				i = MAX_PERSEC
			}
			sending := q.queue[:i]
			q.queue = q.queue[i:]
			go q.sendBatch(sending)
			q.Unlock()
		} else {
			time.Sleep(time.Second)
		}
	}
}

func (q *Queue) sendBatch(batch opentsdb.MultiDataPoint) {
	l.Println("sending", len(batch))
	b, err := batch.Json()
	if err != nil {
		l.Println(err)
		// bad JSON encoding, just give up
		return
	}
	resp, err := http.Post(q.host, "application/json", bytes.NewReader(b))
	if err != nil {
		l.Println(err)
		goto Err
	}
	if resp.StatusCode != http.StatusNoContent {
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			l.Println(err)
		} else if len(body) > 0 {
			l.Println(string(body))
		}
		goto Err
	}
	return
Err:
	l.Println("error, restoring", len(batch))
	for _, dp := range batch {
		q.c <- dp
	}
}
<commit_msg>cmd/scollector: More useful TSDB error prints<commit_after>package queue

import (
	"bytes"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/StackExchange/tcollector/opentsdb"
)

var l = log.New(os.Stdout, "", log.LstdFlags)

type Queue struct {
	sync.Mutex
	host  string
	queue opentsdb.MultiDataPoint
	c     chan *opentsdb.DataPoint
}

// Creates and starts a new Queue.
func New(host string, c chan *opentsdb.DataPoint) *Queue {
	q := Queue{
		host: host,
		c:    c,
	}

	go func() {
		for dp := range c {
			q.Lock()
			q.queue = append(q.queue, dp)
			q.Unlock()
		}
	}()
	go q.send()
	return &q
}

const MAX_PERSEC = 50

func (q *Queue) send() {
	for {
		if len(q.queue) > 0 {
			q.Lock()
			i := len(q.queue)
			if i > MAX_PERSEC {
				i = MAX_PERSEC
			}
			sending := q.queue[:i]
			q.queue = q.queue[i:]
			go q.sendBatch(sending)
			q.Unlock()
		} else {
			time.Sleep(time.Second)
		}
	}
}

func (q *Queue) sendBatch(batch opentsdb.MultiDataPoint) {
	l.Println("sending", len(batch))
	b, err := batch.Json()
	if err != nil {
		l.Println(err)
		// bad JSON encoding, just give up
		return
	}
	resp, err := http.Post(q.host, "application/json", bytes.NewReader(b))
	if err != nil {
		l.Println(err)
		goto Err
	}
	if resp.StatusCode != http.StatusNoContent {
		l.Println("RESP ERR", resp.Status)
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			l.Println("ERR", err)
		}
		if len(body) > 0 {
			l.Println("ERR BODY", string(body))
		}
		l.Println("REQ BODY", string(b))
		goto Err
	}
	return
Err:
	l.Println("error, restoring", len(batch))
	for _, dp := range batch {
		q.c <- dp
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// go command is not available on android

// +build !android

package main

import (
	"fmt"
	"go/build"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

// This file contains a test that compiles and runs each program in testdata
// after generating the string method for its type. The rule is that for testdata/x.go
// we run stringer -type X and then compile and run the program. The resulting
// binary panics if the String method for X is not correct, including for error cases.

func TestEndToEnd(t *testing.T) {
	dir, err := ioutil.TempDir("", "stringer")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)
	// Create stringer in temporary directory.
	stringer := filepath.Join(dir, "stringer.exe")
	err = run("go", "build", "-o", stringer)
	if err != nil {
		t.Fatalf("building stringer: %s", err)
	}
	// Read the testdata directory.
	fd, err := os.Open("testdata")
	if err != nil {
		t.Fatal(err)
	}
	defer fd.Close()
	names, err := fd.Readdirnames(-1)
	if err != nil {
		t.Fatalf("Readdirnames: %s", err)
	}
	// Generate, compile, and run the test programs.
	for _, name := range names {
		if !strings.HasSuffix(name, ".go") {
			t.Errorf("%s is not a Go file", name)
			continue
		}
		if name == "cgo.go" && !build.Default.CgoEnabled {
			t.Logf("cgo is no enabled for %s", name)
			continue
		}
		// Names are known to be ASCII and long enough.
		typeName := fmt.Sprintf("%c%s", name[0]+'A'-'a', name[1:len(name)-len(".go")])
		stringerCompileAndRun(t, dir, stringer, typeName, name)
	}
}

// stringerCompileAndRun runs stringer for the named file and compiles and
// runs the target binary in directory dir. That binary will panic if the String method is incorrect.
func stringerCompileAndRun(t *testing.T, dir, stringer, typeName, fileName string) {
	t.Logf("run: %s %s\n", fileName, typeName)
	source := filepath.Join(dir, fileName)
	err := copy(source, filepath.Join("testdata", fileName))
	if err != nil {
		t.Fatalf("copying file to temporary directory: %s", err)
	}
	stringSource := filepath.Join(dir, typeName+"_string.go")
	// Run stringer in temporary directory.
	err = run(stringer, "-type", typeName, "-output", stringSource, source)
	if err != nil {
		t.Fatal(err)
	}
	// Run the binary in the temporary directory.
	err = run("go", "run", stringSource, source)
	if err != nil {
		t.Fatal(err)
	}
}

// copy copies the from file to the to file.
func copy(to, from string) error {
	toFd, err := os.Create(to)
	if err != nil {
		return err
	}
	defer toFd.Close()
	fromFd, err := os.Open(from)
	if err != nil {
		return err
	}
	defer fromFd.Close()
	_, err = io.Copy(toFd, fromFd)
	return err
}

// run runs a single command and returns an error if it does not succeed.
// os/exec should have this function, to be honest.
func run(name string, arg ...string) error {
	cmd := exec.Command(name, arg...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}
<commit_msg>cmd/stringer: fix typo in test log message ("no" -> "not")<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// go command is not available on android

// +build !android

package main

import (
	"fmt"
	"go/build"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

// This file contains a test that compiles and runs each program in testdata
// after generating the string method for its type. The rule is that for testdata/x.go
// we run stringer -type X and then compile and run the program. The resulting
// binary panics if the String method for X is not correct, including for error cases.

func TestEndToEnd(t *testing.T) {
	dir, err := ioutil.TempDir("", "stringer")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)
	// Create stringer in temporary directory.
	stringer := filepath.Join(dir, "stringer.exe")
	err = run("go", "build", "-o", stringer)
	if err != nil {
		t.Fatalf("building stringer: %s", err)
	}
	// Read the testdata directory.
	fd, err := os.Open("testdata")
	if err != nil {
		t.Fatal(err)
	}
	defer fd.Close()
	names, err := fd.Readdirnames(-1)
	if err != nil {
		t.Fatalf("Readdirnames: %s", err)
	}
	// Generate, compile, and run the test programs.
	for _, name := range names {
		if !strings.HasSuffix(name, ".go") {
			t.Errorf("%s is not a Go file", name)
			continue
		}
		if name == "cgo.go" && !build.Default.CgoEnabled {
			t.Logf("cgo is not enabled for %s", name)
			continue
		}
		// Names are known to be ASCII and long enough.
		typeName := fmt.Sprintf("%c%s", name[0]+'A'-'a', name[1:len(name)-len(".go")])
		stringerCompileAndRun(t, dir, stringer, typeName, name)
	}
}

// stringerCompileAndRun runs stringer for the named file and compiles and
// runs the target binary in directory dir. That binary will panic if the String method is incorrect.
func stringerCompileAndRun(t *testing.T, dir, stringer, typeName, fileName string) {
	t.Logf("run: %s %s\n", fileName, typeName)
	source := filepath.Join(dir, fileName)
	err := copy(source, filepath.Join("testdata", fileName))
	if err != nil {
		t.Fatalf("copying file to temporary directory: %s", err)
	}
	stringSource := filepath.Join(dir, typeName+"_string.go")
	// Run stringer in temporary directory.
	err = run(stringer, "-type", typeName, "-output", stringSource, source)
	if err != nil {
		t.Fatal(err)
	}
	// Run the binary in the temporary directory.
	err = run("go", "run", stringSource, source)
	if err != nil {
		t.Fatal(err)
	}
}

// copy copies the from file to the to file.
func copy(to, from string) error {
	toFd, err := os.Create(to)
	if err != nil {
		return err
	}
	defer toFd.Close()
	fromFd, err := os.Open(from)
	if err != nil {
		return err
	}
	defer fromFd.Close()
	_, err = io.Copy(toFd, fromFd)
	return err
}

// run runs a single command and returns an error if it does not succeed.
// os/exec should have this function, to be honest.
func run(name string, arg ...string) error {
	cmd := exec.Command(name, arg...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 The Syncthing Authors.
//
// 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 main

import (
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"runtime"
	"sort"
	"time"

	"github.com/syncthing/syncthing/lib/config"
	"github.com/syncthing/syncthing/lib/dialer"
	"github.com/syncthing/syncthing/lib/model"
	"github.com/syncthing/syncthing/lib/protocol"
	"github.com/syncthing/syncthing/lib/upgrade"
	"github.com/thejerf/suture"
)

// Current version number of the usage report, for acceptance purposes. If
// fields are added or changed this integer must be incremented so that users
// are prompted for acceptance of the new report.
const usageReportVersion = 2

type usageReportingManager struct {
	cfg   *config.Wrapper
	model *model.Model
	sup   *suture.Supervisor
}

func newUsageReportingManager(cfg *config.Wrapper, m *model.Model) *usageReportingManager {
	mgr := &usageReportingManager{
		cfg:   cfg,
		model: m,
	}

	// Start UR if it's enabled.
	mgr.CommitConfiguration(config.Configuration{}, cfg.Raw())

	// Listen to future config changes so that we can start and stop as
	// appropriate.
	cfg.Subscribe(mgr)

	return mgr
}

func (m *usageReportingManager) VerifyConfiguration(from, to config.Configuration) error {
	return nil
}

func (m *usageReportingManager) CommitConfiguration(from, to config.Configuration) bool {
	if to.Options.URAccepted >= usageReportVersion && m.sup == nil {
		// Usage reporting was turned on; lets start it.
		svc := newUsageReportingService(m.cfg, m.model)
		m.sup = suture.NewSimple("usageReporting")
		m.sup.Add(svc)
		m.sup.ServeBackground()
	} else if to.Options.URAccepted < usageReportVersion && m.sup != nil {
		// Usage reporting was turned off
		m.sup.Stop()
		m.sup = nil
	}

	return true
}

func (m *usageReportingManager) String() string {
	return fmt.Sprintf("usageReportingManager@%p", m)
}

// reportData returns the data to be sent in a usage report. It's used in
// various places, so not part of the usageReportingSvc object.
func reportData(cfg *config.Wrapper, m *model.Model) map[string]interface{} {
	res := make(map[string]interface{})
	res["urVersion"] = usageReportVersion
	res["uniqueID"] = cfg.Options().URUniqueID
	res["version"] = Version
	res["longVersion"] = LongVersion
	res["platform"] = runtime.GOOS + "-" + runtime.GOARCH
	res["numFolders"] = len(cfg.Folders())
	res["numDevices"] = len(cfg.Devices())

	var totFiles, maxFiles int
	var totBytes, maxBytes int64
	for folderID := range cfg.Folders() {
		files, _, bytes := m.GlobalSize(folderID)
		totFiles += files
		totBytes += bytes
		if files > maxFiles {
			maxFiles = files
		}
		if bytes > maxBytes {
			maxBytes = bytes
		}
	}

	res["totFiles"] = totFiles
	res["folderMaxFiles"] = maxFiles
	res["totMiB"] = totBytes / 1024 / 1024
	res["folderMaxMiB"] = maxBytes / 1024 / 1024

	var mem runtime.MemStats
	runtime.ReadMemStats(&mem)
	res["memoryUsageMiB"] = (mem.Sys - mem.HeapReleased) / 1024 / 1024
	res["sha256Perf"] = cpuBench(5, 125*time.Millisecond)

	bytes, err := memorySize()
	if err == nil {
		res["memorySize"] = bytes / 1024 / 1024
	}
	res["numCPU"] = runtime.NumCPU()

	var rescanIntvs []int
	folderUses := map[string]int{
		"readonly":      0,
		"ignorePerms":   0,
		"ignoreDelete":  0,
		"autoNormalize": 0,
	}
	for _, cfg := range cfg.Folders() {
		rescanIntvs = append(rescanIntvs, cfg.RescanIntervalS)

		if cfg.ReadOnly {
			folderUses["readonly"]++
		}
		if cfg.IgnorePerms {
			folderUses["ignorePerms"]++
		}
		if cfg.IgnoreDelete {
			folderUses["ignoreDelete"]++
		}
		if cfg.AutoNormalize {
			folderUses["autoNormalize"]++
		}
	}
	sort.Ints(rescanIntvs)
	res["rescanIntvs"] = rescanIntvs
	res["folderUses"] = folderUses

	deviceUses := map[string]int{
		"introducer":       0,
		"customCertName":   0,
		"compressAlways":   0,
		"compressMetadata": 0,
		"compressNever":    0,
		"dynamicAddr":      0,
		"staticAddr":       0,
	}
	for _, cfg := range cfg.Devices() {
		if cfg.Introducer {
			deviceUses["introducer"]++
		}
		if cfg.CertName != "" && cfg.CertName != "syncthing" {
			deviceUses["customCertName"]++
		}
		if cfg.Compression == protocol.CompressAlways {
			deviceUses["compressAlways"]++
		} else if cfg.Compression == protocol.CompressMetadata {
			deviceUses["compressMetadata"]++
		} else if cfg.Compression == protocol.CompressNever {
			deviceUses["compressNever"]++
		}
		for _, addr := range cfg.Addresses {
			if addr == "dynamic" {
				deviceUses["dynamicAddr"]++
			} else {
				deviceUses["staticAddr"]++
			}
		}
	}
	res["deviceUses"] = deviceUses

	defaultAnnounceServersDNS, defaultAnnounceServersIP, otherAnnounceServers := 0, 0, 0
	for _, addr := range cfg.Options().GlobalAnnServers {
		if addr == "default" {
			defaultAnnounceServersDNS++
		} else if stringIn(addr, config.DefaultDiscoveryServersIP) {
			defaultAnnounceServersIP++
		} else {
			otherAnnounceServers++
		}
	}
	res["announce"] = map[string]interface{}{
		"globalEnabled":     cfg.Options().GlobalAnnEnabled,
		"localEnabled":      cfg.Options().LocalAnnEnabled,
		"defaultServersDNS": defaultAnnounceServersDNS,
		"defaultServersIP":  defaultAnnounceServersIP,
		"otherServers":      otherAnnounceServers,
	}

	defaultRelayServers, otherRelayServers := 0, 0
	for _, addr := range cfg.Options().RelayServers {
		switch addr {
		case "dynamic+https://relays.syncthing.net":
			defaultRelayServers++
		default:
			otherRelayServers++
		}
	}
	res["relays"] = map[string]interface{}{
		"enabled":        cfg.Options().RelaysEnabled,
		"defaultServers": defaultRelayServers,
		"otherServers":   otherRelayServers,
	}

	res["usesRateLimit"] = cfg.Options().MaxRecvKbps > 0 || cfg.Options().MaxSendKbps > 0

	res["upgradeAllowedManual"] = !(upgrade.DisabledByCompilation || noUpgrade)
	res["upgradeAllowedAuto"] = !(upgrade.DisabledByCompilation || noUpgrade) && cfg.Options().AutoUpgradeIntervalH > 0

	return res
}

func stringIn(needle string, haystack []string) bool {
	for _, s := range haystack {
		if needle == s {
			return true
		}
	}
	return false
}

type usageReportingService struct {
	cfg   *config.Wrapper
	model *model.Model
	stop  chan struct{}
}

func newUsageReportingService(cfg *config.Wrapper, model *model.Model) *usageReportingService {
	return &usageReportingService{
		cfg:   cfg,
		model: model,
		stop:  make(chan struct{}),
	}
}

func (s *usageReportingService) sendUsageReport() error {
	d := reportData(s.cfg, s.model)
	var b bytes.Buffer
	json.NewEncoder(&b).Encode(d)

	transp := &http.Transport{}
	client := &http.Client{Transport: transp}
	if BuildEnv == "android" {
		// This works around the lack of DNS resolution on Android... :(
		transp.Dial = func(network, addr string) (net.Conn, error) {
			return dialer.Dial(network, "194.126.249.13:443")
		}
	}

	if s.cfg.Options().URPostInsecurely {
		transp.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}
	_, err := client.Post(s.cfg.Options().URURL, "application/json", &b)
	return err
}

func (s *usageReportingService) Serve() {
	s.stop = make(chan struct{})

	l.Infoln("Starting usage reporting")
	defer l.Infoln("Stopping usage reporting")

	t := time.NewTimer(time.Duration(s.cfg.Options().URInitialDelayS) * time.Second) // time to initial report at start
	for {
		select {
		case <-s.stop:
			return
		case <-t.C:
			err := s.sendUsageReport()
			if err != nil {
				l.Infoln("Usage report:", err)
			}
			t.Reset(24 * time.Hour) // next report tomorrow
		}
	}
}

func (s *usageReportingService) Stop() {
	close(s.stop)
}

// cpuBench returns CPU performance as a measure of single threaded SHA-256 MiB/s
func cpuBench(iterations int, duration time.Duration) float64 {
	var perf float64
	for i := 0; i < iterations; i++ {
		if v := cpuBenchOnce(duration); v > perf {
			perf = v
		}
	}
	return perf
}

func cpuBenchOnce(duration time.Duration) float64 {
	chunkSize := 100 * 1 << 10
	h := sha256.New()
	bs := make([]byte, chunkSize)
	rand.Reader.Read(bs)

	t0 := time.Now()
	b := 0
	for time.Since(t0) < duration {
		h.Write(bs)
		b += chunkSize
	}
	h.Sum(nil)
	d := time.Since(t0)
	return float64(int(float64(b)/d.Seconds()/(1<<20)*100)) / 100
}
<commit_msg>Correctly report the default relay server in usage stats<commit_after>// Copyright (C) 2014 The Syncthing Authors.
//
// 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 main

import (
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"runtime"
	"sort"
	"time"

	"github.com/syncthing/syncthing/lib/config"
	"github.com/syncthing/syncthing/lib/dialer"
	"github.com/syncthing/syncthing/lib/model"
	"github.com/syncthing/syncthing/lib/protocol"
	"github.com/syncthing/syncthing/lib/upgrade"
	"github.com/thejerf/suture"
)

// Current version number of the usage report, for acceptance purposes. If
// fields are added or changed this integer must be incremented so that users
// are prompted for acceptance of the new report.
const usageReportVersion = 2

type usageReportingManager struct {
	cfg   *config.Wrapper
	model *model.Model
	sup   *suture.Supervisor
}

func newUsageReportingManager(cfg *config.Wrapper, m *model.Model) *usageReportingManager {
	mgr := &usageReportingManager{
		cfg:   cfg,
		model: m,
	}

	// Start UR if it's enabled.
	mgr.CommitConfiguration(config.Configuration{}, cfg.Raw())

	// Listen to future config changes so that we can start and stop as
	// appropriate.
	cfg.Subscribe(mgr)

	return mgr
}

func (m *usageReportingManager) VerifyConfiguration(from, to config.Configuration) error {
	return nil
}

func (m *usageReportingManager) CommitConfiguration(from, to config.Configuration) bool {
	if to.Options.URAccepted >= usageReportVersion && m.sup == nil {
		// Usage reporting was turned on; lets start it.
		svc := newUsageReportingService(m.cfg, m.model)
		m.sup = suture.NewSimple("usageReporting")
		m.sup.Add(svc)
		m.sup.ServeBackground()
	} else if to.Options.URAccepted < usageReportVersion && m.sup != nil {
		// Usage reporting was turned off
		m.sup.Stop()
		m.sup = nil
	}

	return true
}

func (m *usageReportingManager) String() string {
	return fmt.Sprintf("usageReportingManager@%p", m)
}

// reportData returns the data to be sent in a usage report. It's used in
// various places, so not part of the usageReportingSvc object.
func reportData(cfg *config.Wrapper, m *model.Model) map[string]interface{} {
	res := make(map[string]interface{})
	res["urVersion"] = usageReportVersion
	res["uniqueID"] = cfg.Options().URUniqueID
	res["version"] = Version
	res["longVersion"] = LongVersion
	res["platform"] = runtime.GOOS + "-" + runtime.GOARCH
	res["numFolders"] = len(cfg.Folders())
	res["numDevices"] = len(cfg.Devices())

	var totFiles, maxFiles int
	var totBytes, maxBytes int64
	for folderID := range cfg.Folders() {
		files, _, bytes := m.GlobalSize(folderID)
		totFiles += files
		totBytes += bytes
		if files > maxFiles {
			maxFiles = files
		}
		if bytes > maxBytes {
			maxBytes = bytes
		}
	}

	res["totFiles"] = totFiles
	res["folderMaxFiles"] = maxFiles
	res["totMiB"] = totBytes / 1024 / 1024
	res["folderMaxMiB"] = maxBytes / 1024 / 1024

	var mem runtime.MemStats
	runtime.ReadMemStats(&mem)
	res["memoryUsageMiB"] = (mem.Sys - mem.HeapReleased) / 1024 / 1024
	res["sha256Perf"] = cpuBench(5, 125*time.Millisecond)

	bytes, err := memorySize()
	if err == nil {
		res["memorySize"] = bytes / 1024 / 1024
	}
	res["numCPU"] = runtime.NumCPU()

	var rescanIntvs []int
	folderUses := map[string]int{
		"readonly":      0,
		"ignorePerms":   0,
		"ignoreDelete":  0,
		"autoNormalize": 0,
	}
	for _, cfg := range cfg.Folders() {
		rescanIntvs = append(rescanIntvs, cfg.RescanIntervalS)

		if cfg.ReadOnly {
			folderUses["readonly"]++
		}
		if cfg.IgnorePerms {
			folderUses["ignorePerms"]++
		}
		if cfg.IgnoreDelete {
			folderUses["ignoreDelete"]++
		}
		if cfg.AutoNormalize {
			folderUses["autoNormalize"]++
		}
	}
	sort.Ints(rescanIntvs)
	res["rescanIntvs"] = rescanIntvs
	res["folderUses"] = folderUses

	deviceUses := map[string]int{
		"introducer":       0,
		"customCertName":   0,
		"compressAlways":   0,
		"compressMetadata": 0,
		"compressNever":    0,
		"dynamicAddr":      0,
		"staticAddr":       0,
	}
	for _, cfg := range cfg.Devices() {
		if cfg.Introducer {
			deviceUses["introducer"]++
		}
		if cfg.CertName != "" && cfg.CertName != "syncthing" {
			deviceUses["customCertName"]++
		}
		if cfg.Compression == protocol.CompressAlways {
			deviceUses["compressAlways"]++
		} else if cfg.Compression == protocol.CompressMetadata {
			deviceUses["compressMetadata"]++
		} else if cfg.Compression == protocol.CompressNever {
			deviceUses["compressNever"]++
		}
		for _, addr := range cfg.Addresses {
			if addr == "dynamic" {
				deviceUses["dynamicAddr"]++
			} else {
				deviceUses["staticAddr"]++
			}
		}
	}
	res["deviceUses"] = deviceUses

	defaultAnnounceServersDNS, defaultAnnounceServersIP, otherAnnounceServers := 0, 0, 0
	for _, addr := range cfg.Options().GlobalAnnServers {
		if addr == "default" {
			defaultAnnounceServersDNS++
		} else if stringIn(addr, config.DefaultDiscoveryServersIP) {
			defaultAnnounceServersIP++
		} else {
			otherAnnounceServers++
		}
	}
	res["announce"] = map[string]interface{}{
		"globalEnabled":     cfg.Options().GlobalAnnEnabled,
		"localEnabled":      cfg.Options().LocalAnnEnabled,
		"defaultServersDNS": defaultAnnounceServersDNS,
		"defaultServersIP":  defaultAnnounceServersIP,
		"otherServers":      otherAnnounceServers,
	}

	defaultRelayServers, otherRelayServers := 0, 0
	for _, addr := range cfg.Options().RelayServers {
		switch addr {
		case "dynamic+https://relays.syncthing.net/endpoint":
			defaultRelayServers++
		default:
			otherRelayServers++
		}
	}
	res["relays"] = map[string]interface{}{
		"enabled":        cfg.Options().RelaysEnabled,
		"defaultServers": defaultRelayServers,
		"otherServers":   otherRelayServers,
	}

	res["usesRateLimit"] = cfg.Options().MaxRecvKbps > 0 || cfg.Options().MaxSendKbps > 0

	res["upgradeAllowedManual"] = !(upgrade.DisabledByCompilation || noUpgrade)
	res["upgradeAllowedAuto"] = !(upgrade.DisabledByCompilation || noUpgrade) && cfg.Options().AutoUpgradeIntervalH > 0

	return res
}

func stringIn(needle string, haystack []string) bool {
	for _, s := range haystack {
		if needle == s {
			return true
		}
	}
	return false
}

type usageReportingService struct {
	cfg   *config.Wrapper
	model *model.Model
	stop  chan struct{}
}

func newUsageReportingService(cfg *config.Wrapper, model *model.Model) *usageReportingService {
	return &usageReportingService{
		cfg:   cfg,
		model: model,
		stop:  make(chan struct{}),
	}
}

func (s *usageReportingService) sendUsageReport() error {
	d := reportData(s.cfg, s.model)
	var b bytes.Buffer
	json.NewEncoder(&b).Encode(d)

	transp := &http.Transport{}
	client := &http.Client{Transport: transp}
	if BuildEnv == "android" {
		// This works around the lack of DNS resolution on Android... :(
		transp.Dial = func(network, addr string) (net.Conn, error) {
			return dialer.Dial(network, "194.126.249.13:443")
		}
	}

	if s.cfg.Options().URPostInsecurely {
		transp.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}
	_, err := client.Post(s.cfg.Options().URURL, "application/json", &b)
	return err
}

func (s *usageReportingService) Serve() {
	s.stop = make(chan struct{})

	l.Infoln("Starting usage reporting")
	defer l.Infoln("Stopping usage reporting")

	t := time.NewTimer(time.Duration(s.cfg.Options().URInitialDelayS) * time.Second) // time to initial report at start
	for {
		select {
		case <-s.stop:
			return
		case <-t.C:
			err := s.sendUsageReport()
			if err != nil {
				l.Infoln("Usage report:", err)
			}
			t.Reset(24 * time.Hour) // next report tomorrow
		}
	}
}

func (s *usageReportingService) Stop() {
	close(s.stop)
}

// cpuBench returns CPU performance as a measure of single threaded SHA-256 MiB/s
func cpuBench(iterations int, duration time.Duration) float64 {
	var perf float64
	for i := 0; i < iterations; i++ {
		if v := cpuBenchOnce(duration); v > perf {
			perf = v
		}
	}
	return perf
}

func cpuBenchOnce(duration time.Duration) float64 {
	chunkSize := 100 * 1 << 10
	h := sha256.New()
	bs := make([]byte, chunkSize)
	rand.Reader.Read(bs)

	t0 := time.Now()
	b := 0
	for time.Since(t0) < duration {
		h.Write(bs)
		b += chunkSize
	}
	h.Sum(nil)
	d := time.Since(t0)
	return float64(int(float64(b)/d.Seconds()/(1<<20)*100)) / 100
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"gopkg.in/natefinch/lumberjack.v2"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	corev1client "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog/v2"
)

func main() {
	os.Exit(run())
}

func run() int {
	terminationLog := flag.String("termination-log-file", "", "Write logs after SIGTERM to this file (in addition to stderr)")
	terminationLock := flag.String("termination-touch-file", "", "Touch this file on SIGTERM and delete on termination")
	kubeconfigPath := flag.String("kubeconfig", "", "Optional kubeconfig used to create events")

	klog.InitFlags(nil)
	flag.Set("v", "9")

	// never log to stderr, only through our termination log writer (which sends it also to stderr)
	flag.Set("logtostderr", "false")
	flag.Set("stderrthreshold", "99")

	flag.Parse()
	args := flag.CommandLine.Args()

	if len(args) == 0 {
		fmt.Println("Missing command line")
		return 1
	}

	// use special tee-like writer when termination log is set
	termCh := make(chan struct{})
	var stderr io.Writer = os.Stderr
	if len(*terminationLog) > 0 {
		stderr = &terminationFileWriter{
			Writer:             os.Stderr,
			fn:                 *terminationLog,
			startFileLoggingCh: termCh,
		}

		// do the klog file writer dance: klog writes to all outputs of lower
		// severity. No idea why. So we discard for anything other than info.
		// Otherwise, we would see errors multiple times.
		klog.SetOutput(ioutil.Discard)
		klog.SetOutputBySeverity("INFO", stderr)
	}

	var client kubernetes.Interface
	if len(*kubeconfigPath) > 0 {
		loader := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(&clientcmd.ClientConfigLoadingRules{ExplicitPath: *kubeconfigPath}, &clientcmd.ConfigOverrides{})
		if cfg, err := loader.ClientConfig(); err != nil {
			klog.Errorf("failed to load kubeconfig %q: %v", *kubeconfigPath, err)
			return 1
		} else {
			client = kubernetes.NewForConfigOrDie(cfg)
		}
	}

	// touch file early. If the file is not removed on termination, we are not
	// terminating cleanly via SIGTERM.
	if len(*terminationLock) > 0 {
		ref, err := eventReference()
		if err != nil {
			klog.Errorf("failed to get event target: %v", err)
			return 1
		}

		if st, err := os.Stat(*terminationLock); err == nil {
			klog.Warningf("Previous pod did not terminate gracefully: %s", st.ModTime().String())
			if client != nil {
				go wait.PollUntil(5*time.Second, func() (bool, error) {
					if err := eventf(client.CoreV1().Events(ref.Namespace), *ref, corev1.EventTypeWarning, "NonGracefulTermination", "Previous pod did not terminate gracefully: %s", st.ModTime().String()); err != nil {
						return false, nil
					}

					select {
					case <-termCh:
					default:
						klog.Infof("Deleting old termination lock file %q", *terminationLock)
						os.Remove(*terminationLock)
					}
					return true, nil
				}, termCh)
			}
		}
		klog.Infof("Touching termination lock file %q", *terminationLock)
		if err := touch(*terminationLock); err != nil {
			klog.Infof("error touching %s: %v", *terminationLock, err)
			// keep going
		}
		defer func() {
			klog.Infof("Deleting termination lock file %q", *terminationLock)
			os.Remove(*terminationLock)
		}()
	}

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = stderr

	// forward SIGTERM and SIGINT to child
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for s := range sigCh {
			select {
			case <-termCh:
			default:
				close(termCh)
			}

			klog.Infof("Received signal %s. Forwarding to sub-process %q.", s, args[0])

			cmd.Process.Signal(s)
		}
	}()

	klog.Infof("Launching sub-process %q", cmd)
	rc := 0
	if err := cmd.Run(); err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			rc = exitError.ExitCode()
		} else {
			klog.Infof("Failed to launch %s: %v", args[0], err)
			return 255
		}
	}

	// remove signal handling
	signal.Stop(sigCh)
	close(sigCh)
	wg.Wait()

	klog.Infof("Termination finished with exit code %d", rc)
	return rc
}

// terminationFileWriter forwards everything to the embedded writer. When
// startFileLoggingCh is closed, everything is appended to the given file name
// in addition.
type terminationFileWriter struct {
	io.Writer
	fn                 string
	startFileLoggingCh <-chan struct{}

	logger io.Writer
}

func (w *terminationFileWriter) Write(bs []byte) (int, error) {
	// temporary hack to avoid logging sensitive tokens.
	// TODO: drop when we moved to a non-sensitive storage format
	if strings.Contains(string(bs), "URI=\"/apis/oauth.openshift.io/v1/oauthaccesstokens/") || strings.Contains(string(bs), "URI=\"/apis/oauth.openshift.io/v1/oauthauthorizetokens/") {
		return len(bs), nil
	}

	select {
	case <-w.startFileLoggingCh:
		if w.logger == nil {
			l := &lumberjack.Logger{
				Filename:   w.fn,
				MaxSize:    100,
				MaxBackups: 3,
				MaxAge:     28,
				Compress:   false,
			}
			w.logger = l
			fmt.Fprintf(os.Stderr, "Copying termination logs to %q\n", w.fn)
		}
		if n, err := w.logger.Write(bs); err != nil {
			return n, err
		} else if n != len(bs) {
			return n, io.ErrShortWrite
		}
	default:
	}

	return w.Writer.Write(bs)
}

func touch(fn string) error {
	_, err := os.Stat(fn)
	if os.IsNotExist(err) {
		file, err := os.Create(fn)
		if err != nil {
			return err
		}
		defer file.Close()
		return nil
	}

	currentTime := time.Now().Local()
	return os.Chtimes(fn, currentTime, currentTime)
}

func eventf(client corev1client.EventInterface, ref corev1.ObjectReference, eventType, reason, messageFmt string, args ...interface{}) error {
	t := metav1.Time{Time: time.Now()}
	host, _ := os.Hostname() // expicitly ignore error. Empty host is fine

	e := &corev1.Event{
		ObjectMeta: metav1.ObjectMeta{
			Name:      fmt.Sprintf("%v.%x", ref.Name, t.UnixNano()),
			Namespace: ref.Namespace,
		},
		InvolvedObject: ref,
		Reason:         reason,
		Message:        fmt.Sprintf(messageFmt, args...),
		FirstTimestamp: t,
		LastTimestamp:  t,
		Count:          1,
		Type:           eventType,
		Source:         corev1.EventSource{Component: "apiserver", Host: host},
	}

	_, err := client.Create(context.TODO(), e, metav1.CreateOptions{})

	if err == nil {
		klog.V(2).Infof("Event(%#v): type: '%v' reason: '%v' %v", e.InvolvedObject, e.Type, e.Reason, e.Message)
	}

	return err
}

func eventReference() (*corev1.ObjectReference, error) {
	ns := os.Getenv("POD_NAMESPACE")
	pod := os.Getenv("POD_NAME")
	if len(ns) == 0 && len(pod) > 0 {
		serviceAccountNamespaceFile := "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
		if _, err := os.Stat(serviceAccountNamespaceFile); err == nil {
			bs, err := ioutil.ReadFile(serviceAccountNamespaceFile)
			if err != nil {
				return nil, err
			}
			ns = string(bs)
		}
	}
	if len(ns) == 0 {
		pod = ""
		ns = "kube-system"
	}
	if len(pod) == 0 {
		return &corev1.ObjectReference{
			Kind:       "Namespace",
			Name:       ns,
			APIVersion: "v1",
		}, nil
	}

	return &corev1.ObjectReference{
		Kind:       "Pod",
		Namespace:  ns,
		Name:       pod,
		APIVersion: "v1",
	}, nil
}
<commit_msg>UPSTREAM: <carry>: watch-termination: fix deletion race and write non-graceful message also to termination.log<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"gopkg.in/natefinch/lumberjack.v2"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	corev1client "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog/v2"
)

func main() {
	os.Exit(run())
}

func run() int {
	terminationLog := flag.String("termination-log-file", "", "Write logs after SIGTERM to this file (in addition to stderr)")
	terminationLock := flag.String("termination-touch-file", "", "Touch this file on SIGTERM and delete on termination")
	kubeconfigPath := flag.String("kubeconfig", "", "Optional kubeconfig used to create events")

	klog.InitFlags(nil)
	flag.Set("v", "9")

	// never log to stderr, only through our termination log writer (which sends it also to stderr)
	flag.Set("logtostderr", "false")
	flag.Set("stderrthreshold", "99")

	flag.Parse()
	args := flag.CommandLine.Args()

	if len(args) == 0 {
		fmt.Println("Missing command line")
		return 1
	}

	// use special tee-like writer when termination log is set
	termCh := make(chan struct{})
	var stderr io.Writer = os.Stderr
	var terminationLogger *terminationFileWriter
	if len(*terminationLog) > 0 {
		terminationLogger = &terminationFileWriter{
			Writer:             os.Stderr,
			fn:                 *terminationLog,
			startFileLoggingCh: termCh,
		}
		stderr = terminationLogger

		// do the klog file writer dance: klog writes to all outputs of lower
		// severity. No idea why. So we discard for anything other than info.
		// Otherwise, we would see errors multiple times.
		klog.SetOutput(ioutil.Discard)
		klog.SetOutputBySeverity("INFO", stderr)
	}

	var client kubernetes.Interface
	if len(*kubeconfigPath) > 0 {
		loader := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(&clientcmd.ClientConfigLoadingRules{ExplicitPath: *kubeconfigPath}, &clientcmd.ConfigOverrides{})
		if cfg, err := loader.ClientConfig(); err != nil {
			klog.Errorf("failed to load kubeconfig %q: %v", *kubeconfigPath, err)
			return 1
		} else {
			client = kubernetes.NewForConfigOrDie(cfg)
		}
	}

	// touch file early. If the file is not removed on termination, we are not
	// terminating cleanly via SIGTERM.
	if len(*terminationLock) > 0 {
		ref, err := eventReference()
		if err != nil {
			klog.Errorf("failed to get event target: %v", err)
			return 1
		}

		if st, err := os.Stat(*terminationLock); err == nil {
			podName := "unknown"
			if v := os.Getenv("POD_NAME"); len(v) > 0 {
				podName = v // pod name is always the same for static pods
			}
			msg := fmt.Sprintf("Previous pod %s started at %s did not terminate gracefully", podName, st.ModTime().String())

			klog.Warning(msg)
			_, _ = terminationLogger.WriteToTerminationLog([]byte(msg + "\n"))

			if client != nil {
				go wait.PollUntil(5*time.Second, func() (bool, error) {
					if err := eventf(client.CoreV1().Events(ref.Namespace), *ref, corev1.EventTypeWarning, "NonGracefulTermination", msg); err != nil {
						return false, nil
					}

					select {
					case <-termCh:
					default:
					}
					return true, nil
				}, termCh)
			}

			klog.Infof("Deleting old termination lock file %q", *terminationLock)
			if err := os.Remove(*terminationLock); err != nil {
				klog.Errorf("Old termination lock file deletion failed: %v", err)
			}
		}

		// separation to see where the new one is starting
		_, _ = terminationLogger.WriteToTerminationLog([]byte("---\n"))

		klog.Infof("Touching termination lock file %q", *terminationLock)
		if err := touch(*terminationLock); err != nil {
			klog.Infof("Error touching %s: %v", *terminationLock, err)
			// keep going
		}

		defer func() {
			klog.Infof("Deleting termination lock file %q", *terminationLock)
			if err := os.Remove(*terminationLock); err != nil {
				klog.Errorf("Termination lock file deletion failed: %v", err)
			}
		}()
	}

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = stderr

	// forward SIGTERM and SIGINT to child
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for s := range sigCh {
			select {
			case <-termCh:
			default:
				close(termCh)
			}

			klog.Infof("Received signal %s. Forwarding to sub-process %q.", s, args[0])

			cmd.Process.Signal(s)
		}
	}()

	klog.Infof("Launching sub-process %q", cmd)
	rc := 0
	if err := cmd.Run(); err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			rc = exitError.ExitCode()
		} else {
			klog.Infof("Failed to launch %s: %v", args[0], err)
			return 255
		}
	}

	// remove signal handling
	signal.Stop(sigCh)
	close(sigCh)
	wg.Wait()

	klog.Infof("Termination finished with exit code %d", rc)
	return rc
}

// terminationFileWriter forwards everything to the embedded writer. When
// startFileLoggingCh is closed, everything is appended to the given file name
// in addition.
type terminationFileWriter struct {
	io.Writer
	fn                 string
	startFileLoggingCh <-chan struct{}

	logger io.Writer
}

func (w *terminationFileWriter) WriteToTerminationLog(bs []byte) (int, error) {
	if w == nil {
		return len(bs), nil
	}

	if w.logger == nil {
		l := &lumberjack.Logger{
			Filename:   w.fn,
			MaxSize:    100,
			MaxBackups: 3,
			MaxAge:     28,
			Compress:   false,
		}
		w.logger = l
		fmt.Fprintf(os.Stderr, "Copying termination logs to %q\n", w.fn)
	}
	if n, err := w.logger.Write(bs); err != nil {
		return n, err
	} else if n != len(bs) {
		return n, io.ErrShortWrite
	}
	return len(bs), nil
}

func (w *terminationFileWriter) Write(bs []byte) (int, error) {
	// temporary hack to avoid logging sensitive tokens.
	// TODO: drop when we moved to a non-sensitive storage format
	if strings.Contains(string(bs), "URI=\"/apis/oauth.openshift.io/v1/oauthaccesstokens/") || strings.Contains(string(bs), "URI=\"/apis/oauth.openshift.io/v1/oauthauthorizetokens/") {
		return len(bs), nil
	}

	select {
	case <-w.startFileLoggingCh:
		if n, err := w.WriteToTerminationLog(bs); err != nil {
			return n, err
		}
	default:
	}

	return w.Writer.Write(bs)
}

func touch(fn string) error {
	_, err := os.Stat(fn)
	if os.IsNotExist(err) {
		file, err := os.Create(fn)
		if err != nil {
			return err
		}
		defer file.Close()
		return nil
	}

	currentTime := time.Now().Local()
	return os.Chtimes(fn, currentTime, currentTime)
}

func eventf(client corev1client.EventInterface, ref corev1.ObjectReference, eventType, reason, messageFmt string, args ...interface{}) error {
	t := metav1.Time{Time: time.Now()}
	host, _ := os.Hostname() // expicitly ignore error. Empty host is fine

	e := &corev1.Event{
		ObjectMeta: metav1.ObjectMeta{
			Name:      fmt.Sprintf("%v.%x", ref.Name, t.UnixNano()),
			Namespace: ref.Namespace,
		},
		InvolvedObject: ref,
		Reason:         reason,
		Message:        fmt.Sprintf(messageFmt, args...),
		FirstTimestamp: t,
		LastTimestamp:  t,
		Count:          1,
		Type:           eventType,
		Source:         corev1.EventSource{Component: "apiserver", Host: host},
	}

	_, err := client.Create(context.TODO(), e, metav1.CreateOptions{})

	if err == nil {
		klog.V(2).Infof("Event(%#v): type: '%v' reason: '%v' %v", e.InvolvedObject, e.Type, e.Reason, e.Message)
	}

	return err
}

func eventReference() (*corev1.ObjectReference, error) {
	ns := os.Getenv("POD_NAMESPACE")
	pod := os.Getenv("POD_NAME")
	if len(ns) == 0 && len(pod) > 0 {
		serviceAccountNamespaceFile := "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
		if _, err := os.Stat(serviceAccountNamespaceFile); err == nil {
			bs, err := ioutil.ReadFile(serviceAccountNamespaceFile)
			if err != nil {
				return nil, err
			}
			ns = string(bs)
		}
	}
	if len(ns) == 0 {
		pod = ""
		ns = "kube-system"
	}
	if len(pod) == 0 {
		return &corev1.ObjectReference{
			Kind:       "Namespace",
			Name:       ns,
			APIVersion: "v1",
		}, nil
	}

	return &corev1.ObjectReference{
		Kind:       "Pod",
		Namespace:  ns,
		Name:       pod,
		APIVersion: "v1",
	}, nil
}
<|endoftext|>