<commit_before>package ebs

import (
	"context"
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/giantswarm/apiextensions/pkg/apis/provider/v1alpha1"
	"github.com/giantswarm/backoff"
	"github.com/giantswarm/microerror"
	"github.com/giantswarm/micrologger"

	"github.com/giantswarm/aws-operator/service/controller/v25/key"
)

const (
	cloudProviderClusterTagValue = "owned"
)

type Config struct {
	Client EC2Client
	Logger micrologger.Logger
}

type EBS struct {
	client EC2Client
	logger micrologger.Logger
}

func New(config Config) (*EBS, error) {
	if config.Client == nil {
		return nil, microerror.Maskf(invalidConfigError, "%T.Client must not be empty", config)
	}
	if config.Logger == nil {
		return nil, microerror.Maskf(invalidConfigError, "%T.Logger must not be empty", config)
	}

	e := &EBS{
		client: config.Client,
		logger: config.Logger,
	}

	return e, nil
}

// DeleteVolume deletes an EBS volume with retry logic.
func (e *EBS) DeleteVolume(ctx context.Context, volumeID string) error {
	e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("deleting EBS volume %s", volumeID))

	o := func() error {
		_, err := e.client.DeleteVolume(&ec2.DeleteVolumeInput{
			VolumeId: aws.String(volumeID),
		})
		if IsVolumeNotFound(err) {
			// Fall through.
			return nil
		} else if err != nil {
			return microerror.Mask(err)
		}

		return nil
	}
	b := backoff.NewExponential(30*time.Second, 5*time.Second)
	n := backoff.NewNotifier(e.logger, context.Background())
	err := backoff.RetryNotify(o, b, n)
	if err != nil {
		return microerror.Mask(err)
	}

	e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("deleted EBS volume %s", volumeID))

	return nil
}

// DetachVolume detaches an EBS volume. If force is specified data loss may occur. If shutdown is
// specified the instance will be stopped first.
func (e *EBS) DetachVolume(ctx context.Context, volumeID string, attachment VolumeAttachment, force bool, shutdown bool, wait bool) error {
	if shutdown {
		e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("stopping instance %s", attachment.InstanceID))

		i := &ec2.StopInstancesInput{
			InstanceIds: []*string{
				aws.String(attachment.InstanceID),
			},
		}
		_, err := e.client.StopInstances(i)
		if err != nil {
			return microerror.Mask(err)
		}

		if wait {
			e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("waiting for instance %s being stopped", attachment.InstanceID))

			i := &ec2.DescribeInstancesInput{
				InstanceIds: []*string{
					aws.String(attachment.InstanceID),
				},
			}
			err := e.client.WaitUntilInstanceStopped(i)
			if err != nil {
				return microerror.Mask(err)
			}

			e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("waited for instance %s being stopped", attachment.InstanceID))
		}

		e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("stopped instance %s", attachment.InstanceID))
	}

	e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("detaching EBS volume %s from instance %s", volumeID, attachment.InstanceID))

	_, err := e.client.DetachVolume(&ec2.DetachVolumeInput{
		Device:     aws.String(attachment.Device),
		InstanceId: aws.String(attachment.InstanceID),
		VolumeId:   aws.String(volumeID),
		Force:      aws.Bool(force),
	})
	if err != nil {
		return microerror.Mask(err)
	}

	e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("detached EBS volume %s from instance %s", volumeID, attachment.InstanceID))

	return nil
}

// ListVolumes lists EBS volumes for a guest cluster. If etcdVolume is true
// the Etcd volume for the master instance will be returned. If persistentVolume
// is set then any Persistent Volumes associated with the cluster will be
// returned.
func (e *EBS) ListVolumes(customObject v1alpha1.AWSConfig, filterFuncs ...func(t *ec2.Tag) bool) ([]Volume, error) {
	var volumes []Volume

	// We filter to only select clusters with the cluster cloud provider tag.
	i := &ec2.DescribeVolumesInput{
		Filters: []*ec2.Filter{
			{
				Name: aws.String(fmt.Sprintf("tag:%s", key.ClusterCloudProviderTag(customObject))),
				Values: []*string{
					aws.String(cloudProviderClusterTagValue),
				},
			},
		},
	}
	o, err := e.client.DescribeVolumes(i)
	if err != nil {
		return nil, microerror.Mask(err)
	}

	for _, v := range o.Volumes {
		if !IsFiltered(v, filterFuncs) {
			continue
		}

		attachments := []VolumeAttachment{}

		if len(v.Attachments) > 0 {
			for _, a := range v.Attachments {
				attachments = append(attachments, VolumeAttachment{
					Device:     *a.Device,
					InstanceID: *a.InstanceId,
				})
			}
		}

		volume := Volume{
			VolumeID:    *v.VolumeId,
			Attachments: attachments,
		}

		volumes = append(volumes, volume)
	}

	return volumes, nil
}
<commit_msg>refactor EBS service (#1509)<commit_after>package ebs

import (
	"context"
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/giantswarm/apiextensions/pkg/apis/provider/v1alpha1"
	"github.com/giantswarm/backoff"
	"github.com/giantswarm/microerror"
	"github.com/giantswarm/micrologger"

	"github.com/giantswarm/aws-operator/service/controller/v25/key"
)

const (
	cloudProviderClusterTagValue = "owned"
)

type Config struct {
	Client EC2Client
	Logger micrologger.Logger
}

type EBS struct {
	client EC2Client
	logger micrologger.Logger
}

func New(config Config) (*EBS, error) {
	if config.Client == nil {
		return nil, microerror.Maskf(invalidConfigError, "%T.Client must not be empty", config)
	}
	if config.Logger == nil {
		return nil, microerror.Maskf(invalidConfigError, "%T.Logger must not be empty", config)
	}

	e := &EBS{
		client: config.Client,
		logger: config.Logger,
	}

	return e, nil
}

// DeleteVolume deletes an EBS volume with retry logic.
func (e *EBS) DeleteVolume(ctx context.Context, volumeID string) error {
	e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("deleting EBS volume %#q", volumeID))

	o := func() error {
		i := &ec2.DeleteVolumeInput{
			VolumeId: aws.String(volumeID),
		}

		_, err := e.client.DeleteVolume(i)
		if IsVolumeNotFound(err) {
			// Fall through.
			return nil
		} else if err != nil {
			return microerror.Mask(err)
		}

		return nil
	}
	b := backoff.NewExponential(30*time.Second, 5*time.Second)
	n := backoff.NewNotifier(e.logger, context.Background())

	err := backoff.RetryNotify(o, b, n)
	if err != nil {
		return microerror.Mask(err)
	}

	e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("deleted EBS volume %#q", volumeID))

	return nil
}

// DetachVolume detaches an EBS volume. If force is specified data loss may occur. If shutdown is
// specified the instance will be stopped first.
func (e *EBS) DetachVolume(ctx context.Context, volumeID string, attachment VolumeAttachment, force bool, shutdown bool, wait bool) error {
	if shutdown {
		e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("requesting to stop instance %#q", attachment.InstanceID))

		i := &ec2.StopInstancesInput{
			InstanceIds: []*string{
				aws.String(attachment.InstanceID),
			},
		}

		_, err := e.client.StopInstances(i)
		if err != nil {
			return microerror.Mask(err)
		}

		e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("requested to stop instance %#q", attachment.InstanceID))
	}

	if shutdown && wait {
		e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("waiting for instance %#q to stop", attachment.InstanceID))

		i := &ec2.DescribeInstancesInput{
			InstanceIds: []*string{
				aws.String(attachment.InstanceID),
			},
		}

		err := e.client.WaitUntilInstanceStopped(i)
		if err != nil {
			return microerror.Mask(err)
		}

		e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("waited for instance %#q to stop", attachment.InstanceID))
	}

	{
		e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("detaching EBS volume %#q from instance %#q", volumeID, attachment.InstanceID))

		i := &ec2.DetachVolumeInput{
			Device:     aws.String(attachment.Device),
			InstanceId: aws.String(attachment.InstanceID),
			VolumeId:   aws.String(volumeID),
			Force:      aws.Bool(force),
		}

		_, err := e.client.DetachVolume(i)
		if err != nil {
			return microerror.Mask(err)
		}

		e.logger.LogCtx(ctx, "level", "debug", "message", fmt.Sprintf("detached EBS volume %#q from instance %#q", volumeID, attachment.InstanceID))
	}

	return nil
}

// ListVolumes lists EBS volumes for a guest cluster. If etcdVolume is true
// the Etcd volume for the master instance will be returned. If persistentVolume
// is set then any Persistent Volumes associated with the cluster will be
// returned.
func (e *EBS) ListVolumes(customObject v1alpha1.AWSConfig, filterFuncs ...func(t *ec2.Tag) bool) ([]Volume, error) {
	var volumes []Volume

	// We filter to only select clusters with the cluster cloud provider tag.
	i := &ec2.DescribeVolumesInput{
		Filters: []*ec2.Filter{
			{
				Name: aws.String(fmt.Sprintf("tag:%s", key.ClusterCloudProviderTag(customObject))),
				Values: []*string{
					aws.String(cloudProviderClusterTagValue),
				},
			},
		},
	}

	o, err := e.client.DescribeVolumes(i)
	if err != nil {
		return nil, microerror.Mask(err)
	}

	for _, v := range o.Volumes {
		if !IsFiltered(v, filterFuncs) {
			continue
		}

		attachments := []VolumeAttachment{}

		if len(v.Attachments) > 0 {
			for _, a := range v.Attachments {
				attachments = append(attachments, VolumeAttachment{
					Device:     *a.Device,
					InstanceID: *a.InstanceId,
				})
			}
		}

		volume := Volume{
			VolumeID:    *v.VolumeId,
			Attachments: attachments,
		}

		volumes = append(volumes, volume)
	}

	return volumes, nil
}
<|endoftext|>


<commit_before><commit_msg>Increased default number of goroutines<commit_after><|endoftext|>


<commit_before>// Copyright 2015 Peter Fern
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file

// Package golifx provides a simple Go interface to the LIFX LAN protocol.
//
// Based on the protocol documentation available at:
// http://lan.developer.lifx.com/
//
// Also included in cmd/lifx is a small CLI utility that allows interacting with
// your LIFX devices on the LAN.
//
// In various parts of this package you may find references to a Device or a
// Light.  The LIFX protocol makes room for future non-light devices by making a
// light a superset of a device, so a Light is a Device, but a Device is not
// necessarily a Light.  At this stage, LIFX only produces lights though, so
// they are the only type of device you will interact with.
package golifx

import (
	"time"

	"github.com/pdf/golifx/common"
)

const (
	// VERSION of this library
	VERSION = `0.3.0`
)

// NewClient returns a pointer to a new Client and any error that occurred
// initializing the client, using the protocol p.  It also kicks off a discovery
// run.
func NewClient(p common.Protocol) (*Client, error) {
	c := &Client{
		protocol:              p,
		devices:               make(map[uint64]common.Device),
		locations:             make(map[string]common.Location),
		groups:                make(map[string]common.Group),
		subscriptions:         make(map[string]*common.Subscription),
		timeout:               common.DefaultTimeout,
		retryInterval:         common.DefaultRetryInterval,
		internalRetryInterval: 10 * time.Millisecond,
		quitChan:              make(chan struct{}, 2),
	}
	p.SetClient(c)
	err := c.discover()
	return c, err
}

// SetLogger allows assigning a custom levelled logger that conforms to the
// common.Logger interface.  To capture logs generated during client creation,
// this should be called before creating a Client. Defaults to
// common.StubLogger, which does no logging at all.
func SetLogger(logger common.Logger) {
	common.SetLogger(logger)
}
<commit_msg>Bump to v0.3.1<commit_after>// Copyright 2015 Peter Fern
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file

// Package golifx provides a simple Go interface to the LIFX LAN protocol.
//
// Based on the protocol documentation available at:
// http://lan.developer.lifx.com/
//
// Also included in cmd/lifx is a small CLI utility that allows interacting with
// your LIFX devices on the LAN.
//
// In various parts of this package you may find references to a Device or a
// Light.  The LIFX protocol makes room for future non-light devices by making a
// light a superset of a device, so a Light is a Device, but a Device is not
// necessarily a Light.  At this stage, LIFX only produces lights though, so
// they are the only type of device you will interact with.
package golifx

import (
	"time"

	"github.com/pdf/golifx/common"
)

const (
	// VERSION of this library
	VERSION = `0.3.1`
)

// NewClient returns a pointer to a new Client and any error that occurred
// initializing the client, using the protocol p.  It also kicks off a discovery
// run.
func NewClient(p common.Protocol) (*Client, error) {
	c := &Client{
		protocol:              p,
		devices:               make(map[uint64]common.Device),
		locations:             make(map[string]common.Location),
		groups:                make(map[string]common.Group),
		subscriptions:         make(map[string]*common.Subscription),
		timeout:               common.DefaultTimeout,
		retryInterval:         common.DefaultRetryInterval,
		internalRetryInterval: 10 * time.Millisecond,
		quitChan:              make(chan struct{}, 2),
	}
	p.SetClient(c)
	err := c.discover()
	return c, err
}

// SetLogger allows assigning a custom levelled logger that conforms to the
// common.Logger interface.  To capture logs generated during client creation,
// this should be called before creating a Client. Defaults to
// common.StubLogger, which does no logging at all.
func SetLogger(logger common.Logger) {
	common.SetLogger(logger)
}
<|endoftext|>


<commit_before>// This package provides a framework for creating powerful lexical
// preprocessors for go code.
package goprep

import (
	"fmt"
	"io"
	"io/ioutil"
	"go/printer"
	"go/parser"
	"go/scanner"
	"go/token"
	"os"
)

// Represents a token as returned by scanner.Scanner.Scan(), with the position,
// token type, and string representation.
type TokenInfo struct {
	Pos   token.Pos
	Token token.Token
	Str   string
}

// StdInit initializes appropriate processing channels for os.Stdin and
// os.Stdout. For the most part, this should be used instead of specific calls
// to Write and Read.
func StdInit() (<-chan TokenInfo, chan<- string, <-chan interface{}) {
	tokIn := Read(os.Stdin)
	tokOut, done := Write(os.Stdout)
	return tokIn, tokOut, done
}

// Write allows writing properly formatted go code to a given io.Writer via a
// series of token strings passed to the returned channel. The second returned
// channel will have a single nil value sent when writing is complete.
func Write(output io.Writer) (chan<- string, <-chan interface{}) {
	tokC := make(chan string)
	done := make(chan interface{})

	reader, writer := io.Pipe()

	// spit the tokens to the write end of the pipe
	go func(output io.WriteCloser, tokC <-chan string) {
		for tok := range tokC {
			fmt.Fprintf(output, " %s", tok)
		}
		output.Close()
	}(writer, tokC)

	// parse the tokens into an AST and write to output
	go func(reader io.ReadCloser, output io.Writer, done chan interface{}) {
		fset := token.NewFileSet()
		file, err := parser.ParseFile(
			fset, "<stdin>", reader, parser.ParseComments)
		if err != nil {
			panic(err)
		}
		printer.Fprint(output, fset, file)
		done <- nil
	}(reader, output, done)

	return tokC, done
}

// Read reads from the given io.Reader and writes a series of TokenInfo objects
// to the returned channel.
func Read(input io.Reader) <-chan TokenInfo {
	// start reading
	src, err := ioutil.ReadAll(input)
	if err != nil { panic(err) }

	fset := token.NewFileSet()
	file := fset.AddFile("<stdin>", fset.Base(), len(src))

	s := scanner.Scanner{}
	s.Init(file, src, nil, scanner.InsertSemis | scanner.ScanComments)

	tokC := make(chan TokenInfo)

	go func(s scanner.Scanner, tokC chan<- TokenInfo) {
		pos, tok, str := s.Scan()
		for tok != token.EOF {
			if tok == token.COMMENT {
				str = str + "\n"
			}
			tokC <- TokenInfo{pos, tok, str}
			pos, tok, str = s.Scan()
		}
		close(tokC)
	}(s, tokC)

	return tokC
}

// True takes a TokenInfo and always returns true.
func True(TokenInfo) bool { return true }

// False takes a TokenInfo and always returns false.
func False(TokenInfo) bool { return false}

// Ignore produces a modified input stream that does not include any tokens for
// which f evaluates to true, thus discarding a certain class of tokens.
func Ignore(tIn <-chan TokenInfo, out chan<- string, f func(TokenInfo) bool) <-chan TokenInfo {
	tOut := make(chan TokenInfo)
	go func() {
		for tok := range tIn {
			if !f(tok) {
				tOut <- tok
			}
		}
		close(tOut)
	}()
	return tOut
}

// IgnoreToken is like Ignore, discarding all tokens whose string content is
// equal to the given string.
func IgnoreToken(tIn <-chan TokenInfo, out chan<- string, str string) <-chan TokenInfo {
	return Ignore(tIn, out, func(ti TokenInfo) bool {
		return ti.Str == str
	})
}

// IgnoreType is like Ignore, discarding all tokens of a certain type.
func IgnoreType(tIn <-chan TokenInfo, out chan<- string, tok token.Token) <-chan TokenInfo {
	return Ignore(tIn, out, func(ti TokenInfo) bool {
		return ti.Token == tok
	})
}

// Pass redirects all tokens for which f evaluates to true to the output
// channel, returning the altered input channel.
func Pass(tIn <-chan TokenInfo, out chan<- string, f func(TokenInfo) bool) <-chan TokenInfo {
	tOut := make(chan TokenInfo)
	go func() {
		for tok := range tIn {
			if f(tok) {
				out <- tok.Str
			} else {
				tOut <- tok
			}
		}
		close(tOut)
	}()
	return tOut
}

// PassToken is like Pass, passing all tokens whose string content is equal to
// the given string.
func PassToken(tIn <-chan TokenInfo, out chan<- string, str string) <-chan TokenInfo {
	return Pass(tIn, out, func(ti TokenInfo) bool {
		return ti.Str == str
	})
}

// PassType is like Pass, passing all tokens of a certain type.
func PassType(tIn <-chan TokenInfo, out chan<- string, tok token.Token) <-chan TokenInfo {
	return Pass(tIn, out, func(ti TokenInfo) bool {
		return ti.Token == tok
	})
}

// Discard discards whatever tokens remain and waits for the channel to close,
// then closing the output channel as well.  It is an appropriate way to end a
// long list of processing declarations.
func Discard(tIn <-chan TokenInfo, out chan<- string) {
	for _ = range tIn {}
	close(out)
}

<commit_msg>Adding TODO notice<commit_after>// This package provides a framework for creating powerful lexical
// preprocessors for go code.
//
// TODO add 'line' directives
package goprep

import (
	"fmt"
	"io"
	"io/ioutil"
	"go/printer"
	"go/parser"
	"go/scanner"
	"go/token"
	"os"
)

// Represents a token as returned by scanner.Scanner.Scan(), with the position,
// token type, and string representation.
type TokenInfo struct {
	Pos   token.Pos
	Token token.Token
	Str   string
}

// StdInit initializes appropriate processing channels for os.Stdin and
// os.Stdout. For the most part, this should be used instead of specific calls
// to Write and Read.
func StdInit() (<-chan TokenInfo, chan<- string, <-chan interface{}) {
	tokIn := Read(os.Stdin)
	tokOut, done := Write(os.Stdout)
	return tokIn, tokOut, done
}

// Write allows writing properly formatted go code to a given io.Writer via a
// series of token strings passed to the returned channel. The second returned
// channel will have a single nil value sent when writing is complete.
func Write(output io.Writer) (chan<- string, <-chan interface{}) {
	tokC := make(chan string)
	done := make(chan interface{})

	reader, writer := io.Pipe()

	// spit the tokens to the write end of the pipe
	go func(output io.WriteCloser, tokC <-chan string) {
		for tok := range tokC {
			fmt.Fprintf(output, " %s", tok)
		}
		output.Close()
	}(writer, tokC)

	// parse the tokens into an AST and write to output
	go func(reader io.ReadCloser, output io.Writer, done chan interface{}) {
		fset := token.NewFileSet()
		file, err := parser.ParseFile(
			fset, "<stdin>", reader, parser.ParseComments)
		if err != nil {
			panic(err)
		}
		printer.Fprint(output, fset, file)
		done <- nil
	}(reader, output, done)

	return tokC, done
}

// Read reads from the given io.Reader and writes a series of TokenInfo objects
// to the returned channel.
func Read(input io.Reader) <-chan TokenInfo {
	// start reading
	src, err := ioutil.ReadAll(input)
	if err != nil { panic(err) }

	fset := token.NewFileSet()
	file := fset.AddFile("<stdin>", fset.Base(), len(src))

	s := scanner.Scanner{}
	s.Init(file, src, nil, scanner.InsertSemis | scanner.ScanComments)

	tokC := make(chan TokenInfo)

	go func(s scanner.Scanner, tokC chan<- TokenInfo) {
		pos, tok, str := s.Scan()
		for tok != token.EOF {
			if tok == token.COMMENT {
				str = str + "\n"
			}
			tokC <- TokenInfo{pos, tok, str}
			pos, tok, str = s.Scan()
		}
		close(tokC)
	}(s, tokC)

	return tokC
}

// True takes a TokenInfo and always returns true.
func True(TokenInfo) bool { return true }

// False takes a TokenInfo and always returns false.
func False(TokenInfo) bool { return false}

// Ignore produces a modified input stream that does not include any tokens for
// which f evaluates to true, thus discarding a certain class of tokens.
func Ignore(tIn <-chan TokenInfo, out chan<- string, f func(TokenInfo) bool) <-chan TokenInfo {
	tOut := make(chan TokenInfo)
	go func() {
		for tok := range tIn {
			if !f(tok) {
				tOut <- tok
			}
		}
		close(tOut)
	}()
	return tOut
}

// IgnoreToken is like Ignore, discarding all tokens whose string content is
// equal to the given string.
func IgnoreToken(tIn <-chan TokenInfo, out chan<- string, str string) <-chan TokenInfo {
	return Ignore(tIn, out, func(ti TokenInfo) bool {
		return ti.Str == str
	})
}

// IgnoreType is like Ignore, discarding all tokens of a certain type.
func IgnoreType(tIn <-chan TokenInfo, out chan<- string, tok token.Token) <-chan TokenInfo {
	return Ignore(tIn, out, func(ti TokenInfo) bool {
		return ti.Token == tok
	})
}

// Pass redirects all tokens for which f evaluates to true to the output
// channel, returning the altered input channel.
func Pass(tIn <-chan TokenInfo, out chan<- string, f func(TokenInfo) bool) <-chan TokenInfo {
	tOut := make(chan TokenInfo)
	go func() {
		for tok := range tIn {
			if f(tok) {
				out <- tok.Str
			} else {
				tOut <- tok
			}
		}
		close(tOut)
	}()
	return tOut
}

// PassToken is like Pass, passing all tokens whose string content is equal to
// the given string.
func PassToken(tIn <-chan TokenInfo, out chan<- string, str string) <-chan TokenInfo {
	return Pass(tIn, out, func(ti TokenInfo) bool {
		return ti.Str == str
	})
}

// PassType is like Pass, passing all tokens of a certain type.
func PassType(tIn <-chan TokenInfo, out chan<- string, tok token.Token) <-chan TokenInfo {
	return Pass(tIn, out, func(ti TokenInfo) bool {
		return ti.Token == tok
	})
}

// Discard discards whatever tokens remain and waits for the channel to close,
// then closing the output channel as well.  It is an appropriate way to end a
// long list of processing declarations.
func Discard(tIn <-chan TokenInfo, out chan<- string) {
	for _ = range tIn {}
	close(out)
}

<|endoftext|>


<commit_before>package server_manager

import (
	"log"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/consul/consul/server_details"
	"github.com/hashicorp/consul/lib"
)

type consulServerEventTypes int

const (
	// clientRPCJitterFraction determines the amount of jitter added to
	// clientRPCMinReuseDuration before a connection is expired and a new
	// connection is established in order to rebalance load across consul
	// servers.  The cluster-wide number of connections per second from
	// rebalancing is applied after this jitter to ensure the CPU impact
	// is always finite.  See newRebalanceConnsPerSecPerServer's comment
	// for additional commentary.
	//
	// For example, in a 10K consul cluster with 5x servers, this default
	// averages out to ~13 new connections from rebalancing per server
	// per second (each connection is reused for 120s to 180s).
	clientRPCJitterFraction = 2

	// clientRPCMinReuseDuration controls the minimum amount of time RPC
	// queries are sent over an established connection to a single server
	clientRPCMinReuseDuration = 120 * time.Second

	// initialRebalanceTimeout is the initial value for the
	// rebalanceTimer.  This value is discarded immediately after the
	// client becomes aware of the first server.
	initialRebalanceTimeout = 24 * time.Hour

	// Limit the number of new connections a server receives per second
	// for connection rebalancing.  This limit caps the load caused by
	// continual rebalancing efforts when a cluster is in equilibrium.  A
	// lower value comes at the cost of increased recovery time after a
	// partition.  This parameter begins to take effect when there are
	// more than ~48K clients querying 5x servers or at lower server
	// values when there is a partition.
	//
	// For example, in a 100K consul cluster with 5x servers, it will
	// take ~5min for all servers to rebalance their connections.  If
	// 99,995 agents are in the minority talking to only one server, it
	// will take ~26min for all servers to rebalance.  A 10K cluster in
	// the same scenario will take ~2.6min to rebalance.
	newRebalanceConnsPerSecPerServer = 64
)

type ConsulClusterInfo interface {
	NumNodes() int
}

// serverCfg is the thread-safe configuration struct used to maintain the
// list of Consul servers in ServerManager.
//
// NOTE(sean@): We are explicitly relying on the fact that serverConfig will
// be copied onto the stack.  Please keep this structure light.
type serverConfig struct {
	// servers tracks the locally known servers.  List membership is
	// maintained by Serf.
	servers []*server_details.ServerDetails
}

type ServerManager struct {
	// serverConfig provides the necessary load/store semantics for the
	// server list.
	serverConfigValue atomic.Value
	serverConfigLock  sync.Mutex

	// shutdownCh is a copy of the channel in consul.Client
	shutdownCh chan struct{}

	logger *log.Logger

	// clusterInfo is used to estimate the approximate number of nodes in
	// a cluster and limit the rate at which it rebalances server
	// connections.  ConsulClusterInfo is an interface that wraps serf.
	clusterInfo ConsulClusterInfo

	// notifyFailedServersBarrier is acts as a barrier to prevent
	// queueing behind serverConfigLog and acts as a TryLock().
	notifyFailedBarrier int32
}

// AddServer takes out an internal write lock and adds a new server.  If the
// server is not known, appends the server to the list.  The new server will
// begin seeing use after the rebalance timer fires or enough servers fail
// organically.  If the server is already known, merge the new server
// details.
func (sm *ServerManager) AddServer(server *server_details.ServerDetails) {
	sm.serverConfigLock.Lock()
	defer sm.serverConfigLock.Unlock()
	serverCfg := sm.getServerConfig()

	// Check if this server is known
	found := false
	for idx, existing := range serverCfg.servers {
		if existing.Name == server.Name {
			newServers := make([]*server_details.ServerDetails, len(serverCfg.servers))
			copy(newServers, serverCfg.servers)

			// Overwrite the existing server details in order to
			// possibly update metadata (e.g. server version)
			newServers[idx] = server

			serverCfg.servers = newServers
			found = true
			break
		}
	}

	// Add to the list if not known
	if !found {
		newServers := make([]*server_details.ServerDetails, len(serverCfg.servers), len(serverCfg.servers)+1)
		copy(newServers, serverCfg.servers)
		newServers = append(newServers, server)
		serverCfg.servers = newServers
	}

	sm.saveServerConfig(serverCfg)
}

// cycleServers returns a new list of servers that has dequeued the first
// server and enqueued it at the end of the list.  cycleServers assumes the
// caller is holding the serverConfigLock.
func (sc *serverConfig) cycleServer() (servers []*server_details.ServerDetails) {
	numServers := len(sc.servers)
	if numServers < 2 {
		return servers // No action required
	}

	newServers := make([]*server_details.ServerDetails, 0, numServers)
	newServers = append(newServers, sc.servers[1:]...)
	newServers = append(newServers, sc.servers[0])
	return newServers
}

// FindServer takes out an internal "read lock" and searches through the list
// of servers to find a "healthy" server.  If the server is actually
// unhealthy, we rely on Serf to detect this and remove the node from the
// server list.  If the server at the front of the list has failed or fails
// during an RPC call, it is rotated to the end of the list.  If there are no
// servers available, return nil.
func (sm *ServerManager) FindServer() *server_details.ServerDetails {
	serverCfg := sm.getServerConfig()
	numServers := len(serverCfg.servers)
	if numServers == 0 {
		sm.logger.Printf("[WARN] consul: No servers found in the server config")
		return nil
	} else {
		// Return whatever is at the front of the list because it is
		// assumed to be the oldest in the server list (unless -
		// hypothetically - the server list was rotated right after a
		// server was added).
		return serverCfg.servers[0]
	}
}

// getServerConfig is a convenience method which hides the locking semantics
// of atomic.Value from the caller.
func (sm *ServerManager) getServerConfig() serverConfig {
	return sm.serverConfigValue.Load().(serverConfig)
}

// saveServerConfig is a convenience method which hides the locking semantics
// of atomic.Value from the caller.
func (sm *ServerManager) saveServerConfig(sc serverConfig) {
	sm.serverConfigValue.Store(sc)
}

// New is the only way to safely create a new ServerManager struct.
func New(logger *log.Logger, shutdownCh chan struct{}, clusterInfo ConsulClusterInfo) (sm *ServerManager) {
	// NOTE(sean@): Can't pass *consul.Client due to an import cycle
	sm = new(ServerManager)
	sm.logger = logger
	sm.clusterInfo = clusterInfo
	sm.shutdownCh = shutdownCh

	sc := serverConfig{}
	sc.servers = make([]*server_details.ServerDetails, 0)
	sm.saveServerConfig(sc)
	return sm
}

// NotifyFailedServer marks the passed in server as "failed" by rotating it
// to the end of the server list.
func (sm *ServerManager) NotifyFailedServer(server *server_details.ServerDetails) {
	serverCfg := sm.getServerConfig()

	// If the server being failed is not the first server on the list,
	// this is a noop.  If, however, the server is failed and first on
	// the list, acquire the lock, retest, and take the penalty of moving
	// the server to the end of the list.

	// Use atomic.CAS to emulate a TryLock().
	if len(serverCfg.servers) > 0 && serverCfg.servers[0] == server &&
		atomic.CompareAndSwapInt32(&sm.notifyFailedBarrier, 0, 1) {
		defer atomic.StoreInt32(&sm.notifyFailedBarrier, 0)

		// Grab a lock, retest, and take the hit of cycling the first
		// server to the end.
		sm.serverConfigLock.Lock()
		defer sm.serverConfigLock.Unlock()
		serverCfg = sm.getServerConfig()

		if len(serverCfg.servers) > 0 && serverCfg.servers[0] == server {
			serverCfg.servers = serverCfg.cycleServer()
			sm.saveServerConfig(serverCfg)
		}
	}
}

// NumServers takes out an internal "read lock" and returns the number of
// servers.  numServers includes both healthy and unhealthy servers.
func (sm *ServerManager) NumServers() (numServers int) {
	serverCfg := sm.getServerConfig()
	numServers = len(serverCfg.servers)
	return numServers
}

// RebalanceServers takes out an internal write lock and shuffles the list of
// servers on this agent.  This allows for a redistribution of work across
// consul servers and provides a guarantee that the order of the server list
// isn't related to the age at which the node was added to the cluster.
// Elsewhere we rely on the position in the server list as a hint regarding
// the stability of a server relative to its position in the server list.
// Servers at or near the front of the list are more stable than servers near
// the end of the list.
func (sm *ServerManager) RebalanceServers() {
	sm.serverConfigLock.Lock()
	defer sm.serverConfigLock.Unlock()
	serverCfg := sm.getServerConfig()

	newServers := make([]*server_details.ServerDetails, len(serverCfg.servers))
	copy(newServers, serverCfg.servers)

	// Shuffle the server list on server join.  Servers are selected from
	// the head of the list and are moved to the end of the list on
	// failure.
	for i := len(serverCfg.servers) - 1; i > 0; i-- {
		j := rand.Int31n(int32(i + 1))
		newServers[i], newServers[j] = newServers[j], newServers[i]
	}
	serverCfg.servers = newServers

	sm.saveServerConfig(serverCfg)
}

// RemoveServer takes out an internal write lock and removes a server from
// the server list.
func (sm *ServerManager) RemoveServer(server *server_details.ServerDetails) {
	sm.serverConfigLock.Lock()
	defer sm.serverConfigLock.Unlock()
	serverCfg := sm.getServerConfig()

	// Remove the server if known
	n := len(serverCfg.servers)
	for i := 0; i < n; i++ {
		if serverCfg.servers[i].Name == server.Name {
			newServers := make([]*server_details.ServerDetails, 0, len(serverCfg.servers)-1)
			newServers = append(newServers, serverCfg.servers[:i]...)
			newServers = append(newServers, serverCfg.servers[i+1:]...)
			serverCfg.servers = newServers

			sm.saveServerConfig(serverCfg)
			return
		}
	}
}

// refreshServerRebalanceTimer is only called once the rebalanceTimer
// expires.  Historically this was an expensive routine and is intended to be
// run in isolation in a dedicated, non-concurrent task.
func (sm *ServerManager) refreshServerRebalanceTimer(timer *time.Timer) time.Duration {
	serverCfg := sm.getServerConfig()
	numConsulServers := len(serverCfg.servers)
	// Limit this connection's life based on the size (and health) of the
	// cluster.  Never rebalance a connection more frequently than
	// connReuseLowWatermarkDuration, and make sure we never exceed
	// clusterWideRebalanceConnsPerSec operations/s across numLANMembers.
	clusterWideRebalanceConnsPerSec := float64(numConsulServers * newRebalanceConnsPerSecPerServer)
	connReuseLowWatermarkDuration := clientRPCMinReuseDuration + lib.RandomStagger(clientRPCMinReuseDuration/clientRPCJitterFraction)
	numLANMembers := sm.clusterInfo.NumNodes()
	connRebalanceTimeout := lib.RateScaledInterval(clusterWideRebalanceConnsPerSec, connReuseLowWatermarkDuration, numLANMembers)

	timer.Reset(connRebalanceTimeout)
	return connRebalanceTimeout
}

// Start is used to start and manage the task of automatically shuffling and
// rebalancing the list of consul servers.  This maintenance only happens
// periodically based on the expiration of the timer.  Failed servers are
// automatically cycled to the end of the list.  New servers are appended to
// the list.  The order of the server list must be shuffled periodically to
// distribute load across all known and available consul servers.
func (sm *ServerManager) Start() {
	var rebalanceTimer *time.Timer = time.NewTimer(initialRebalanceTimeout)
	var rebalanceTaskDispatched int32

	func() {
		sm.serverConfigLock.Lock()
		defer sm.serverConfigLock.Unlock()

		serverCfgPtr := sm.serverConfigValue.Load()
		if serverCfgPtr == nil {
			panic("server config has not been initialized")
		}
		var serverCfg serverConfig
		serverCfg = serverCfgPtr.(serverConfig)
		sm.saveServerConfig(serverCfg)
	}()

	for {
		select {
		case <-rebalanceTimer.C:
			sm.logger.Printf("[INFO] server manager: server rebalance timeout")
			sm.RebalanceServers()

			// Only run one rebalance task at a time, but do
			// allow for the channel to be drained
			if atomic.CompareAndSwapInt32(&rebalanceTaskDispatched, 0, 1) {
				sm.logger.Printf("[INFO] server manager: Launching rebalance duration task")
				go func() {
					defer atomic.StoreInt32(&rebalanceTaskDispatched, 0)
					sm.refreshServerRebalanceTimer(rebalanceTimer)
				}()
			}

		case <-sm.shutdownCh:
			sm.logger.Printf("[INFO] server manager: shutting down")
			return
		}
	}
}
<commit_msg>Only rotate server list with more than one server<commit_after>package server_manager

import (
	"log"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/consul/consul/server_details"
	"github.com/hashicorp/consul/lib"
)

type consulServerEventTypes int

const (
	// clientRPCJitterFraction determines the amount of jitter added to
	// clientRPCMinReuseDuration before a connection is expired and a new
	// connection is established in order to rebalance load across consul
	// servers.  The cluster-wide number of connections per second from
	// rebalancing is applied after this jitter to ensure the CPU impact
	// is always finite.  See newRebalanceConnsPerSecPerServer's comment
	// for additional commentary.
	//
	// For example, in a 10K consul cluster with 5x servers, this default
	// averages out to ~13 new connections from rebalancing per server
	// per second (each connection is reused for 120s to 180s).
	clientRPCJitterFraction = 2

	// clientRPCMinReuseDuration controls the minimum amount of time RPC
	// queries are sent over an established connection to a single server
	clientRPCMinReuseDuration = 120 * time.Second

	// initialRebalanceTimeout is the initial value for the
	// rebalanceTimer.  This value is discarded immediately after the
	// client becomes aware of the first server.
	initialRebalanceTimeout = 24 * time.Hour

	// Limit the number of new connections a server receives per second
	// for connection rebalancing.  This limit caps the load caused by
	// continual rebalancing efforts when a cluster is in equilibrium.  A
	// lower value comes at the cost of increased recovery time after a
	// partition.  This parameter begins to take effect when there are
	// more than ~48K clients querying 5x servers or at lower server
	// values when there is a partition.
	//
	// For example, in a 100K consul cluster with 5x servers, it will
	// take ~5min for all servers to rebalance their connections.  If
	// 99,995 agents are in the minority talking to only one server, it
	// will take ~26min for all servers to rebalance.  A 10K cluster in
	// the same scenario will take ~2.6min to rebalance.
	newRebalanceConnsPerSecPerServer = 64
)

type ConsulClusterInfo interface {
	NumNodes() int
}

// serverCfg is the thread-safe configuration struct used to maintain the
// list of Consul servers in ServerManager.
//
// NOTE(sean@): We are explicitly relying on the fact that serverConfig will
// be copied onto the stack.  Please keep this structure light.
type serverConfig struct {
	// servers tracks the locally known servers.  List membership is
	// maintained by Serf.
	servers []*server_details.ServerDetails
}

type ServerManager struct {
	// serverConfig provides the necessary load/store semantics for the
	// server list.
	serverConfigValue atomic.Value
	serverConfigLock  sync.Mutex

	// shutdownCh is a copy of the channel in consul.Client
	shutdownCh chan struct{}

	logger *log.Logger

	// clusterInfo is used to estimate the approximate number of nodes in
	// a cluster and limit the rate at which it rebalances server
	// connections.  ConsulClusterInfo is an interface that wraps serf.
	clusterInfo ConsulClusterInfo

	// notifyFailedServersBarrier is acts as a barrier to prevent
	// queueing behind serverConfigLog and acts as a TryLock().
	notifyFailedBarrier int32
}

// AddServer takes out an internal write lock and adds a new server.  If the
// server is not known, appends the server to the list.  The new server will
// begin seeing use after the rebalance timer fires or enough servers fail
// organically.  If the server is already known, merge the new server
// details.
func (sm *ServerManager) AddServer(server *server_details.ServerDetails) {
	sm.serverConfigLock.Lock()
	defer sm.serverConfigLock.Unlock()
	serverCfg := sm.getServerConfig()

	// Check if this server is known
	found := false
	for idx, existing := range serverCfg.servers {
		if existing.Name == server.Name {
			newServers := make([]*server_details.ServerDetails, len(serverCfg.servers))
			copy(newServers, serverCfg.servers)

			// Overwrite the existing server details in order to
			// possibly update metadata (e.g. server version)
			newServers[idx] = server

			serverCfg.servers = newServers
			found = true
			break
		}
	}

	// Add to the list if not known
	if !found {
		newServers := make([]*server_details.ServerDetails, len(serverCfg.servers), len(serverCfg.servers)+1)
		copy(newServers, serverCfg.servers)
		newServers = append(newServers, server)
		serverCfg.servers = newServers
	}

	sm.saveServerConfig(serverCfg)
}

// cycleServers returns a new list of servers that has dequeued the first
// server and enqueued it at the end of the list.  cycleServers assumes the
// caller is holding the serverConfigLock.
func (sc *serverConfig) cycleServer() (servers []*server_details.ServerDetails) {
	numServers := len(sc.servers)
	if numServers < 2 {
		return servers // No action required
	}

	newServers := make([]*server_details.ServerDetails, 0, numServers)
	newServers = append(newServers, sc.servers[1:]...)
	newServers = append(newServers, sc.servers[0])
	return newServers
}

// FindServer takes out an internal "read lock" and searches through the list
// of servers to find a "healthy" server.  If the server is actually
// unhealthy, we rely on Serf to detect this and remove the node from the
// server list.  If the server at the front of the list has failed or fails
// during an RPC call, it is rotated to the end of the list.  If there are no
// servers available, return nil.
func (sm *ServerManager) FindServer() *server_details.ServerDetails {
	serverCfg := sm.getServerConfig()
	numServers := len(serverCfg.servers)
	if numServers == 0 {
		sm.logger.Printf("[WARN] consul: No servers found in the server config")
		return nil
	} else {
		// Return whatever is at the front of the list because it is
		// assumed to be the oldest in the server list (unless -
		// hypothetically - the server list was rotated right after a
		// server was added).
		return serverCfg.servers[0]
	}
}

// getServerConfig is a convenience method which hides the locking semantics
// of atomic.Value from the caller.
func (sm *ServerManager) getServerConfig() serverConfig {
	return sm.serverConfigValue.Load().(serverConfig)
}

// saveServerConfig is a convenience method which hides the locking semantics
// of atomic.Value from the caller.
func (sm *ServerManager) saveServerConfig(sc serverConfig) {
	sm.serverConfigValue.Store(sc)
}

// New is the only way to safely create a new ServerManager struct.
func New(logger *log.Logger, shutdownCh chan struct{}, clusterInfo ConsulClusterInfo) (sm *ServerManager) {
	// NOTE(sean@): Can't pass *consul.Client due to an import cycle
	sm = new(ServerManager)
	sm.logger = logger
	sm.clusterInfo = clusterInfo
	sm.shutdownCh = shutdownCh

	sc := serverConfig{}
	sc.servers = make([]*server_details.ServerDetails, 0)
	sm.saveServerConfig(sc)
	return sm
}

// NotifyFailedServer marks the passed in server as "failed" by rotating it
// to the end of the server list.
func (sm *ServerManager) NotifyFailedServer(server *server_details.ServerDetails) {
	serverCfg := sm.getServerConfig()

	// If the server being failed is not the first server on the list,
	// this is a noop.  If, however, the server is failed and first on
	// the list, acquire the lock, retest, and take the penalty of moving
	// the server to the end of the list.

	// Only rotate the server list when there is more than one server
	if len(serverCfg.servers) > 1 && serverCfg.servers[0] == server &&
		// Use atomic.CAS to emulate a TryLock().
		atomic.CompareAndSwapInt32(&sm.notifyFailedBarrier, 0, 1) {
		defer atomic.StoreInt32(&sm.notifyFailedBarrier, 0)

		// Grab a lock, retest, and take the hit of cycling the first
		// server to the end.
		sm.serverConfigLock.Lock()
		defer sm.serverConfigLock.Unlock()
		serverCfg = sm.getServerConfig()

		if len(serverCfg.servers) > 1 && serverCfg.servers[0] == server {
			serverCfg.servers = serverCfg.cycleServer()
			sm.saveServerConfig(serverCfg)
		}
	}
}

// NumServers takes out an internal "read lock" and returns the number of
// servers.  numServers includes both healthy and unhealthy servers.
func (sm *ServerManager) NumServers() (numServers int) {
	serverCfg := sm.getServerConfig()
	numServers = len(serverCfg.servers)
	return numServers
}

// RebalanceServers takes out an internal write lock and shuffles the list of
// servers on this agent.  This allows for a redistribution of work across
// consul servers and provides a guarantee that the order of the server list
// isn't related to the age at which the node was added to the cluster.
// Elsewhere we rely on the position in the server list as a hint regarding
// the stability of a server relative to its position in the server list.
// Servers at or near the front of the list are more stable than servers near
// the end of the list.
func (sm *ServerManager) RebalanceServers() {
	sm.serverConfigLock.Lock()
	defer sm.serverConfigLock.Unlock()
	serverCfg := sm.getServerConfig()

	newServers := make([]*server_details.ServerDetails, len(serverCfg.servers))
	copy(newServers, serverCfg.servers)

	// Shuffle the server list on server join.  Servers are selected from
	// the head of the list and are moved to the end of the list on
	// failure.
	for i := len(serverCfg.servers) - 1; i > 0; i-- {
		j := rand.Int31n(int32(i + 1))
		newServers[i], newServers[j] = newServers[j], newServers[i]
	}
	serverCfg.servers = newServers

	sm.saveServerConfig(serverCfg)
}

// RemoveServer takes out an internal write lock and removes a server from
// the server list.
func (sm *ServerManager) RemoveServer(server *server_details.ServerDetails) {
	sm.serverConfigLock.Lock()
	defer sm.serverConfigLock.Unlock()
	serverCfg := sm.getServerConfig()

	// Remove the server if known
	n := len(serverCfg.servers)
	for i := 0; i < n; i++ {
		if serverCfg.servers[i].Name == server.Name {
			newServers := make([]*server_details.ServerDetails, 0, len(serverCfg.servers)-1)
			newServers = append(newServers, serverCfg.servers[:i]...)
			newServers = append(newServers, serverCfg.servers[i+1:]...)
			serverCfg.servers = newServers

			sm.saveServerConfig(serverCfg)
			return
		}
	}
}

// refreshServerRebalanceTimer is only called once the rebalanceTimer
// expires.  Historically this was an expensive routine and is intended to be
// run in isolation in a dedicated, non-concurrent task.
func (sm *ServerManager) refreshServerRebalanceTimer(timer *time.Timer) time.Duration {
	serverCfg := sm.getServerConfig()
	numConsulServers := len(serverCfg.servers)
	// Limit this connection's life based on the size (and health) of the
	// cluster.  Never rebalance a connection more frequently than
	// connReuseLowWatermarkDuration, and make sure we never exceed
	// clusterWideRebalanceConnsPerSec operations/s across numLANMembers.
	clusterWideRebalanceConnsPerSec := float64(numConsulServers * newRebalanceConnsPerSecPerServer)
	connReuseLowWatermarkDuration := clientRPCMinReuseDuration + lib.RandomStagger(clientRPCMinReuseDuration/clientRPCJitterFraction)
	numLANMembers := sm.clusterInfo.NumNodes()
	connRebalanceTimeout := lib.RateScaledInterval(clusterWideRebalanceConnsPerSec, connReuseLowWatermarkDuration, numLANMembers)

	timer.Reset(connRebalanceTimeout)
	return connRebalanceTimeout
}

// Start is used to start and manage the task of automatically shuffling and
// rebalancing the list of consul servers.  This maintenance only happens
// periodically based on the expiration of the timer.  Failed servers are
// automatically cycled to the end of the list.  New servers are appended to
// the list.  The order of the server list must be shuffled periodically to
// distribute load across all known and available consul servers.
func (sm *ServerManager) Start() {
	var rebalanceTimer *time.Timer = time.NewTimer(initialRebalanceTimeout)
	var rebalanceTaskDispatched int32

	func() {
		sm.serverConfigLock.Lock()
		defer sm.serverConfigLock.Unlock()

		serverCfgPtr := sm.serverConfigValue.Load()
		if serverCfgPtr == nil {
			panic("server config has not been initialized")
		}
		var serverCfg serverConfig
		serverCfg = serverCfgPtr.(serverConfig)
		sm.saveServerConfig(serverCfg)
	}()

	for {
		select {
		case <-rebalanceTimer.C:
			sm.logger.Printf("[INFO] server manager: server rebalance timeout")
			sm.RebalanceServers()

			// Only run one rebalance task at a time, but do
			// allow for the channel to be drained
			if atomic.CompareAndSwapInt32(&rebalanceTaskDispatched, 0, 1) {
				sm.logger.Printf("[INFO] server manager: Launching rebalance duration task")
				go func() {
					defer atomic.StoreInt32(&rebalanceTaskDispatched, 0)
					sm.refreshServerRebalanceTimer(rebalanceTimer)
				}()
			}

		case <-sm.shutdownCh:
			sm.logger.Printf("[INFO] server manager: shutting down")
			return
		}
	}
}
<|endoftext|>


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

// Package repository implements a deployment repository using a map.
// It can be easily replaced by a deployment repository that uses some
// form of persistent storage.
package repository

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/kubernetes/deployment-manager/manager/manager"
)

// deploymentTypeInstanceMap stores type instances mapped by deployment name.
// This allows for simple updating and deleting of per-deployment instances
// when deployments are created/updated/deleted.
type deploymentTypeInstanceMap map[string][]*manager.TypeInstance
type typeInstanceMap map[string]deploymentTypeInstanceMap

type mapBasedRepository struct {
	sync.RWMutex
	deployments map[string]manager.Deployment
	manifests   map[string]map[string]*manager.Manifest
	instances   typeInstanceMap
}

// NewMapBasedRepository returns a new map based repository.
func NewMapBasedRepository() manager.Repository {
	return &mapBasedRepository{
		deployments: make(map[string]manager.Deployment, 0),
		manifests:   make(map[string]map[string]*manager.Manifest, 0),
		instances:   typeInstanceMap{},
	}
}

// ListDeployments returns of all of the deployments in the repository.
func (r *mapBasedRepository) ListDeployments() ([]manager.Deployment, error) {
	r.RLock()
	defer r.RUnlock()

	l := []manager.Deployment{}
	for _, deployment := range r.deployments {
		l = append(l, deployment)
	}

	return l, nil
}

// GetDeployment returns the deployment with the supplied name.
// If the deployment is not found, it returns an error.
func (r *mapBasedRepository) GetDeployment(name string) (*manager.Deployment, error) {
	d, ok := r.deployments[name]
	if !ok {
		return nil, fmt.Errorf("deployment %s not found", name)
	}
	return &d, nil
}

// GetValidDeployment returns the deployment with the supplied name.
// If the deployment is not found or marked as deleted, it returns an error.
func (r *mapBasedRepository) GetValidDeployment(name string) (*manager.Deployment, error) {
	d, err := r.GetDeployment(name)
	if err != nil {
		return nil, err
	}

	if d.Status == manager.DeletedStatus {
		return nil, fmt.Errorf("deployment %s is deleted", name)
	}

	return d, nil
}

// SetDeploymentStatus sets the DeploymentStatus of the deployment and updates ModifiedAt
func (r *mapBasedRepository) SetDeploymentStatus(name string, status manager.DeploymentStatus) error {
	return func() error {
		r.Lock()
		defer r.Unlock()

		d, err := r.GetValidDeployment(name)
		if err != nil {
			return err
		}

		d.Status = status
		d.ModifiedAt = time.Now()
		r.deployments[name] = *d
		return nil
	}()
}

// CreateDeployment creates a new deployment and stores it in the repository.
func (r *mapBasedRepository) CreateDeployment(name string) (*manager.Deployment, error) {
	d, err := func() (*manager.Deployment, error) {
		r.Lock()
		defer r.Unlock()

		exists, _ := r.GetValidDeployment(name)
		if exists != nil {
			return nil, fmt.Errorf("Deployment %s already exists", name)
		}

		d := manager.NewDeployment(name)
		d.Status = manager.CreatedStatus
		d.DeployedAt = time.Now()
		r.deployments[name] = *d
		return d, nil
	}()

	if err != nil {
		return nil, err
	}

	log.Printf("created deployment: %v", d)
	return d, nil
}

func (r *mapBasedRepository) AddManifest(deploymentName string, manifest *manager.Manifest) error {
	err := func() error {
		r.Lock()
		defer r.Unlock()

		d, err := r.GetValidDeployment(deploymentName)
		if err != nil {
			return err
		}

		l, err := r.listManifestsForDeployment(deploymentName)
		if err != nil {
			return err
		}

		// Make sure the manifest doesn't already exist, and if not, add the manifest to
		// map of manifests this deployment has
		if _, ok := l[manifest.Name]; ok {
			return fmt.Errorf("Manifest %s already exists in deployment %s", manifest.Name, deploymentName)
		}

		l[manifest.Name] = manifest
		d.LatestManifest = manifest.Name
		return nil
	}()

	if err != nil {
		return err
	}

	log.Printf("Added manifest %s to deployment: %s", manifest.Name, deploymentName)
	return nil
}

// DeleteDeployment deletes the deployment with the supplied name.
// If forget is true, then the deployment is removed from the repository.
// Otherwise, it is marked as deleted and retained.
func (r *mapBasedRepository) DeleteDeployment(name string, forget bool) (*manager.Deployment, error) {
	d, err := func() (*manager.Deployment, error) {
		r.Lock()
		defer r.Unlock()

		d, err := r.GetValidDeployment(name)
		if err != nil {
			return nil, err
		}

		if !forget {
			d.DeletedAt = time.Now()
			d.Status = manager.DeletedStatus
			r.deployments[name] = *d
		} else {
			delete(r.deployments, name)
			delete(r.manifests, name)
			d.LatestManifest = ""
		}

		return d, nil
	}()

	if err != nil {
		return nil, err
	}

	log.Printf("deleted deployment: %v", d)
	return d, nil
}

func (r *mapBasedRepository) ListManifests(deploymentName string) (map[string]*manager.Manifest, error) {
	r.Lock()
	defer r.Unlock()

	_, err := r.GetValidDeployment(deploymentName)
	if err != nil {
		return nil, err
	}

	return r.listManifestsForDeployment(deploymentName)
}

func (r *mapBasedRepository) listManifestsForDeployment(deploymentName string) (map[string]*manager.Manifest, error) {
	l, ok := r.manifests[deploymentName]
	if !ok {
		l = make(map[string]*manager.Manifest, 0)
		r.manifests[deploymentName] = l
	}

	return l, nil
}

func (r *mapBasedRepository) GetManifest(deploymentName string, manifestName string) (*manager.Manifest, error) {
	r.Lock()
	defer r.Unlock()

	_, err := r.GetValidDeployment(deploymentName)
	if err != nil {
		return nil, err
	}

	return r.getManifestForDeployment(deploymentName, manifestName)
}

func (r *mapBasedRepository) getManifestForDeployment(deploymentName string, manifestName string) (*manager.Manifest, error) {
	l, err := r.listManifestsForDeployment(deploymentName)
	if err != nil {
		return nil, err
	}

	m, ok := l[manifestName]
	if !ok {
		return nil, fmt.Errorf("manifest %s not found in deployment %s", manifestName, deploymentName)
	}

	return m, nil
}

// GetLatestManifest returns the latest manifest for a given deployment,
// which by definition is the manifest with the largest time stamp.
func (r *mapBasedRepository) GetLatestManifest(deploymentName string) (*manager.Manifest, error) {
	r.Lock()
	defer r.Unlock()

	d, err := r.GetValidDeployment(deploymentName)
	if err != nil {
		return nil, err
	}

	return r.getManifestForDeployment(deploymentName, d.LatestManifest)
}

// ListTypes returns all types known from existing instances.
func (r *mapBasedRepository) ListTypes() []string {
	var keys []string
	for k := range r.instances {
		keys = append(keys, k)
	}

	return keys
}

// GetTypeInstances returns all instances of a given type. If type is empty,
// returns all instances for all types.
func (r *mapBasedRepository) GetTypeInstances(typeName string) []*manager.TypeInstance {
	r.Lock()
	defer r.Unlock()

	var instances []*manager.TypeInstance
	for t, dInstMap := range r.instances {
		if t == typeName || typeName == "all" {
			for _, i := range dInstMap {
				instances = append(instances, i...)
			}
		}
	}

	return instances
}

// ClearTypeInstances deletes all instances associated with the given
// deployment name from the type instance repository.
func (r *mapBasedRepository) ClearTypeInstances(deploymentName string) {
	r.Lock()
	defer r.Unlock()

	for t, dMap := range r.instances {
		delete(dMap, deploymentName)
		if len(dMap) == 0 {
			delete(r.instances, t)
		}
	}
}

// SetTypeInstances sets all type instances for a given deployment name.
//
// To clear the current set of instances first, caller should first use
// ClearTypeInstances().
func (r *mapBasedRepository) SetTypeInstances(deploymentName string, instances map[string][]*manager.TypeInstance) {
	r.Lock()
	defer r.Unlock()

	// Add each instance list to the appropriate type map.
	for t, is := range instances {
		if r.instances[t] == nil {
			r.instances[t] = make(deploymentTypeInstanceMap)
		}

		r.instances[t][deploymentName] = is
	}
}
<commit_msg>Repository now stores latest manifest in deployment.<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package repository implements a deployment repository using a map.
// It can be easily replaced by a deployment repository that uses some
// form of persistent storage.
package repository

import (
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/kubernetes/deployment-manager/manager/manager"
)

// deploymentTypeInstanceMap stores type instances mapped by deployment name.
// This allows for simple updating and deleting of per-deployment instances
// when deployments are created/updated/deleted.
type deploymentTypeInstanceMap map[string][]*manager.TypeInstance
type typeInstanceMap map[string]deploymentTypeInstanceMap

type mapBasedRepository struct {
	sync.RWMutex
	deployments map[string]manager.Deployment
	manifests   map[string]map[string]*manager.Manifest
	instances   typeInstanceMap
}

// NewMapBasedRepository returns a new map based repository.
func NewMapBasedRepository() manager.Repository {
	return &mapBasedRepository{
		deployments: make(map[string]manager.Deployment, 0),
		manifests:   make(map[string]map[string]*manager.Manifest, 0),
		instances:   typeInstanceMap{},
	}
}

// ListDeployments returns of all of the deployments in the repository.
func (r *mapBasedRepository) ListDeployments() ([]manager.Deployment, error) {
	r.RLock()
	defer r.RUnlock()

	l := []manager.Deployment{}
	for _, deployment := range r.deployments {
		l = append(l, deployment)
	}

	return l, nil
}

// GetDeployment returns the deployment with the supplied name.
// If the deployment is not found, it returns an error.
func (r *mapBasedRepository) GetDeployment(name string) (*manager.Deployment, error) {
	d, ok := r.deployments[name]
	if !ok {
		return nil, fmt.Errorf("deployment %s not found", name)
	}
	return &d, nil
}

// GetValidDeployment returns the deployment with the supplied name.
// If the deployment is not found or marked as deleted, it returns an error.
func (r *mapBasedRepository) GetValidDeployment(name string) (*manager.Deployment, error) {
	d, err := r.GetDeployment(name)
	if err != nil {
		return nil, err
	}

	if d.Status == manager.DeletedStatus {
		return nil, fmt.Errorf("deployment %s is deleted", name)
	}

	return d, nil
}

// SetDeploymentStatus sets the DeploymentStatus of the deployment and updates ModifiedAt
func (r *mapBasedRepository) SetDeploymentStatus(name string, status manager.DeploymentStatus) error {
	return func() error {
		r.Lock()
		defer r.Unlock()

		d, err := r.GetValidDeployment(name)
		if err != nil {
			return err
		}

		d.Status = status
		d.ModifiedAt = time.Now()
		r.deployments[name] = *d
		return nil
	}()
}

// CreateDeployment creates a new deployment and stores it in the repository.
func (r *mapBasedRepository) CreateDeployment(name string) (*manager.Deployment, error) {
	d, err := func() (*manager.Deployment, error) {
		r.Lock()
		defer r.Unlock()

		exists, _ := r.GetValidDeployment(name)
		if exists != nil {
			return nil, fmt.Errorf("Deployment %s already exists", name)
		}

		d := manager.NewDeployment(name)
		d.Status = manager.CreatedStatus
		d.DeployedAt = time.Now()
		r.deployments[name] = *d
		return d, nil
	}()

	if err != nil {
		return nil, err
	}

	log.Printf("created deployment: %v", d)
	return d, nil
}

func (r *mapBasedRepository) AddManifest(deploymentName string, manifest *manager.Manifest) error {
	err := func() error {
		r.Lock()
		defer r.Unlock()

		d, err := r.GetValidDeployment(deploymentName)
		if err != nil {
			return err
		}

		l, err := r.listManifestsForDeployment(deploymentName)
		if err != nil {
			return err
		}

		// Make sure the manifest doesn't already exist, and if not, add the manifest to
		// map of manifests this deployment has
		if _, ok := l[manifest.Name]; ok {
			return fmt.Errorf("Manifest %s already exists in deployment %s", manifest.Name, deploymentName)
		}

		l[manifest.Name] = manifest
		d.LatestManifest = manifest.Name
		r.deployments[deploymentName] = *d

		return nil
	}()

	if err != nil {
		return err
	}

	log.Printf("Added manifest %s to deployment: %s", manifest.Name, deploymentName)
	return nil
}

// DeleteDeployment deletes the deployment with the supplied name.
// If forget is true, then the deployment is removed from the repository.
// Otherwise, it is marked as deleted and retained.
func (r *mapBasedRepository) DeleteDeployment(name string, forget bool) (*manager.Deployment, error) {
	d, err := func() (*manager.Deployment, error) {
		r.Lock()
		defer r.Unlock()

		d, err := r.GetValidDeployment(name)
		if err != nil {
			return nil, err
		}

		if !forget {
			d.DeletedAt = time.Now()
			d.Status = manager.DeletedStatus
			r.deployments[name] = *d
		} else {
			delete(r.deployments, name)
			delete(r.manifests, name)
			d.LatestManifest = ""
		}

		return d, nil
	}()

	if err != nil {
		return nil, err
	}

	log.Printf("deleted deployment: %v", d)
	return d, nil
}

func (r *mapBasedRepository) ListManifests(deploymentName string) (map[string]*manager.Manifest, error) {
	r.Lock()
	defer r.Unlock()

	_, err := r.GetValidDeployment(deploymentName)
	if err != nil {
		return nil, err
	}

	return r.listManifestsForDeployment(deploymentName)
}

func (r *mapBasedRepository) listManifestsForDeployment(deploymentName string) (map[string]*manager.Manifest, error) {
	l, ok := r.manifests[deploymentName]
	if !ok {
		l = make(map[string]*manager.Manifest, 0)
		r.manifests[deploymentName] = l
	}

	return l, nil
}

func (r *mapBasedRepository) GetManifest(deploymentName string, manifestName string) (*manager.Manifest, error) {
	r.Lock()
	defer r.Unlock()

	_, err := r.GetValidDeployment(deploymentName)
	if err != nil {
		return nil, err
	}

	return r.getManifestForDeployment(deploymentName, manifestName)
}

func (r *mapBasedRepository) getManifestForDeployment(deploymentName string, manifestName string) (*manager.Manifest, error) {
	l, err := r.listManifestsForDeployment(deploymentName)
	if err != nil {
		return nil, err
	}

	m, ok := l[manifestName]
	if !ok {
		return nil, fmt.Errorf("manifest %s not found in deployment %s", manifestName, deploymentName)
	}

	return m, nil
}

// GetLatestManifest returns the latest manifest for a given deployment,
// which by definition is the manifest with the largest time stamp.
func (r *mapBasedRepository) GetLatestManifest(deploymentName string) (*manager.Manifest, error) {
	r.Lock()
	defer r.Unlock()

	d, err := r.GetValidDeployment(deploymentName)
	if err != nil {
		return nil, err
	}

	return r.getManifestForDeployment(deploymentName, d.LatestManifest)
}

// ListTypes returns all types known from existing instances.
func (r *mapBasedRepository) ListTypes() []string {
	var keys []string
	for k := range r.instances {
		keys = append(keys, k)
	}

	return keys
}

// GetTypeInstances returns all instances of a given type. If type is empty,
// returns all instances for all types.
func (r *mapBasedRepository) GetTypeInstances(typeName string) []*manager.TypeInstance {
	r.Lock()
	defer r.Unlock()

	var instances []*manager.TypeInstance
	for t, dInstMap := range r.instances {
		if t == typeName || typeName == "all" {
			for _, i := range dInstMap {
				instances = append(instances, i...)
			}
		}
	}

	return instances
}

// ClearTypeInstances deletes all instances associated with the given
// deployment name from the type instance repository.
func (r *mapBasedRepository) ClearTypeInstances(deploymentName string) {
	r.Lock()
	defer r.Unlock()

	for t, dMap := range r.instances {
		delete(dMap, deploymentName)
		if len(dMap) == 0 {
			delete(r.instances, t)
		}
	}
}

// SetTypeInstances sets all type instances for a given deployment name.
//
// To clear the current set of instances first, caller should first use
// ClearTypeInstances().
func (r *mapBasedRepository) SetTypeInstances(deploymentName string, instances map[string][]*manager.TypeInstance) {
	r.Lock()
	defer r.Unlock()

	// Add each instance list to the appropriate type map.
	for t, is := range instances {
		if r.instances[t] == nil {
			r.instances[t] = make(deploymentTypeInstanceMap)
		}

		r.instances[t][deploymentName] = is
	}
}
<|endoftext|>


<commit_before>
package main

import (
    "database/sql"
    "fmt"
)

func Up_20130106222315(txn *sql.Tx) {
    fmt.Println("Hello from migration 20130106222315 Up!")
}

func Down_20130106222315(txn *sql.Tx) {
    fmt.Println("Hello from migration 20130106222315 Down!")
}
<commit_msg>gofmt examples<commit_after>package main

import (
	"database/sql"
	"fmt"
)

func Up_20130106222315(txn *sql.Tx) {
	fmt.Println("Hello from migration 20130106222315 Up!")
}

func Down_20130106222315(txn *sql.Tx) {
	fmt.Println("Hello from migration 20130106222315 Down!")
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package manager

import (
	"errors"
	"io"
	"mime/multipart"
	"path"
	"path/filepath"
	"strings"

	uploadClient "github.com/webx-top/client/upload"
	_ "github.com/webx-top/client/upload/driver"
	"github.com/webx-top/echo"

	"github.com/admpub/nging/application/handler"
	"github.com/admpub/nging/application/library/common"
	modelFile "github.com/admpub/nging/application/model/file"
	"github.com/admpub/nging/application/registry/upload"
	"github.com/admpub/nging/application/registry/upload/driver/filesystem"
	"github.com/admpub/nging/application/registry/upload/helper"
	"github.com/admpub/qrcode"
)

// 文件上传保存路径规则：
// 表名称/表行ID/文件名

// ResponseDataForUpload 根据不同的上传方式响应不同的数据格式
func ResponseDataForUpload(ctx echo.Context, field string, err error, imageURLs []string) (result echo.H, embed bool) {
	return upload.ResponserGet(field)(ctx, field, err, imageURLs)
}

var (
	StorerEngine   = filesystem.Name
	DefaultChecker = func(r *uploadClient.Result) error {
		return nil
	}
)

func File(ctx echo.Context) error {
	uploadType := ctx.Param(`type`)
	typ, _, _ := getTableInfo(uploadType)
	file := ctx.Param(`*`)
	file = filepath.Join(helper.UploadDir, typ, file)
	return ctx.File(file)
}

func SaveFilename(subdir, name, postFilename string) (string, error) {
	ext := filepath.Ext(postFilename)
	fname := name
	if len(fname) == 0 {
		var err error
		fname, err = common.UniqueID()
		if err != nil {
			return ``, err
		}
	}
	fname += ext
	return subdir + fname, nil
}

// Upload 上传文件
func Upload(ctx echo.Context) error {
	ownerType := `user`
	user := handler.User(ctx)
	var ownerID uint64
	if user != nil {
		ownerID = uint64(user.Id)
	}
	if ownerID < 1 {
		ctx.Data().SetError(ctx.E(`请先登录`))
		return ctx.Redirect(handler.URLFor(`/login`))
	}
	return UploadByOwner(ctx, ownerType, ownerID)
}

// UploadByOwner 上传文件
func UploadByOwner(ctx echo.Context, ownerType string, ownerID uint64) error {
	uploadType := ctx.Param(`type`)
	field := ctx.Query(`field`) // 上传表单file输入框名称
	pipe := ctx.Form(`pipe`)
	var (
		err      error
		fileURLs []string
	)
	if len(uploadType) == 0 {
		err = ctx.E(`请提供参数“%s”`, ctx.Path())
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}
	tableName, fieldName, defaults := getTableInfo(uploadType)
	if !upload.SubdirIsAllowed(uploadType, defaults...) {
		err = ctx.E(`参数“%s”未被登记`, uploadType)
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}
	//echo.Dump(ctx.Forms())
	newStore := upload.StorerGet(StorerEngine)
	if newStore == nil {
		err := errors.New(ctx.T(`存储引擎“%s”未被登记`, StorerEngine))
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}

	fileM := modelFile.NewFile(ctx)
	fileM.StorerName = StorerEngine
	fileM.TableId = ``
	fileM.SetFieldName(fieldName)
	fileM.SetTableName(tableName)
	fileM.OwnerId = ownerID
	fileM.OwnerType = ownerType
	fileType := ctx.Form(`filetype`)
	fileM.Type = fileType

	storer := newStore(ctx, tableName) // 使用表名称作为文件夹名
	defer storer.Close()
	var subdir, name string
	subdir, name, err = upload.CheckerGet(uploadType, defaults...)(ctx, fileM)
	if err != nil {
		return err
	}
	dbsaver := upload.DBSaverGet(uploadType, defaults...)
	checker := func(r *uploadClient.Result) error {
		extension := path.Ext(r.FileName)
		if len(r.FileType) > 0 {
			if !uploadClient.CheckTypeExtension(fileType, extension) {
				return ctx.E(`不支持将扩展名为“%v”的文件作为“%v”类型的文件来进行上传`, extension, fileType)
			}
		} else {
			r.FileType = uploadClient.FileType(uploadClient.DetectType(extension))
		}
		return DefaultChecker(r) //fileM.FnGetByMd5()
	}

	clientName := ctx.Form(`client`)
	if len(clientName) > 0 {
		result := &uploadClient.Result{}
		result.SetFileNameGenerator(func(filename string) (string, error) {
			return SaveFilename(subdir, name, filename)
		})

		client := uploadClient.Upload(ctx, clientName, result, storer, watermarkFile, checker)
		if client.GetError() != nil {
			if client.GetError() == upload.ErrExistsFile {
				client.SetError(nil)
			}
			return client.Response()
		}

		fileM.SetByUploadResult(result)

		var reader io.ReadCloser
		reader, err = storer.Get(result.SavePath)
		if reader != nil {
			defer reader.Close()
		}
		if err != nil {
			return client.SetError(err).Response()
		}
		err = dbsaver(fileM, result, reader)
		return client.SetError(err).Response()
	}
	var results uploadClient.Results
	results, err = upload.BatchUpload(
		ctx,
		`files[]`,
		func(r *uploadClient.Result) (string, error) {
			if err := checker(r); err != nil {
				return ``, err
			}
			return SaveFilename(subdir, name, r.FileName)
		},
		storer,
		func(result *uploadClient.Result, file multipart.File) error {
			fileM.Id = 0
			fileM.SetByUploadResult(result)
			return dbsaver(fileM, result, file)
		},
		watermarkFile,
	)
	datax, embed := ResponseDataForUpload(ctx, field, err, results.FileURLs())
	if err != nil {
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}

	if pipe == `deqr` { //解析二维码
		if len(results) > 0 {
			reader, err := storer.Get(results[0].SavePath)
			if reader != nil {
				defer reader.Close()
			}
			if err != nil {
				if !embed {
					datax[`raw`] = err.Error()
					return ctx.JSON(datax)
				}
				return err
			}
			raw, err := qrcode.Decode(reader, strings.TrimPrefix(path.Ext(results[0].SavePath), `.`))
			if err != nil {
				raw = err.Error()
			}
			datax[`raw`] = raw
		}
	}
	if !embed {
		return ctx.JSON(datax)
	}
	data := ctx.Data()
	data.SetData(datax)
	return ctx.JSON(data)
}

func getTableInfo(uploadType string) (tableName string, fieldName string, defaults []string) {
	return upload.GetTableInfo(uploadType)
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package manager

import (
	"errors"
	"io"
	"mime/multipart"
	"path"
	"path/filepath"
	"strings"

	uploadClient "github.com/webx-top/client/upload"
	_ "github.com/webx-top/client/upload/driver"
	"github.com/webx-top/echo"

	"github.com/admpub/nging/application/handler"
	"github.com/admpub/nging/application/library/common"
	modelFile "github.com/admpub/nging/application/model/file"
	"github.com/admpub/nging/application/registry/upload"
	"github.com/admpub/nging/application/registry/upload/driver/filesystem"
	"github.com/admpub/nging/application/registry/upload/helper"
	"github.com/admpub/qrcode"
)

// 文件上传保存路径规则：
// 表名称/表行ID/文件名

// ResponseDataForUpload 根据不同的上传方式响应不同的数据格式
func ResponseDataForUpload(ctx echo.Context, field string, err error, imageURLs []string) (result echo.H, embed bool) {
	return upload.ResponserGet(field)(ctx, field, err, imageURLs)
}

var (
	StorerEngine   = filesystem.Name
	DefaultChecker = func(r *uploadClient.Result) error {
		return nil
	}
)

func File(ctx echo.Context) error {
	uploadType := ctx.Param(`type`)
	typ, _, _ := getTableInfo(uploadType)
	file := ctx.Param(`*`)
	file = filepath.Join(helper.UploadDir, typ, file)
	return ctx.File(file)
}

// SaveFilename SaveFilename(`0/`,``,`img.jpg`)
func SaveFilename(subdir, name, postFilename string) (string, error) {
	ext := filepath.Ext(postFilename)
	fname := name
	if len(fname) == 0 {
		var err error
		fname, err = common.UniqueID()
		if err != nil {
			return ``, err
		}
	}
	fname += ext
	return subdir + fname, nil
}

// Upload 上传文件
func Upload(ctx echo.Context) error {
	ownerType := `user`
	user := handler.User(ctx)
	var ownerID uint64
	if user != nil {
		ownerID = uint64(user.Id)
	}
	if ownerID < 1 {
		ctx.Data().SetError(ctx.E(`请先登录`))
		return ctx.Redirect(handler.URLFor(`/login`))
	}
	return UploadByOwner(ctx, ownerType, ownerID)
}

// UploadByOwner 上传文件
func UploadByOwner(ctx echo.Context, ownerType string, ownerID uint64) error {
	uploadType := ctx.Param(`type`)
	field := ctx.Query(`field`) // 上传表单file输入框名称
	pipe := ctx.Form(`pipe`)
	var (
		err      error
		fileURLs []string
	)
	if len(uploadType) == 0 {
		err = ctx.E(`请提供参数“%s”`, ctx.Path())
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}
	tableName, fieldName, defaults := getTableInfo(uploadType)
	if !upload.SubdirIsAllowed(uploadType, defaults...) {
		err = ctx.E(`参数“%s”未被登记`, uploadType)
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}
	//echo.Dump(ctx.Forms())
	newStore := upload.StorerGet(StorerEngine)
	if newStore == nil {
		err := errors.New(ctx.T(`存储引擎“%s”未被登记`, StorerEngine))
		datax, embed := ResponseDataForUpload(ctx, field, err, fileURLs)
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}

	fileM := modelFile.NewFile(ctx)
	fileM.StorerName = StorerEngine
	fileM.TableId = ``
	fileM.SetFieldName(fieldName)
	fileM.SetTableName(tableName)
	fileM.OwnerId = ownerID
	fileM.OwnerType = ownerType
	fileType := ctx.Form(`filetype`)
	fileM.Type = fileType

	storer := newStore(ctx, tableName) // 使用表名称作为文件夹名
	defer storer.Close()
	var subdir, name string
	subdir, name, err = upload.CheckerGet(uploadType, defaults...)(ctx, fileM)
	if err != nil {
		return err
	}
	dbsaver := upload.DBSaverGet(uploadType, defaults...)
	checker := func(r *uploadClient.Result) error {
		extension := path.Ext(r.FileName)
		if len(r.FileType) > 0 {
			if !uploadClient.CheckTypeExtension(fileType, extension) {
				return ctx.E(`不支持将扩展名为“%v”的文件作为“%v”类型的文件来进行上传`, extension, fileType)
			}
		} else {
			r.FileType = uploadClient.FileType(uploadClient.DetectType(extension))
		}
		return DefaultChecker(r) //fileM.FnGetByMd5()
	}

	clientName := ctx.Form(`client`)
	if len(clientName) > 0 {
		result := &uploadClient.Result{}
		result.SetFileNameGenerator(func(filename string) (string, error) {
			return SaveFilename(subdir, name, filename)
		})

		client := uploadClient.Upload(ctx, clientName, result, storer, watermarkFile, checker)
		if client.GetError() != nil {
			if client.GetError() == upload.ErrExistsFile {
				client.SetError(nil)
			}
			return client.Response()
		}

		fileM.SetByUploadResult(result)

		var reader io.ReadCloser
		reader, err = storer.Get(result.SavePath)
		if reader != nil {
			defer reader.Close()
		}
		if err != nil {
			return client.SetError(err).Response()
		}
		err = dbsaver(fileM, result, reader)
		return client.SetError(err).Response()
	}
	var results uploadClient.Results
	results, err = upload.BatchUpload(
		ctx,
		`files[]`,
		func(r *uploadClient.Result) (string, error) {
			if err := checker(r); err != nil {
				return ``, err
			}
			return SaveFilename(subdir, name, r.FileName)
		},
		storer,
		func(result *uploadClient.Result, file multipart.File) error {
			fileM.Id = 0
			fileM.SetByUploadResult(result)
			return dbsaver(fileM, result, file)
		},
		watermarkFile,
	)
	datax, embed := ResponseDataForUpload(ctx, field, err, results.FileURLs())
	if err != nil {
		if !embed {
			return ctx.JSON(datax)
		}
		return err
	}

	if pipe == `deqr` { //解析二维码
		if len(results) > 0 {
			reader, err := storer.Get(results[0].SavePath)
			if reader != nil {
				defer reader.Close()
			}
			if err != nil {
				if !embed {
					datax[`raw`] = err.Error()
					return ctx.JSON(datax)
				}
				return err
			}
			raw, err := qrcode.Decode(reader, strings.TrimPrefix(path.Ext(results[0].SavePath), `.`))
			if err != nil {
				raw = err.Error()
			}
			datax[`raw`] = raw
		}
	}
	if !embed {
		return ctx.JSON(datax)
	}
	data := ctx.Data()
	data.SetData(datax)
	return ctx.JSON(data)
}

func getTableInfo(uploadType string) (tableName string, fieldName string, defaults []string) {
	return upload.GetTableInfo(uploadType)
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package broker

import (
	"fmt"
	"regexp"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	jwt "github.com/dgrijalva/jwt-go"
	"golang.org/x/net/context"
)

// ValidateToken implements the core.BrokerManagerServer interface
func (b component) ValidateToken(bctx context.Context, req *core.ValidateTokenBrokerReq) (*core.ValidateTokenBrokerRes, error) {
	b.Ctx.Debug("Handle ValidateToken request")
	if len(req.AppEUI) != 8 {
		err := errors.New(errors.Structural, "Invalid request parameters")
		b.Ctx.WithError(err).Debug("Unable to handle ValidateToken request")
		return new(core.ValidateTokenBrokerRes), err
	}
	if err := b.validateToken(bctx, req.Token, req.AppEUI); err != nil {
		b.Ctx.WithError(err).Debug("Unable to handle ValidateToken request")
		return new(core.ValidateTokenBrokerRes), err
	}
	return new(core.ValidateTokenBrokerRes), nil
}

// ValidateOTAA implements the core.BrokerManager interface
func (b component) ValidateOTAA(bctx context.Context, req *core.ValidateOTAABrokerReq) (*core.ValidateOTAABrokerRes, error) {
	b.Ctx.Debug("Handle ValidateOTAA request")

	// 1. Validate the request
	re := regexp.MustCompile("^([-\\w]+\\.?)+:\\d+$")
	if len(req.AppEUI) != 8 || !re.Match([]byte(req.NetAddress)) {
		err := errors.New(errors.Structural, "Invalid request parameters")
		b.Ctx.WithError(err).Debug("Unable to validate OTAA request")
		return new(core.ValidateOTAABrokerRes), err
	}

	// 2. Verify and validate the token
	if err := b.validateToken(bctx, req.Token, req.AppEUI); err != nil {
		return new(core.ValidateOTAABrokerRes), err
	}

	// 3. Update the internal storage
	b.Ctx.WithField("AppEUI", req.AppEUI).Debug("Request accepted by broker. Registering / Updating App.")
	err := b.AppStorage.upsert(appEntry{
		Dialer: NewDialer([]byte(req.NetAddress)),
		AppEUI: req.AppEUI,
	})
	if err != nil {
		b.Ctx.WithError(err).Debug("Error while trying to save valid request")
		return new(core.ValidateOTAABrokerRes), errors.New(errors.Operational, err)
	}

	// 4. Done.
	return new(core.ValidateOTAABrokerRes), nil
}

// UpsertABP implements the core.BrokerManager interface
func (b component) UpsertABP(bctx context.Context, req *core.UpsertABPBrokerReq) (*core.UpsertABPBrokerRes, error) {
	b.Ctx.Debug("Handle UpsertABP request")

	// 1. Validate the request
	re := regexp.MustCompile("^([-\\w]+\\.?)+:\\d+$")
	if len(req.AppEUI) != 8 || !re.Match([]byte(req.NetAddress)) || len(req.DevAddr) != 4 || len(req.NwkSKey) != 16 {
		err := errors.New(errors.Structural, "Invalid request parameters")
		b.Ctx.WithError(err).Debug("Unable to proceed Upsert ABP request")
		return new(core.UpsertABPBrokerRes), err
	}

	// 2. Verify and validate the token
	if err := b.validateToken(bctx, req.Token, req.AppEUI); err != nil {
		return new(core.UpsertABPBrokerRes), err
	}

	// 3. Update the internal storage
	b.Ctx.WithField("AppEUI", req.AppEUI).WithField("DevAddr", req.DevAddr).Debug("Request accepted by broker. Registering device.")
	var nwkSKey [16]byte
	copy(nwkSKey[:], req.NwkSKey)
	err := b.NetworkController.upsert(devEntry{
		Dialer:  NewDialer([]byte(req.NetAddress)),
		AppEUI:  req.AppEUI,
		DevEUI:  append([]byte{0, 0, 0, 0}, req.DevAddr...),
		DevAddr: req.DevAddr,
		NwkSKey: nwkSKey,
		FCntUp:  0,
	})
	if err != nil {
		b.Ctx.WithError(err).Debug("Error while trying to save valid request")
		return new(core.UpsertABPBrokerRes), errors.New(errors.Operational, err)
	}

	// 4. Done.
	return new(core.UpsertABPBrokerRes), nil
}

// validateToken verify an OAuth Bearer token pass through metadata during RPC
func (b component) validateToken(ctx context.Context, token string, appEUI []byte) error {
	parsed, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
		if b.TokenKeyProvider == nil {
			return nil, errors.New(errors.Structural, "No token provider configured")
		}
		k, err := b.TokenKeyProvider.Get(false)
		if err != nil {
			return nil, err
		}
		if k.Algorithm != token.Header["alg"] {
			return nil, errors.New(errors.Structural, fmt.Sprintf("Expected algorithm %v but got %v", k.Algorithm, token.Header["alg"]))
		}
		return []byte(k.Key), nil
	})
	if err != nil {
		return errors.New(errors.Structural, fmt.Sprintf("Unable to parse token: %s", err.Error()))
	}
	if !parsed.Valid {
		return errors.New(errors.Operational, "The token is not valid or is expired")
	}

	apps, ok := parsed.Claims["apps"].([]interface{})
	if !ok {
		return fmt.Errorf("Invalid type of apps claim: %T", parsed.Claims["apps"])
	}

	for _, a := range apps {
		if s, ok := a.(string); ok && s == fmt.Sprintf("%X", appEUI) {
			return nil
		}
	}

	return errors.New(errors.Operational, "Unauthorized")
}
<commit_msg>Clarified token invalidation error message<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package broker

import (
	"fmt"
	"regexp"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	jwt "github.com/dgrijalva/jwt-go"
	"golang.org/x/net/context"
)

// ValidateToken implements the core.BrokerManagerServer interface
func (b component) ValidateToken(bctx context.Context, req *core.ValidateTokenBrokerReq) (*core.ValidateTokenBrokerRes, error) {
	b.Ctx.Debug("Handle ValidateToken request")
	if len(req.AppEUI) != 8 {
		err := errors.New(errors.Structural, "Invalid request parameters")
		b.Ctx.WithError(err).Debug("Unable to handle ValidateToken request")
		return new(core.ValidateTokenBrokerRes), err
	}
	if err := b.validateToken(bctx, req.Token, req.AppEUI); err != nil {
		b.Ctx.WithError(err).Debug("The token is invalid")
		return new(core.ValidateTokenBrokerRes), err
	}
	return new(core.ValidateTokenBrokerRes), nil
}

// ValidateOTAA implements the core.BrokerManager interface
func (b component) ValidateOTAA(bctx context.Context, req *core.ValidateOTAABrokerReq) (*core.ValidateOTAABrokerRes, error) {
	b.Ctx.Debug("Handle ValidateOTAA request")

	// 1. Validate the request
	re := regexp.MustCompile("^([-\\w]+\\.?)+:\\d+$")
	if len(req.AppEUI) != 8 || !re.Match([]byte(req.NetAddress)) {
		err := errors.New(errors.Structural, "Invalid request parameters")
		b.Ctx.WithError(err).Debug("Unable to validate OTAA request")
		return new(core.ValidateOTAABrokerRes), err
	}

	// 2. Verify and validate the token
	if err := b.validateToken(bctx, req.Token, req.AppEUI); err != nil {
		return new(core.ValidateOTAABrokerRes), err
	}

	// 3. Update the internal storage
	b.Ctx.WithField("AppEUI", req.AppEUI).Debug("Request accepted by broker. Registering / Updating App.")
	err := b.AppStorage.upsert(appEntry{
		Dialer: NewDialer([]byte(req.NetAddress)),
		AppEUI: req.AppEUI,
	})
	if err != nil {
		b.Ctx.WithError(err).Debug("Error while trying to save valid request")
		return new(core.ValidateOTAABrokerRes), errors.New(errors.Operational, err)
	}

	// 4. Done.
	return new(core.ValidateOTAABrokerRes), nil
}

// UpsertABP implements the core.BrokerManager interface
func (b component) UpsertABP(bctx context.Context, req *core.UpsertABPBrokerReq) (*core.UpsertABPBrokerRes, error) {
	b.Ctx.Debug("Handle UpsertABP request")

	// 1. Validate the request
	re := regexp.MustCompile("^([-\\w]+\\.?)+:\\d+$")
	if len(req.AppEUI) != 8 || !re.Match([]byte(req.NetAddress)) || len(req.DevAddr) != 4 || len(req.NwkSKey) != 16 {
		err := errors.New(errors.Structural, "Invalid request parameters")
		b.Ctx.WithError(err).Debug("Unable to proceed Upsert ABP request")
		return new(core.UpsertABPBrokerRes), err
	}

	// 2. Verify and validate the token
	if err := b.validateToken(bctx, req.Token, req.AppEUI); err != nil {
		return new(core.UpsertABPBrokerRes), err
	}

	// 3. Update the internal storage
	b.Ctx.WithField("AppEUI", req.AppEUI).WithField("DevAddr", req.DevAddr).Debug("Request accepted by broker. Registering device.")
	var nwkSKey [16]byte
	copy(nwkSKey[:], req.NwkSKey)
	err := b.NetworkController.upsert(devEntry{
		Dialer:  NewDialer([]byte(req.NetAddress)),
		AppEUI:  req.AppEUI,
		DevEUI:  append([]byte{0, 0, 0, 0}, req.DevAddr...),
		DevAddr: req.DevAddr,
		NwkSKey: nwkSKey,
		FCntUp:  0,
	})
	if err != nil {
		b.Ctx.WithError(err).Debug("Error while trying to save valid request")
		return new(core.UpsertABPBrokerRes), errors.New(errors.Operational, err)
	}

	// 4. Done.
	return new(core.UpsertABPBrokerRes), nil
}

// validateToken verify an OAuth Bearer token pass through metadata during RPC
func (b component) validateToken(ctx context.Context, token string, appEUI []byte) error {
	parsed, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
		if b.TokenKeyProvider == nil {
			return nil, errors.New(errors.Structural, "No token provider configured")
		}
		k, err := b.TokenKeyProvider.Get(false)
		if err != nil {
			return nil, err
		}
		if k.Algorithm != token.Header["alg"] {
			return nil, errors.New(errors.Structural, fmt.Sprintf("Expected algorithm %v but got %v", k.Algorithm, token.Header["alg"]))
		}
		return []byte(k.Key), nil
	})
	if err != nil {
		return errors.New(errors.Structural, fmt.Sprintf("Unable to parse token: %s", err.Error()))
	}
	if !parsed.Valid {
		return errors.New(errors.Operational, "The token is not valid or is expired")
	}

	apps, ok := parsed.Claims["apps"].([]interface{})
	if !ok {
		return fmt.Errorf("Invalid type of apps claim: %T", parsed.Claims["apps"])
	}

	for _, a := range apps {
		if s, ok := a.(string); ok && s == fmt.Sprintf("%X", appEUI) {
			return nil
		}
	}

	return errors.New(errors.Operational, "Unauthorized")
}
<|endoftext|>