<commit_before>/*
Copyright 2019 The Knative Authors.

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

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

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

package gc

import (
	"context"
	"sort"
	"time"

	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"knative.dev/pkg/logging"
	pkgreconciler "knative.dev/pkg/reconciler"
	"knative.dev/serving/pkg/apis/serving"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	"knative.dev/serving/pkg/apis/serving/v1beta1"
	clientset "knative.dev/serving/pkg/client/clientset/versioned"
	configreconciler "knative.dev/serving/pkg/client/injection/reconciler/serving/v1/configuration"
	listers "knative.dev/serving/pkg/client/listers/serving/v1"
	configns "knative.dev/serving/pkg/reconciler/gc/config"
)

// reconciler implements controller.Reconciler for Garbage Collection resources.
type reconciler struct {
	client clientset.Interface

	// listers index properties about resources
	revisionLister listers.RevisionLister
}

// Check that our reconciler implements configreconciler.Interface
var _ configreconciler.Interface = (*reconciler)(nil)

func (c *reconciler) ReconcileKind(ctx context.Context, config *v1.Configuration) pkgreconciler.Event {
	cfg := configns.FromContext(ctx).RevisionGC
	logger := logging.FromContext(ctx)

	selector := labels.SelectorFromSet(labels.Set{serving.ConfigurationLabelKey: config.Name})
	revs, err := c.revisionLister.Revisions(config.Namespace).List(selector)
	if err != nil {
		return err
	}

	gcSkipOffset := cfg.StaleRevisionMinimumGenerations

	if gcSkipOffset >= int64(len(revs)) {
		return nil
	}

	// Sort by creation timestamp descending
	sort.Slice(revs, func(i, j int) bool {
		return revs[j].CreationTimestamp.Before(&revs[i].CreationTimestamp)
	})

	for _, rev := range revs[gcSkipOffset:] {
		if isRevisionStale(ctx, rev, config) {
			err := c.client.ServingV1().Revisions(rev.Namespace).Delete(rev.Name, &metav1.DeleteOptions{})
			if err != nil {
				logger.With(zap.Error(err)).Errorf("Failed to delete stale revision %q", rev.Name)
				continue
			}
		}
	}
	return nil
}

func isRevisionStale(ctx context.Context, rev *v1.Revision, config *v1.Configuration) bool {
	if config.Status.LatestReadyRevisionName == rev.Name {
		return false
	}

	cfg := configns.FromContext(ctx).RevisionGC
	logger := logging.FromContext(ctx)

	curTime := time.Now()
	if rev.ObjectMeta.CreationTimestamp.Add(cfg.StaleRevisionCreateDelay).After(curTime) {
		// Revision was created sooner than staleRevisionCreateDelay. Ignore it.
		return false
	}

	lastPin, err := rev.GetLastPinned()
	if err != nil {
		if err.(v1.LastPinnedParseError).Type != v1.AnnotationParseErrorTypeMissing {
			logger.Errorw("Failed to determine revision last pinned", zap.Error(err))
		} else {
			// Revision was never pinned and its RevisionConditionReady is not true after staleRevisionCreateDelay.
			// It usually happens when ksvc was deployed with wrong configuration.
			rc := rev.Status.GetCondition(v1beta1.RevisionConditionReady)
			if rc == nil || rc.Status != corev1.ConditionTrue {
				return true
			}
		}
		return false
	}

	ret := lastPin.Add(cfg.StaleRevisionTimeout).Before(curTime)
	if ret {
		logger.Infof("Detected stale revision %v with creation time %v and lastPinned time %v.", rev.ObjectMeta.Name, rev.ObjectMeta.CreationTimestamp, lastPin)
	}
	return ret
}
<commit_msg>drop v1beta1 usage (#8404)<commit_after>/*
Copyright 2019 The Knative Authors.

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

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

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

package gc

import (
	"context"
	"sort"
	"time"

	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"knative.dev/pkg/logging"
	pkgreconciler "knative.dev/pkg/reconciler"
	"knative.dev/serving/pkg/apis/serving"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	clientset "knative.dev/serving/pkg/client/clientset/versioned"
	configreconciler "knative.dev/serving/pkg/client/injection/reconciler/serving/v1/configuration"
	listers "knative.dev/serving/pkg/client/listers/serving/v1"
	configns "knative.dev/serving/pkg/reconciler/gc/config"
)

// reconciler implements controller.Reconciler for Garbage Collection resources.
type reconciler struct {
	client clientset.Interface

	// listers index properties about resources
	revisionLister listers.RevisionLister
}

// Check that our reconciler implements configreconciler.Interface
var _ configreconciler.Interface = (*reconciler)(nil)

func (c *reconciler) ReconcileKind(ctx context.Context, config *v1.Configuration) pkgreconciler.Event {
	cfg := configns.FromContext(ctx).RevisionGC
	logger := logging.FromContext(ctx)

	selector := labels.SelectorFromSet(labels.Set{serving.ConfigurationLabelKey: config.Name})
	revs, err := c.revisionLister.Revisions(config.Namespace).List(selector)
	if err != nil {
		return err
	}

	gcSkipOffset := cfg.StaleRevisionMinimumGenerations

	if gcSkipOffset >= int64(len(revs)) {
		return nil
	}

	// Sort by creation timestamp descending
	sort.Slice(revs, func(i, j int) bool {
		return revs[j].CreationTimestamp.Before(&revs[i].CreationTimestamp)
	})

	for _, rev := range revs[gcSkipOffset:] {
		if isRevisionStale(ctx, rev, config) {
			err := c.client.ServingV1().Revisions(rev.Namespace).Delete(rev.Name, &metav1.DeleteOptions{})
			if err != nil {
				logger.With(zap.Error(err)).Errorf("Failed to delete stale revision %q", rev.Name)
				continue
			}
		}
	}
	return nil
}

func isRevisionStale(ctx context.Context, rev *v1.Revision, config *v1.Configuration) bool {
	if config.Status.LatestReadyRevisionName == rev.Name {
		return false
	}

	cfg := configns.FromContext(ctx).RevisionGC
	logger := logging.FromContext(ctx)

	curTime := time.Now()
	if rev.ObjectMeta.CreationTimestamp.Add(cfg.StaleRevisionCreateDelay).After(curTime) {
		// Revision was created sooner than staleRevisionCreateDelay. Ignore it.
		return false
	}

	lastPin, err := rev.GetLastPinned()
	if err != nil {
		if err.(v1.LastPinnedParseError).Type != v1.AnnotationParseErrorTypeMissing {
			logger.Errorw("Failed to determine revision last pinned", zap.Error(err))
		} else {
			// Revision was never pinned and its RevisionConditionReady is not true after staleRevisionCreateDelay.
			// It usually happens when ksvc was deployed with wrong configuration.
			rc := rev.Status.GetCondition(v1.RevisionConditionReady)
			if rc == nil || rc.Status != corev1.ConditionTrue {
				return true
			}
		}
		return false
	}

	ret := lastPin.Add(cfg.StaleRevisionTimeout).Before(curTime)
	if ret {
		logger.Infof("Detected stale revision %v with creation time %v and lastPinned time %v.", rev.ObjectMeta.Name, rev.ObjectMeta.CreationTimestamp, lastPin)
	}
	return ret
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/go-kit/kit/log/term"
	metrics "github.com/rcrowley/go-metrics"

	tmtypes "github.com/tendermint/tendermint/types"
	"github.com/tendermint/tmlibs/log"
	"github.com/tendermint/tools/tm-monitor/monitor"
)

var version = "0.2.0"

var logger = log.NewNopLogger()

type statistics struct {
	BlockTimeSample    metrics.Histogram
	TxThroughputSample metrics.Histogram
	BlockLatency       metrics.Histogram
}

func main() {
	var duration, txsRate, connections int
	var verbose bool

	flag.IntVar(&connections, "c", 1, "Connections to keep open per endpoint")
	flag.IntVar(&duration, "T", 10, "Exit after the specified amount of time in seconds")
	flag.IntVar(&txsRate, "r", 1000, "Txs per second to send in a connection")
	flag.BoolVar(&verbose, "v", false, "Verbose output")

	flag.Usage = func() {
		fmt.Println(`Tendermint blockchain benchmarking tool.

Usage:
	tm-bench [-c 1] [-T 10] [-r 1000] [endpoints]

Examples:
	tm-bench localhost:46657`)
		fmt.Println("Flags:")
		flag.PrintDefaults()
	}

	flag.Parse()

	if flag.NArg() == 0 {
		flag.Usage()
		os.Exit(1)
	}

	if verbose {
		// Color errors red
		colorFn := func(keyvals ...interface{}) term.FgBgColor {
			for i := 1; i < len(keyvals); i += 2 {
				if _, ok := keyvals[i].(error); ok {
					return term.FgBgColor{Fg: term.White, Bg: term.Red}
				}
			}
			return term.FgBgColor{}
		}
		logger = log.NewTMLoggerWithColorFn(log.NewSyncWriter(os.Stdout), colorFn)
	}

	fmt.Printf("Running %ds test @ %s\n", duration, flag.Arg(0))

	endpoints := strings.Split(flag.Arg(0), ",")

	blockCh := make(chan tmtypes.Header, 100)
	blockLatencyCh := make(chan float64, 100)

	nodes := startNodes(endpoints, blockCh, blockLatencyCh)

	transacters := startTransacters(endpoints, connections, txsRate)

	stats := &statistics{
		BlockTimeSample:    metrics.NewHistogram(metrics.NewUniformSample(1000)),
		TxThroughputSample: metrics.NewHistogram(metrics.NewUniformSample(1000)),
		BlockLatency:       metrics.NewHistogram(metrics.NewUniformSample(1000)),
	}

	lastBlockHeight := -1

	durationTimer := time.After(time.Duration(duration) * time.Second)
	ticker := time.NewTicker(1 * time.Second)
	var blocks, txs int
	for {
		select {
		case b := <-blockCh:
			if lastBlockHeight < b.Height {
				blocks++
				txs += b.NumTxs
				lastBlockHeight = b.Height
			}
		case l := <-blockLatencyCh:
			stats.BlockLatency.Update(int64(l))
		case <-ticker.C:
			stats.BlockTimeSample.Update(int64(blocks))
			stats.TxThroughputSample.Update(int64(txs))
			blocks = 0
			txs = 0
		case <-durationTimer:
			for _, t := range transacters {
				t.Stop()
			}

			printStatistics(stats)

			for _, n := range nodes {
				n.Stop()
			}
			return
		}
	}
}

func startNodes(endpoints []string, blockCh chan<- tmtypes.Header, blockLatencyCh chan<- float64) []*monitor.Node {
	nodes := make([]*monitor.Node, len(endpoints))

	for i, e := range endpoints {
		n := monitor.NewNode(e)
		n.SetLogger(logger.With("node", e))
		n.SendBlocksTo(blockCh)
		n.SendBlockLatenciesTo(blockLatencyCh)
		if err := n.Start(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		nodes[i] = n
	}

	return nodes
}

func startTransacters(endpoints []string, connections int, txsRate int) []*transacter {
	transacters := make([]*transacter, len(endpoints))

	for i, e := range endpoints {
		t := newTransacter(e, connections, txsRate)
		t.SetLogger(logger)
		if err := t.Start(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		transacters[i] = t
	}

	return transacters
}

func printStatistics(stats *statistics) {
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 5, ' ', 0)
	fmt.Fprintln(w, "Stats\tAvg\tStdev\tMax\t")
	fmt.Fprintln(w, fmt.Sprintf("Block latency\t%.2fms\t%.2fms\t%dms\t",
		stats.BlockLatency.Mean()/1000000.0,
		stats.BlockLatency.StdDev()/1000000.0,
		stats.BlockLatency.Max()/1000000))
	fmt.Fprintln(w, fmt.Sprintf("Blocks/sec\t%.3f\t%.3f\t%d\t",
		stats.BlockTimeSample.Mean(),
		stats.BlockTimeSample.StdDev(),
		stats.BlockTimeSample.Max()))
	fmt.Fprintln(w, fmt.Sprintf("Txs/sec\t%.0f\t%.0f\t%d\t",
		stats.TxThroughputSample.Mean(),
		stats.TxThroughputSample.StdDev(),
		stats.TxThroughputSample.Max()))
	w.Flush()
}
<commit_msg>update tm-bench version to 0.2.1<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/go-kit/kit/log/term"
	metrics "github.com/rcrowley/go-metrics"

	tmtypes "github.com/tendermint/tendermint/types"
	"github.com/tendermint/tmlibs/log"
	"github.com/tendermint/tools/tm-monitor/monitor"
)

var version = "0.2.1"

var logger = log.NewNopLogger()

type statistics struct {
	BlockTimeSample    metrics.Histogram
	TxThroughputSample metrics.Histogram
	BlockLatency       metrics.Histogram
}

func main() {
	var duration, txsRate, connections int
	var verbose bool

	flag.IntVar(&connections, "c", 1, "Connections to keep open per endpoint")
	flag.IntVar(&duration, "T", 10, "Exit after the specified amount of time in seconds")
	flag.IntVar(&txsRate, "r", 1000, "Txs per second to send in a connection")
	flag.BoolVar(&verbose, "v", false, "Verbose output")

	flag.Usage = func() {
		fmt.Println(`Tendermint blockchain benchmarking tool.

Usage:
	tm-bench [-c 1] [-T 10] [-r 1000] [endpoints]

Examples:
	tm-bench localhost:46657`)
		fmt.Println("Flags:")
		flag.PrintDefaults()
	}

	flag.Parse()

	if flag.NArg() == 0 {
		flag.Usage()
		os.Exit(1)
	}

	if verbose {
		// Color errors red
		colorFn := func(keyvals ...interface{}) term.FgBgColor {
			for i := 1; i < len(keyvals); i += 2 {
				if _, ok := keyvals[i].(error); ok {
					return term.FgBgColor{Fg: term.White, Bg: term.Red}
				}
			}
			return term.FgBgColor{}
		}
		logger = log.NewTMLoggerWithColorFn(log.NewSyncWriter(os.Stdout), colorFn)
	}

	fmt.Printf("Running %ds test @ %s\n", duration, flag.Arg(0))

	endpoints := strings.Split(flag.Arg(0), ",")

	blockCh := make(chan tmtypes.Header, 100)
	blockLatencyCh := make(chan float64, 100)

	nodes := startNodes(endpoints, blockCh, blockLatencyCh)

	transacters := startTransacters(endpoints, connections, txsRate)

	stats := &statistics{
		BlockTimeSample:    metrics.NewHistogram(metrics.NewUniformSample(1000)),
		TxThroughputSample: metrics.NewHistogram(metrics.NewUniformSample(1000)),
		BlockLatency:       metrics.NewHistogram(metrics.NewUniformSample(1000)),
	}

	lastBlockHeight := -1

	durationTimer := time.After(time.Duration(duration) * time.Second)
	ticker := time.NewTicker(1 * time.Second)
	var blocks, txs int
	for {
		select {
		case b := <-blockCh:
			if lastBlockHeight < b.Height {
				blocks++
				txs += b.NumTxs
				lastBlockHeight = b.Height
			}
		case l := <-blockLatencyCh:
			stats.BlockLatency.Update(int64(l))
		case <-ticker.C:
			stats.BlockTimeSample.Update(int64(blocks))
			stats.TxThroughputSample.Update(int64(txs))
			blocks = 0
			txs = 0
		case <-durationTimer:
			for _, t := range transacters {
				t.Stop()
			}

			printStatistics(stats)

			for _, n := range nodes {
				n.Stop()
			}
			return
		}
	}
}

func startNodes(endpoints []string, blockCh chan<- tmtypes.Header, blockLatencyCh chan<- float64) []*monitor.Node {
	nodes := make([]*monitor.Node, len(endpoints))

	for i, e := range endpoints {
		n := monitor.NewNode(e)
		n.SetLogger(logger.With("node", e))
		n.SendBlocksTo(blockCh)
		n.SendBlockLatenciesTo(blockLatencyCh)
		if err := n.Start(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		nodes[i] = n
	}

	return nodes
}

func startTransacters(endpoints []string, connections int, txsRate int) []*transacter {
	transacters := make([]*transacter, len(endpoints))

	for i, e := range endpoints {
		t := newTransacter(e, connections, txsRate)
		t.SetLogger(logger)
		if err := t.Start(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		transacters[i] = t
	}

	return transacters
}

func printStatistics(stats *statistics) {
	w := tabwriter.NewWriter(os.Stdout, 0, 0, 5, ' ', 0)
	fmt.Fprintln(w, "Stats\tAvg\tStdev\tMax\t")
	fmt.Fprintln(w, fmt.Sprintf("Block latency\t%.2fms\t%.2fms\t%dms\t",
		stats.BlockLatency.Mean()/1000000.0,
		stats.BlockLatency.StdDev()/1000000.0,
		stats.BlockLatency.Max()/1000000))
	fmt.Fprintln(w, fmt.Sprintf("Blocks/sec\t%.3f\t%.3f\t%d\t",
		stats.BlockTimeSample.Mean(),
		stats.BlockTimeSample.StdDev(),
		stats.BlockTimeSample.Max()))
	fmt.Fprintln(w, fmt.Sprintf("Txs/sec\t%.0f\t%.0f\t%d\t",
		stats.TxThroughputSample.Mean(),
		stats.TxThroughputSample.StdDev(),
		stats.TxThroughputSample.Max()))
	w.Flush()
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSVpnGateway_basic(t *testing.T) {
	var v, v2 ec2.VpnGateway

	testNotEqual := func(*terraform.State) error {
		if len(v.VpcAttachments) == 0 {
			return fmt.Errorf("VPN Gateway A is not attached")
		}
		if len(v2.VpcAttachments) == 0 {
			return fmt.Errorf("VPN Gateway B is not attached")
		}

		id1 := v.VpcAttachments[0].VpcId
		id2 := v2.VpcAttachments[0].VpcId
		if id1 == id2 {
			return fmt.Errorf("Both attachment IDs are the same")
		}

		return nil
	}

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccVpnGatewayConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &v),
				),
			},

			resource.TestStep{
				Config: testAccVpnGatewayConfigChangeVPC,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &v2),
					testNotEqual,
				),
			},
		},
	})
}

func TestAccAWSVpnGateway_disappears(t *testing.T) {
	var v ec2.VpnGateway

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccVpnGatewayConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists("aws_vpn_gateway.foo", &v),
					testAccAWSVpnGatewayDisappears(&v),
				),
				ExpectNonEmptyPlan: true,
			},
		},
	})
}

func TestAccAWSVpnGateway_reattach(t *testing.T) {
	var vpc1, vpc2 ec2.Vpc
	var vgw1, vgw2 ec2.VpnGateway

	testAttachmentFunc := func(vgw *ec2.VpnGateway, vpc *ec2.Vpc) func(*terraform.State) error {
		return func(*terraform.State) error {
			if len(vgw.VpcAttachments) == 0 {
				return fmt.Errorf("VPN Gateway %q has no VPC attachments.",
					*vgw.VpnGatewayId)
			}

			if len(vgw.VpcAttachments) > 1 {
				count := 0
				for _, v := range vgw.VpcAttachments {
					if *v.State == "attached" {
						count += 1
					}
				}
				if count > 1 {
					return fmt.Errorf(
						"VPN Gateway %q has an unexpected number of VPC attachments (more than 1): %#v",
						*vgw.VpnGatewayId, vgw.VpcAttachments)
				}
			}

			if *vgw.VpcAttachments[0].State != "attached" {
				return fmt.Errorf("Expected VPN Gateway %q to be attached.",
					*vgw.VpnGatewayId)
			}

			if *vgw.VpcAttachments[0].VpcId != *vpc.VpcId {
				return fmt.Errorf("Expected VPN Gateway %q to be attached to VPC %q, but got: %q",
					*vgw.VpnGatewayId, *vpc.VpcId, *vgw.VpcAttachments[0].VpcId)
			}
			return nil
		}
	}

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigReattach,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpcExists("aws_vpc.foo", &vpc1),
					testAccCheckVpcExists("aws_vpc.bar", &vpc2),
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &vgw1),
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.bar", &vgw2),
					testAttachmentFunc(&vgw1, &vpc1),
					testAttachmentFunc(&vgw2, &vpc2),
				),
			},
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigReattachChange,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &vgw1),
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.bar", &vgw2),
					testAttachmentFunc(&vgw2, &vpc1),
					testAttachmentFunc(&vgw1, &vpc2),
				),
			},
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigReattach,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &vgw1),
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.bar", &vgw2),
					testAttachmentFunc(&vgw1, &vpc1),
					testAttachmentFunc(&vgw2, &vpc2),
				),
			},
		},
	})
}

func TestAccAWSVpnGateway_delete(t *testing.T) {
	var vpnGateway ec2.VpnGateway

	testDeleted := func(r string) resource.TestCheckFunc {
		return func(s *terraform.State) error {
			_, ok := s.RootModule().Resources[r]
			if ok {
				return fmt.Errorf("VPN Gateway %q should have been deleted.", r)
			}
			return nil
		}
	}

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccVpnGatewayConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists("aws_vpn_gateway.foo", &vpnGateway)),
			},
			resource.TestStep{
				Config: testAccNoVpnGatewayConfig,
				Check:  resource.ComposeTestCheckFunc(testDeleted("aws_vpn_gateway.foo")),
			},
		},
	})
}

func TestAccAWSVpnGateway_tags(t *testing.T) {
	var v ec2.VpnGateway

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigTags,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists("aws_vpn_gateway.foo", &v),
					testAccCheckTags(&v.Tags, "foo", "bar"),
				),
			},
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigTagsUpdate,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists("aws_vpn_gateway.foo", &v),
					testAccCheckTags(&v.Tags, "foo", ""),
					testAccCheckTags(&v.Tags, "bar", "baz"),
				),
			},
		},
	})
}

func testAccAWSVpnGatewayDisappears(gateway *ec2.VpnGateway) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.Meta().(*AWSClient).ec2conn

		_, err := conn.DetachVpnGateway(&ec2.DetachVpnGatewayInput{
			VpnGatewayId: gateway.VpnGatewayId,
			VpcId:        gateway.VpcAttachments[0].VpcId,
		})
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok {
				if ec2err.Code() == "InvalidVpnGatewayID.NotFound" {
					return nil
				} else if ec2err.Code() == "InvalidVpnGatewayAttachment.NotFound" {
					return nil
				}
			}

			if err != nil {
				return err
			}
		}

		opts := &ec2.DeleteVpnGatewayInput{
			VpnGatewayId: gateway.VpnGatewayId,
		}
		if _, err := conn.DeleteVpnGateway(opts); err != nil {
			return err
		}
		return resource.Retry(40*time.Minute, func() *resource.RetryError {
			opts := &ec2.DescribeVpnGatewaysInput{
				VpnGatewayIds: []*string{gateway.VpnGatewayId},
			}
			resp, err := conn.DescribeVpnGateways(opts)
			if err != nil {
				cgw, ok := err.(awserr.Error)
				if ok && cgw.Code() == "InvalidVpnGatewayID.NotFound" {
					return nil
				}
				if ok && cgw.Code() == "IncorrectState" {
					return resource.RetryableError(fmt.Errorf(
						"Waiting for VPN Gateway to be in the correct state: %v", gateway.VpnGatewayId))
				}
				return resource.NonRetryableError(
					fmt.Errorf("Error retrieving VPN Gateway: %s", err))
			}
			if *resp.VpnGateways[0].State == "deleted" {
				return nil
			}
			return resource.RetryableError(fmt.Errorf(
				"Waiting for VPN Gateway: %v", gateway.VpnGatewayId))
		})
	}
}

func testAccCheckVpnGatewayDestroy(s *terraform.State) error {
	ec2conn := testAccProvider.Meta().(*AWSClient).ec2conn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_vpn_gateway" {
			continue
		}

		// Try to find the resource
		resp, err := ec2conn.DescribeVpnGateways(&ec2.DescribeVpnGatewaysInput{
			VpnGatewayIds: []*string{aws.String(rs.Primary.ID)},
		})
		if err == nil {
			var v *ec2.VpnGateway
			for _, g := range resp.VpnGateways {
				if *g.VpnGatewayId == rs.Primary.ID {
					v = g
				}
			}

			if v == nil {
				// wasn't found
				return nil
			}

			if *v.State != "deleted" {
				return fmt.Errorf("Expected VPN Gateway to be in deleted state, but was not: %s", v)
			}
			return nil
		}

		// Verify the error is what we want
		ec2err, ok := err.(awserr.Error)
		if !ok {
			return err
		}
		if ec2err.Code() != "InvalidVpnGatewayID.NotFound" {
			return err
		}
	}

	return nil
}

func testAccCheckVpnGatewayExists(n string, ig *ec2.VpnGateway) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No ID is set")
		}

		ec2conn := testAccProvider.Meta().(*AWSClient).ec2conn
		resp, err := ec2conn.DescribeVpnGateways(&ec2.DescribeVpnGatewaysInput{
			VpnGatewayIds: []*string{aws.String(rs.Primary.ID)},
		})
		if err != nil {
			return err
		}
		if len(resp.VpnGateways) == 0 {
			return fmt.Errorf("VPN Gateway not found")
		}

		*ig = *resp.VpnGateways[0]

		return nil
	}
}

const testAccNoVpnGatewayConfig = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}
`

const testAccVpnGatewayConfig = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.foo.id}"
}
`

const testAccVpnGatewayConfigChangeVPC = `
resource "aws_vpc" "bar" {
	cidr_block = "10.2.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.bar.id}"
}
`

const testAccCheckVpnGatewayConfigTags = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.foo.id}"
	tags {
		foo = "bar"
	}
}
`

const testAccCheckVpnGatewayConfigTagsUpdate = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.foo.id}"
	tags {
		bar = "baz"
	}
}
`

const testAccCheckVpnGatewayConfigReattach = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpc" "bar" {
	cidr_block = "10.2.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.foo.id}"
}

resource "aws_vpn_gateway" "bar" {
	vpc_id = "${aws_vpc.bar.id}"
}
`

const testAccCheckVpnGatewayConfigReattachChange = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpc" "bar" {
	cidr_block = "10.2.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.bar.id}"
}

resource "aws_vpn_gateway" "bar" {
	vpc_id = "${aws_vpc.foo.id}"
}
`
<commit_msg>provider/aws: Clear up test failure for TestAccAWSVpnGateway_disappears<commit_after>package aws

import (
	"fmt"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSVpnGateway_basic(t *testing.T) {
	var v, v2 ec2.VpnGateway

	testNotEqual := func(*terraform.State) error {
		if len(v.VpcAttachments) == 0 {
			return fmt.Errorf("VPN Gateway A is not attached")
		}
		if len(v2.VpcAttachments) == 0 {
			return fmt.Errorf("VPN Gateway B is not attached")
		}

		id1 := v.VpcAttachments[0].VpcId
		id2 := v2.VpcAttachments[0].VpcId
		if id1 == id2 {
			return fmt.Errorf("Both attachment IDs are the same")
		}

		return nil
	}

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccVpnGatewayConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &v),
				),
			},

			resource.TestStep{
				Config: testAccVpnGatewayConfigChangeVPC,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &v2),
					testNotEqual,
				),
			},
		},
	})
}

func TestAccAWSVpnGateway_disappears(t *testing.T) {
	var v ec2.VpnGateway

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccVpnGatewayConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists("aws_vpn_gateway.foo", &v),
					testAccAWSVpnGatewayDisappears(&v),
				),
				ExpectNonEmptyPlan: true,
			},
		},
	})
}

func TestAccAWSVpnGateway_reattach(t *testing.T) {
	var vpc1, vpc2 ec2.Vpc
	var vgw1, vgw2 ec2.VpnGateway

	testAttachmentFunc := func(vgw *ec2.VpnGateway, vpc *ec2.Vpc) func(*terraform.State) error {
		return func(*terraform.State) error {
			if len(vgw.VpcAttachments) == 0 {
				return fmt.Errorf("VPN Gateway %q has no VPC attachments.",
					*vgw.VpnGatewayId)
			}

			if len(vgw.VpcAttachments) > 1 {
				count := 0
				for _, v := range vgw.VpcAttachments {
					if *v.State == "attached" {
						count += 1
					}
				}
				if count > 1 {
					return fmt.Errorf(
						"VPN Gateway %q has an unexpected number of VPC attachments (more than 1): %#v",
						*vgw.VpnGatewayId, vgw.VpcAttachments)
				}
			}

			if *vgw.VpcAttachments[0].State != "attached" {
				return fmt.Errorf("Expected VPN Gateway %q to be attached.",
					*vgw.VpnGatewayId)
			}

			if *vgw.VpcAttachments[0].VpcId != *vpc.VpcId {
				return fmt.Errorf("Expected VPN Gateway %q to be attached to VPC %q, but got: %q",
					*vgw.VpnGatewayId, *vpc.VpcId, *vgw.VpcAttachments[0].VpcId)
			}
			return nil
		}
	}

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigReattach,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpcExists("aws_vpc.foo", &vpc1),
					testAccCheckVpcExists("aws_vpc.bar", &vpc2),
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &vgw1),
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.bar", &vgw2),
					testAttachmentFunc(&vgw1, &vpc1),
					testAttachmentFunc(&vgw2, &vpc2),
				),
			},
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigReattachChange,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &vgw1),
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.bar", &vgw2),
					testAttachmentFunc(&vgw2, &vpc1),
					testAttachmentFunc(&vgw1, &vpc2),
				),
			},
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigReattach,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.foo", &vgw1),
					testAccCheckVpnGatewayExists(
						"aws_vpn_gateway.bar", &vgw2),
					testAttachmentFunc(&vgw1, &vpc1),
					testAttachmentFunc(&vgw2, &vpc2),
				),
			},
		},
	})
}

func TestAccAWSVpnGateway_delete(t *testing.T) {
	var vpnGateway ec2.VpnGateway

	testDeleted := func(r string) resource.TestCheckFunc {
		return func(s *terraform.State) error {
			_, ok := s.RootModule().Resources[r]
			if ok {
				return fmt.Errorf("VPN Gateway %q should have been deleted.", r)
			}
			return nil
		}
	}

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccVpnGatewayConfig,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists("aws_vpn_gateway.foo", &vpnGateway)),
			},
			resource.TestStep{
				Config: testAccNoVpnGatewayConfig,
				Check:  resource.ComposeTestCheckFunc(testDeleted("aws_vpn_gateway.foo")),
			},
		},
	})
}

func TestAccAWSVpnGateway_tags(t *testing.T) {
	var v ec2.VpnGateway

	resource.Test(t, resource.TestCase{
		PreCheck:      func() { testAccPreCheck(t) },
		IDRefreshName: "aws_vpn_gateway.foo",
		Providers:     testAccProviders,
		CheckDestroy:  testAccCheckVpnGatewayDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigTags,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists("aws_vpn_gateway.foo", &v),
					testAccCheckTags(&v.Tags, "foo", "bar"),
				),
			},
			resource.TestStep{
				Config: testAccCheckVpnGatewayConfigTagsUpdate,
				Check: resource.ComposeTestCheckFunc(
					testAccCheckVpnGatewayExists("aws_vpn_gateway.foo", &v),
					testAccCheckTags(&v.Tags, "foo", ""),
					testAccCheckTags(&v.Tags, "bar", "baz"),
				),
			},
		},
	})
}

func testAccAWSVpnGatewayDisappears(gateway *ec2.VpnGateway) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		conn := testAccProvider.Meta().(*AWSClient).ec2conn

		_, err := conn.DetachVpnGateway(&ec2.DetachVpnGatewayInput{
			VpnGatewayId: gateway.VpnGatewayId,
			VpcId:        gateway.VpcAttachments[0].VpcId,
		})
		if err != nil {
			ec2err, ok := err.(awserr.Error)
			if ok {
				if ec2err.Code() == "InvalidVpnGatewayID.NotFound" {
					return nil
				} else if ec2err.Code() == "InvalidVpnGatewayAttachment.NotFound" {
					return nil
				}
			}

			if err != nil {
				return err
			}
		}

		opts := &ec2.DeleteVpnGatewayInput{
			VpnGatewayId: gateway.VpnGatewayId,
		}
		if _, err := conn.DeleteVpnGateway(opts); err != nil {
			return err
		}
		return resource.Retry(40*time.Minute, func() *resource.RetryError {
			opts := &ec2.DescribeVpnGatewaysInput{
				VpnGatewayIds: []*string{gateway.VpnGatewayId},
			}
			resp, err := conn.DescribeVpnGateways(opts)
			if err != nil {
				cgw, ok := err.(awserr.Error)
				if ok && cgw.Code() == "InvalidVpnGatewayID.NotFound" {
					return nil
				}
				if ok && cgw.Code() == "IncorrectState" {
					return resource.RetryableError(fmt.Errorf(
						"Waiting for VPN Gateway to be in the correct state: %v", gateway.VpnGatewayId))
				}
				return resource.NonRetryableError(
					fmt.Errorf("Error retrieving VPN Gateway: %s", err))
			}
			if *resp.VpnGateways[0].State == "deleted" {
				return nil
			}
			return resource.RetryableError(fmt.Errorf(
				"Waiting for VPN Gateway: %v", gateway.VpnGatewayId))
		})
	}
}

func testAccCheckVpnGatewayDestroy(s *terraform.State) error {
	ec2conn := testAccProvider.Meta().(*AWSClient).ec2conn

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_vpn_gateway" {
			continue
		}

		// Try to find the resource
		resp, err := ec2conn.DescribeVpnGateways(&ec2.DescribeVpnGatewaysInput{
			VpnGatewayIds: []*string{aws.String(rs.Primary.ID)},
		})
		if err == nil {
			var v *ec2.VpnGateway
			for _, g := range resp.VpnGateways {
				if *g.VpnGatewayId == rs.Primary.ID {
					v = g
				}
			}

			if v == nil {
				// wasn't found
				return nil
			}

			if *v.State != "deleted" {
				return fmt.Errorf("Expected VPN Gateway to be in deleted state, but was not: %s", v)
			}
			return nil
		}

		// Verify the error is what we want
		ec2err, ok := err.(awserr.Error)
		if !ok {
			return err
		}
		if ec2err.Code() != "InvalidVpnGatewayID.NotFound" {
			return err
		}
	}

	return nil
}

func testAccCheckVpnGatewayExists(n string, ig *ec2.VpnGateway) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No ID is set")
		}

		ec2conn := testAccProvider.Meta().(*AWSClient).ec2conn
		resp, err := ec2conn.DescribeVpnGateways(&ec2.DescribeVpnGatewaysInput{
			VpnGatewayIds: []*string{aws.String(rs.Primary.ID)},
		})
		if err != nil {
			return err
		}
		if len(resp.VpnGateways) == 0 {
			return fmt.Errorf("VPN Gateway not found")
		}

		*ig = *resp.VpnGateways[0]

		return nil
	}
}

const testAccNoVpnGatewayConfig = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}
`

const testAccVpnGatewayConfig = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.foo.id}"
}
`

const testAccVpnGatewayConfigChangeVPC = `
resource "aws_vpc" "bar" {
	cidr_block = "10.2.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.bar.id}"
}
`

const testAccCheckVpnGatewayConfigTags = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.foo.id}"
	tags {
		foo = "bar"
	}
}
`

const testAccCheckVpnGatewayConfigTagsUpdate = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.foo.id}"
	tags {
		bar = "baz"
	}
}
`

const testAccCheckVpnGatewayConfigReattach = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpc" "bar" {
	cidr_block = "10.2.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.foo.id}"
}

resource "aws_vpn_gateway" "bar" {
	vpc_id = "${aws_vpc.bar.id}"
}
`

const testAccCheckVpnGatewayConfigReattachChange = `
resource "aws_vpc" "foo" {
	cidr_block = "10.1.0.0/16"
}

resource "aws_vpc" "bar" {
	cidr_block = "10.2.0.0/16"
}

resource "aws_vpn_gateway" "foo" {
	vpc_id = "${aws_vpc.bar.id}"
}

resource "aws_vpn_gateway" "bar" {
	vpc_id = "${aws_vpc.foo.id}"
}
`
<|endoftext|>


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

package server

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/parser"
	"github.com/pingcap/parser/ast"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/parser/terror"
	"github.com/pingcap/tidb/domain"
	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/statistics/handle"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/util/sqlexec"
)

type sqlInfoFetcher struct {
	store tikv.Storage
	do    *domain.Domain
	s     session.Session
}

type tableNamePair struct {
	DBName    string
	TableName string
}

type tableNameExtractor struct {
	curDB string
	names map[tableNamePair]struct{}
}

func (tne *tableNameExtractor) Enter(in ast.Node) (ast.Node, bool) {
	if _, ok := in.(*ast.TableName); ok {
		return in, true
	}
	return in, false
}

func (tne *tableNameExtractor) Leave(in ast.Node) (ast.Node, bool) {
	if t, ok := in.(*ast.TableName); ok {
		tp := tableNamePair{DBName: t.Schema.L, TableName: t.Name.L}
		if tp.DBName == "" {
			tp.DBName = tne.curDB
		}
		if _, ok := tne.names[tp]; !ok {
			tne.names[tp] = struct{}{}
		}
	}
	return in, true
}

func (sh *sqlInfoFetcher) zipInfoForSQL(w http.ResponseWriter, r *http.Request) {
	var err error
	sh.s, err = session.CreateSession(sh.store)
	if err != nil {
		serveError(w, http.StatusInternalServerError, fmt.Sprintf("create session failed, err: %v", err))
		return
	}
	defer sh.s.Close()
	sh.do = domain.GetDomain(sh.s)
	reqCtx := r.Context()
	sql := r.FormValue("sql")
	pprofTimeString := r.FormValue("pprof_time")
	timeoutString := r.FormValue("timeout")
	curDB := strings.ToLower(r.FormValue("current_db"))
	if curDB != "" {
		_, err = sh.s.Execute(reqCtx, fmt.Sprintf("use %v", curDB))
		if err != nil {
			serveError(w, http.StatusInternalServerError, fmt.Sprintf("use database %v failed, err: %v", curDB, err))
			return
		}
	}
	var (
		pprofTime int
		timeout   int
	)
	if pprofTimeString != "" {
		pprofTime, err = strconv.Atoi(pprofTimeString)
		if err != nil {
			serveError(w, http.StatusBadRequest, "invalid value for pprof_time, please input a int value larger than 5")
			return
		}
	}
	if pprofTimeString != "" && pprofTime < 5 {
		serveError(w, http.StatusBadRequest, "pprof time is too short, please input a int value larger than 5")
	}
	if timeoutString != "" {
		timeout, err = strconv.Atoi(timeoutString)
		if err != nil {
			serveError(w, http.StatusBadRequest, "invalid value for timeout")
			return
		}
	}
	if timeout < pprofTime {
		timeout = pprofTime
	}
	pairs, err := sh.extractTableNames(sql, curDB)
	if err != nil {
		serveError(w, http.StatusBadRequest, fmt.Sprintf("invalid SQL text, err: %v", err))
		return
	}
	zw := zip.NewWriter(w)
	defer func() {
		terror.Log(zw.Close())
	}()
	for pair := range pairs {
		jsonTbl, err := sh.getStatsForTable(pair)
		if err != nil {
			err = sh.writeErrFile(zw, fmt.Sprintf("%v.%v.stats.err.txt", pair.DBName, pair.TableName), err)
			terror.Log(err)
			continue
		}
		statsFw, err := zw.Create(fmt.Sprintf("%v.%v.json", pair.DBName, pair.TableName))
		if err != nil {
			terror.Log(err)
			continue
		}
		data, err := json.Marshal(jsonTbl)
		if err != nil {
			err = sh.writeErrFile(zw, fmt.Sprintf("%v.%v.stats.err.txt", pair.DBName, pair.TableName), err)
			terror.Log(err)
			continue
		}
		_, err = statsFw.Write(data)
		if err != nil {
			err = sh.writeErrFile(zw, fmt.Sprintf("%v.%v.stats.err.txt", pair.DBName, pair.TableName), err)
			terror.Log(err)
			continue
		}
	}
	for pair := range pairs {
		err = sh.getShowCreateTable(pair, zw)
		if err != nil {
			err = sh.writeErrFile(zw, fmt.Sprintf("%v.%v.schema.err.txt", pair.DBName, pair.TableName), err)
			terror.Log(err)
			return
		}
	}
	// If we don't catch profile. We just get a explain result.
	if pprofTime == 0 {
		recordSets, err := sh.s.(sqlexec.SQLExecutor).Execute(reqCtx, fmt.Sprintf("explain %s", sql))
		if len(recordSets) > 0 {
			defer terror.Call(recordSets[0].Close)
		}
		if err != nil {
			err = sh.writeErrFile(zw, "explain.err.txt", err)
			terror.Log(err)
			return
		}
		sRows, err := session.ResultSetToStringSlice(reqCtx, sh.s, recordSets[0])
		if err != nil {
			err = sh.writeErrFile(zw, "explain.err.txt", err)
			terror.Log(err)
			return
		}
		fw, err := zw.Create("explain.txt")
		if err != nil {
			terror.Log(err)
			return
		}
		for _, row := range sRows {
			fmt.Fprintf(fw, "%s\n", strings.Join(row, "\t"))
		}
	} else {
		// Otherwise we catch a profile and run `EXPLAIN ANALYZE` result.
		ctx, cancelFunc := context.WithCancel(reqCtx)
		timer := time.NewTimer(time.Second * time.Duration(timeout))
		resultChan := make(chan *explainAnalyzeResult)
		go sh.getExplainAnalyze(ctx, sql, resultChan)
		errChan := make(chan error)
		var buf bytes.Buffer
		go sh.catchCPUProfile(reqCtx, pprofTime, &buf, errChan)
		select {
		case result := <-resultChan:
			timer.Stop()
			cancelFunc()
			if result.err != nil {
				err = sh.writeErrFile(zw, "explain_analyze.err.txt", result.err)
				terror.Log(err)
				return
			}
			if len(result.rows) == 0 {
				break
			}
			fw, err := zw.Create("explain_analyze.txt")
			if err != nil {
				terror.Log(err)
				break
			}
			for _, row := range result.rows {
				fmt.Fprintf(fw, "%s\n", strings.Join(row, "\t"))
			}
		case <-timer.C:
			cancelFunc()
		}
		err = dumpCPUProfile(errChan, &buf, zw)
		if err != nil {
			err = sh.writeErrFile(zw, "profile.err.txt", err)
			terror.Log(err)
			return
		}
	}
}

func dumpCPUProfile(errChan chan error, buf *bytes.Buffer, zw *zip.Writer) error {
	err := <-errChan
	if err != nil {
		return err
	}
	fw, err := zw.Create("profile")
	if err != nil {
		return err
	}
	_, err = fw.Write(buf.Bytes())
	if err != nil {
		return err
	}
	return nil
}

func (sh *sqlInfoFetcher) writeErrFile(zw *zip.Writer, name string, err error) error {
	fw, err1 := zw.Create(name)
	if err1 != nil {
		return err1
	}
	fmt.Fprintf(fw, "error: %v", err)
	return nil
}

type explainAnalyzeResult struct {
	rows [][]string
	err  error
}

func (sh *sqlInfoFetcher) getExplainAnalyze(ctx context.Context, sql string, resultChan chan<- *explainAnalyzeResult) {
	recordSets, err := sh.s.(sqlexec.SQLExecutor).Execute(ctx, fmt.Sprintf("explain analyze %s", sql))
	if err != nil {
		resultChan <- &explainAnalyzeResult{err: err}
		return
	}
	rows, err := session.ResultSetToStringSlice(ctx, sh.s, recordSets[0])
	if err != nil {
		terror.Log(err)
		return
	}
	if len(recordSets) > 0 {
		terror.Call(recordSets[0].Close)
	}
	resultChan <- &explainAnalyzeResult{rows: rows}
}

func (sh *sqlInfoFetcher) catchCPUProfile(ctx context.Context, sec int, buf *bytes.Buffer, errChan chan<- error) {
	if err := pprof.StartCPUProfile(buf); err != nil {
		errChan <- err
		return
	}
	sleepWithCtx(ctx, time.Duration(sec)*time.Second)
	pprof.StopCPUProfile()
	errChan <- nil
}

func (sh *sqlInfoFetcher) getStatsForTable(pair tableNamePair) (*handle.JSONTable, error) {
	is := sh.do.InfoSchema()
	h := sh.do.StatsHandle()
	tbl, err := is.TableByName(model.NewCIStr(pair.DBName), model.NewCIStr(pair.TableName))
	if err != nil {
		return nil, err
	}
	js, err := h.DumpStatsToJSON(pair.DBName, tbl.Meta(), nil)
	return js, err
}

func (sh *sqlInfoFetcher) getShowCreateTable(pair tableNamePair, zw *zip.Writer) error {
	recordSets, err := sh.s.(sqlexec.SQLExecutor).Execute(context.TODO(), fmt.Sprintf("show create table `%v`.`%v`", pair.DBName, pair.TableName))
	if len(recordSets) > 0 {
		defer terror.Call(recordSets[0].Close)
	}
	if err != nil {
		return err
	}
	sRows, err := session.ResultSetToStringSlice(context.Background(), sh.s, recordSets[0])
	if err != nil {
		terror.Log(err)
		return nil
	}
	fw, err := zw.Create(fmt.Sprintf("%v.%v.schema.txt", pair.DBName, pair.TableName))
	if err != nil {
		terror.Log(err)
		return nil
	}
	for _, row := range sRows {
		fmt.Fprintf(fw, "%s\n", strings.Join(row, "\t"))
	}
	return nil
}

func (sh *sqlInfoFetcher) extractTableNames(sql, curDB string) (map[tableNamePair]struct{}, error) {
	p := parser.New()
	charset, collation := sh.s.GetSessionVars().GetCharsetInfo()
	stmts, _, err := p.Parse(sql, charset, collation)
	if err != nil {
		return nil, err
	}
	if len(stmts) > 1 {
		return nil, errors.Errorf("Only 1 statement is allowed")
	}
	extractor := &tableNameExtractor{
		curDB: curDB,
		names: make(map[tableNamePair]struct{}),
	}
	stmts[0].Accept(extractor)
	return extractor.names, nil
}
<commit_msg>server: use ParseWithParams to avoid injection (#22545)<commit_after>// Copyright 2019 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/parser"
	"github.com/pingcap/parser/ast"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/parser/terror"
	"github.com/pingcap/tidb/domain"
	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/statistics/handle"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/util/sqlexec"
)

type sqlInfoFetcher struct {
	store tikv.Storage
	do    *domain.Domain
	s     session.Session
}

type tableNamePair struct {
	DBName    string
	TableName string
}

type tableNameExtractor struct {
	curDB string
	names map[tableNamePair]struct{}
}

func (tne *tableNameExtractor) Enter(in ast.Node) (ast.Node, bool) {
	if _, ok := in.(*ast.TableName); ok {
		return in, true
	}
	return in, false
}

func (tne *tableNameExtractor) Leave(in ast.Node) (ast.Node, bool) {
	if t, ok := in.(*ast.TableName); ok {
		tp := tableNamePair{DBName: t.Schema.L, TableName: t.Name.L}
		if tp.DBName == "" {
			tp.DBName = tne.curDB
		}
		if _, ok := tne.names[tp]; !ok {
			tne.names[tp] = struct{}{}
		}
	}
	return in, true
}

func (sh *sqlInfoFetcher) zipInfoForSQL(w http.ResponseWriter, r *http.Request) {
	var err error
	sh.s, err = session.CreateSession(sh.store)
	if err != nil {
		serveError(w, http.StatusInternalServerError, fmt.Sprintf("create session failed, err: %v", err))
		return
	}
	defer sh.s.Close()
	sh.do = domain.GetDomain(sh.s)
	reqCtx := r.Context()
	sql := r.FormValue("sql")
	pprofTimeString := r.FormValue("pprof_time")
	timeoutString := r.FormValue("timeout")
	curDB := strings.ToLower(r.FormValue("current_db"))
	if curDB != "" {
		stmts, err := sh.s.ParseWithParams(reqCtx, "use %n", curDB)
		if err != nil {
			serveError(w, http.StatusInternalServerError, fmt.Sprintf("use database %v failed, err: %v", curDB, err))
			return
		}
		_, err = sh.s.ExecuteStmt(reqCtx, stmts[0])
		if err != nil {
			serveError(w, http.StatusInternalServerError, fmt.Sprintf("use database %v failed, err: %v", curDB, err))
			return
		}
	}
	var (
		pprofTime int
		timeout   int
	)
	if pprofTimeString != "" {
		pprofTime, err = strconv.Atoi(pprofTimeString)
		if err != nil {
			serveError(w, http.StatusBadRequest, "invalid value for pprof_time, please input a int value larger than 5")
			return
		}
	}
	if pprofTimeString != "" && pprofTime < 5 {
		serveError(w, http.StatusBadRequest, "pprof time is too short, please input a int value larger than 5")
	}
	if timeoutString != "" {
		timeout, err = strconv.Atoi(timeoutString)
		if err != nil {
			serveError(w, http.StatusBadRequest, "invalid value for timeout")
			return
		}
	}
	if timeout < pprofTime {
		timeout = pprofTime
	}
	pairs, err := sh.extractTableNames(sql, curDB)
	if err != nil {
		serveError(w, http.StatusBadRequest, fmt.Sprintf("invalid SQL text, err: %v", err))
		return
	}
	zw := zip.NewWriter(w)
	defer func() {
		terror.Log(zw.Close())
	}()
	for pair := range pairs {
		jsonTbl, err := sh.getStatsForTable(pair)
		if err != nil {
			err = sh.writeErrFile(zw, fmt.Sprintf("%v.%v.stats.err.txt", pair.DBName, pair.TableName), err)
			terror.Log(err)
			continue
		}
		statsFw, err := zw.Create(fmt.Sprintf("%v.%v.json", pair.DBName, pair.TableName))
		if err != nil {
			terror.Log(err)
			continue
		}
		data, err := json.Marshal(jsonTbl)
		if err != nil {
			err = sh.writeErrFile(zw, fmt.Sprintf("%v.%v.stats.err.txt", pair.DBName, pair.TableName), err)
			terror.Log(err)
			continue
		}
		_, err = statsFw.Write(data)
		if err != nil {
			err = sh.writeErrFile(zw, fmt.Sprintf("%v.%v.stats.err.txt", pair.DBName, pair.TableName), err)
			terror.Log(err)
			continue
		}
	}
	for pair := range pairs {
		err = sh.getShowCreateTable(pair, zw)
		if err != nil {
			err = sh.writeErrFile(zw, fmt.Sprintf("%v.%v.schema.err.txt", pair.DBName, pair.TableName), err)
			terror.Log(err)
			return
		}
	}
	// If we don't catch profile. We just get a explain result.
	if pprofTime == 0 {
		recordSets, err := sh.s.(sqlexec.SQLExecutor).Execute(reqCtx, fmt.Sprintf("explain %s", sql))
		if len(recordSets) > 0 {
			defer terror.Call(recordSets[0].Close)
		}
		if err != nil {
			err = sh.writeErrFile(zw, "explain.err.txt", err)
			terror.Log(err)
			return
		}
		sRows, err := session.ResultSetToStringSlice(reqCtx, sh.s, recordSets[0])
		if err != nil {
			err = sh.writeErrFile(zw, "explain.err.txt", err)
			terror.Log(err)
			return
		}
		fw, err := zw.Create("explain.txt")
		if err != nil {
			terror.Log(err)
			return
		}
		for _, row := range sRows {
			fmt.Fprintf(fw, "%s\n", strings.Join(row, "\t"))
		}
	} else {
		// Otherwise we catch a profile and run `EXPLAIN ANALYZE` result.
		ctx, cancelFunc := context.WithCancel(reqCtx)
		timer := time.NewTimer(time.Second * time.Duration(timeout))
		resultChan := make(chan *explainAnalyzeResult)
		go sh.getExplainAnalyze(ctx, sql, resultChan)
		errChan := make(chan error)
		var buf bytes.Buffer
		go sh.catchCPUProfile(reqCtx, pprofTime, &buf, errChan)
		select {
		case result := <-resultChan:
			timer.Stop()
			cancelFunc()
			if result.err != nil {
				err = sh.writeErrFile(zw, "explain_analyze.err.txt", result.err)
				terror.Log(err)
				return
			}
			if len(result.rows) == 0 {
				break
			}
			fw, err := zw.Create("explain_analyze.txt")
			if err != nil {
				terror.Log(err)
				break
			}
			for _, row := range result.rows {
				fmt.Fprintf(fw, "%s\n", strings.Join(row, "\t"))
			}
		case <-timer.C:
			cancelFunc()
		}
		err = dumpCPUProfile(errChan, &buf, zw)
		if err != nil {
			err = sh.writeErrFile(zw, "profile.err.txt", err)
			terror.Log(err)
			return
		}
	}
}

func dumpCPUProfile(errChan chan error, buf *bytes.Buffer, zw *zip.Writer) error {
	err := <-errChan
	if err != nil {
		return err
	}
	fw, err := zw.Create("profile")
	if err != nil {
		return err
	}
	_, err = fw.Write(buf.Bytes())
	if err != nil {
		return err
	}
	return nil
}

func (sh *sqlInfoFetcher) writeErrFile(zw *zip.Writer, name string, err error) error {
	fw, err1 := zw.Create(name)
	if err1 != nil {
		return err1
	}
	fmt.Fprintf(fw, "error: %v", err)
	return nil
}

type explainAnalyzeResult struct {
	rows [][]string
	err  error
}

func (sh *sqlInfoFetcher) getExplainAnalyze(ctx context.Context, sql string, resultChan chan<- *explainAnalyzeResult) {
	recordSets, err := sh.s.(sqlexec.SQLExecutor).Execute(ctx, fmt.Sprintf("explain analyze %s", sql))
	if err != nil {
		resultChan <- &explainAnalyzeResult{err: err}
		return
	}
	rows, err := session.ResultSetToStringSlice(ctx, sh.s, recordSets[0])
	if err != nil {
		terror.Log(err)
		return
	}
	if len(recordSets) > 0 {
		terror.Call(recordSets[0].Close)
	}
	resultChan <- &explainAnalyzeResult{rows: rows}
}

func (sh *sqlInfoFetcher) catchCPUProfile(ctx context.Context, sec int, buf *bytes.Buffer, errChan chan<- error) {
	if err := pprof.StartCPUProfile(buf); err != nil {
		errChan <- err
		return
	}
	sleepWithCtx(ctx, time.Duration(sec)*time.Second)
	pprof.StopCPUProfile()
	errChan <- nil
}

func (sh *sqlInfoFetcher) getStatsForTable(pair tableNamePair) (*handle.JSONTable, error) {
	is := sh.do.InfoSchema()
	h := sh.do.StatsHandle()
	tbl, err := is.TableByName(model.NewCIStr(pair.DBName), model.NewCIStr(pair.TableName))
	if err != nil {
		return nil, err
	}
	js, err := h.DumpStatsToJSON(pair.DBName, tbl.Meta(), nil)
	return js, err
}

func (sh *sqlInfoFetcher) getShowCreateTable(pair tableNamePair, zw *zip.Writer) error {
	recordSets, err := sh.s.(sqlexec.SQLExecutor).Execute(context.TODO(), fmt.Sprintf("show create table `%v`.`%v`", pair.DBName, pair.TableName))
	if len(recordSets) > 0 {
		defer terror.Call(recordSets[0].Close)
	}
	if err != nil {
		return err
	}
	sRows, err := session.ResultSetToStringSlice(context.Background(), sh.s, recordSets[0])
	if err != nil {
		terror.Log(err)
		return nil
	}
	fw, err := zw.Create(fmt.Sprintf("%v.%v.schema.txt", pair.DBName, pair.TableName))
	if err != nil {
		terror.Log(err)
		return nil
	}
	for _, row := range sRows {
		fmt.Fprintf(fw, "%s\n", strings.Join(row, "\t"))
	}
	return nil
}

func (sh *sqlInfoFetcher) extractTableNames(sql, curDB string) (map[tableNamePair]struct{}, error) {
	p := parser.New()
	charset, collation := sh.s.GetSessionVars().GetCharsetInfo()
	stmts, _, err := p.Parse(sql, charset, collation)
	if err != nil {
		return nil, err
	}
	if len(stmts) > 1 {
		return nil, errors.Errorf("Only 1 statement is allowed")
	}
	extractor := &tableNameExtractor{
		curDB: curDB,
		names: make(map[tableNamePair]struct{}),
	}
	stmts[0].Accept(extractor)
	return extractor.names, nil
}
<|endoftext|>


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

package server

import (
	"strconv"

	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/pd/server/core"
	"github.com/pingcap/pd/server/namespace"
)

type storeStatistics struct {
	opt             *scheduleOption
	namespace       string
	Up              int
	Disconnect      int
	Unhealth        int
	Down            int
	Offline         int
	Tombstone       int
	LowSpace        int
	StorageSize     uint64
	StorageCapacity uint64
	RegionCount     int
	LeaderCount     int
}

func newStoreStatistics(opt *scheduleOption, namespace string) *storeStatistics {
	return &storeStatistics{
		opt:       opt,
		namespace: namespace,
	}
}

func (s *storeStatistics) Observe(store *core.StoreInfo) {
	// Store state.
	switch store.GetState() {
	case metapb.StoreState_Up:
		if store.DownTime() >= s.opt.GetMaxStoreDownTime() {
			s.Down++
		} else if store.IsUnhealth() {
			s.Unhealth++
		} else if store.IsDisconnected() {
			s.Disconnect++
		} else {
			s.Up++
		}
	case metapb.StoreState_Offline:
		s.Offline++
	case metapb.StoreState_Tombstone:
		s.Tombstone++
		return
	}
	if store.IsLowSpace(s.opt.GetLowSpaceRatio()) {
		s.LowSpace++
	}

	// Store stats.
	s.StorageSize += store.StorageSize()
	s.StorageCapacity += store.Stats.GetCapacity()
	s.RegionCount += store.RegionCount
	s.LeaderCount += store.LeaderCount

	id := strconv.FormatUint(store.GetId(), 10)
	storeStatusGauge.WithLabelValues(s.namespace, id, "region_score").Set(store.RegionScore(s.opt.GetHighSpaceRatio(), s.opt.GetLowSpaceRatio(), 0))
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_score").Set(store.LeaderScore(0))
	storeStatusGauge.WithLabelValues(s.namespace, id, "region_size").Set(float64(store.RegionSize))
	storeStatusGauge.WithLabelValues(s.namespace, id, "region_count").Set(float64(store.RegionCount))
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_size").Set(float64(store.LeaderSize))
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_count").Set(float64(store.LeaderCount))
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_available").Set(float64(store.Stats.GetAvailable()))
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_used").Set(float64(store.Stats.GetUsedSize()))
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_capacity").Set(float64(store.Stats.GetCapacity()))
}

func (s *storeStatistics) Collect() {
	metrics := make(map[string]float64)
	metrics["store_up_count"] = float64(s.Up)
	metrics["store_disconnected_count"] = float64(s.Disconnect)
	metrics["store_down_count"] = float64(s.Down)
	metrics["store_unhealth_count"] = float64(s.Unhealth)
	metrics["store_offline_count"] = float64(s.Offline)
	metrics["store_tombstone_count"] = float64(s.Tombstone)
	metrics["store_low_space_count"] = float64(s.LowSpace)
	metrics["region_count"] = float64(s.RegionCount)
	metrics["leader_count"] = float64(s.LeaderCount)
	metrics["storage_size"] = float64(s.StorageSize)
	metrics["storage_capacity"] = float64(s.StorageCapacity)

	for label, value := range metrics {
		clusterStatusGauge.WithLabelValues(label, s.namespace).Set(value)
	}
}

type storeStatisticsMap struct {
	opt        *scheduleOption
	classifier namespace.Classifier
	stats      map[string]*storeStatistics
}

func newStoreStatisticsMap(opt *scheduleOption, classifier namespace.Classifier) *storeStatisticsMap {
	return &storeStatisticsMap{
		opt:        opt,
		classifier: classifier,
		stats:      make(map[string]*storeStatistics),
	}
}

func (m *storeStatisticsMap) Observe(store *core.StoreInfo) {
	namespace := m.classifier.GetStoreNamespace(store)
	stat, ok := m.stats[namespace]
	if !ok {
		stat = newStoreStatistics(m.opt, namespace)
		m.stats[namespace] = stat
	}
	stat.Observe(store)
}

func (m *storeStatisticsMap) Collect() {
	for _, s := range m.stats {
		s.Collect()
	}
}
<commit_msg>fix store metrics (#1261)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"strconv"

	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/pd/server/core"
	"github.com/pingcap/pd/server/namespace"
)

type storeStatistics struct {
	opt             *scheduleOption
	namespace       string
	Up              int
	Disconnect      int
	Unhealth        int
	Down            int
	Offline         int
	Tombstone       int
	LowSpace        int
	StorageSize     uint64
	StorageCapacity uint64
	RegionCount     int
	LeaderCount     int
}

func newStoreStatistics(opt *scheduleOption, namespace string) *storeStatistics {
	return &storeStatistics{
		opt:       opt,
		namespace: namespace,
	}
}

func (s *storeStatistics) Observe(store *core.StoreInfo) {
	id := strconv.FormatUint(store.GetId(), 10)
	// Store state.
	switch store.GetState() {
	case metapb.StoreState_Up:
		if store.DownTime() >= s.opt.GetMaxStoreDownTime() {
			s.Down++
		} else if store.IsUnhealth() {
			s.Unhealth++
		} else if store.IsDisconnected() {
			s.Disconnect++
		} else {
			s.Up++
		}
	case metapb.StoreState_Offline:
		s.Offline++
	case metapb.StoreState_Tombstone:
		s.Tombstone++
		s.resetStoreStatistics(id)
		return
	}
	if store.IsLowSpace(s.opt.GetLowSpaceRatio()) {
		s.LowSpace++
	}

	// Store stats.
	s.StorageSize += store.StorageSize()
	s.StorageCapacity += store.Stats.GetCapacity()
	s.RegionCount += store.RegionCount
	s.LeaderCount += store.LeaderCount

	storeStatusGauge.WithLabelValues(s.namespace, id, "region_score").Set(store.RegionScore(s.opt.GetHighSpaceRatio(), s.opt.GetLowSpaceRatio(), 0))
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_score").Set(store.LeaderScore(0))
	storeStatusGauge.WithLabelValues(s.namespace, id, "region_size").Set(float64(store.RegionSize))
	storeStatusGauge.WithLabelValues(s.namespace, id, "region_count").Set(float64(store.RegionCount))
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_size").Set(float64(store.LeaderSize))
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_count").Set(float64(store.LeaderCount))
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_available").Set(float64(store.Stats.GetAvailable()))
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_used").Set(float64(store.Stats.GetUsedSize()))
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_capacity").Set(float64(store.Stats.GetCapacity()))
}

func (s *storeStatistics) Collect() {
	metrics := make(map[string]float64)
	metrics["store_up_count"] = float64(s.Up)
	metrics["store_disconnected_count"] = float64(s.Disconnect)
	metrics["store_down_count"] = float64(s.Down)
	metrics["store_unhealth_count"] = float64(s.Unhealth)
	metrics["store_offline_count"] = float64(s.Offline)
	metrics["store_tombstone_count"] = float64(s.Tombstone)
	metrics["store_low_space_count"] = float64(s.LowSpace)
	metrics["region_count"] = float64(s.RegionCount)
	metrics["leader_count"] = float64(s.LeaderCount)
	metrics["storage_size"] = float64(s.StorageSize)
	metrics["storage_capacity"] = float64(s.StorageCapacity)

	for label, value := range metrics {
		clusterStatusGauge.WithLabelValues(label, s.namespace).Set(value)
	}
}

func (s *storeStatistics) resetStoreStatistics(id string) {
	storeStatusGauge.WithLabelValues(s.namespace, id, "region_score").Set(0)
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_score").Set(0)
	storeStatusGauge.WithLabelValues(s.namespace, id, "region_size").Set(0)
	storeStatusGauge.WithLabelValues(s.namespace, id, "region_count").Set(0)
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_size").Set(0)
	storeStatusGauge.WithLabelValues(s.namespace, id, "leader_count").Set(0)
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_available").Set(0)
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_used").Set(0)
	storeStatusGauge.WithLabelValues(s.namespace, id, "store_capacity").Set(0)
}

type storeStatisticsMap struct {
	opt        *scheduleOption
	classifier namespace.Classifier
	stats      map[string]*storeStatistics
}

func newStoreStatisticsMap(opt *scheduleOption, classifier namespace.Classifier) *storeStatisticsMap {
	return &storeStatisticsMap{
		opt:        opt,
		classifier: classifier,
		stats:      make(map[string]*storeStatistics),
	}
}

func (m *storeStatisticsMap) Observe(store *core.StoreInfo) {
	namespace := m.classifier.GetStoreNamespace(store)
	stat, ok := m.stats[namespace]
	if !ok {
		stat = newStoreStatistics(m.opt, namespace)
		m.stats[namespace] = stat
	}
	stat.Observe(store)
}

func (m *storeStatisticsMap) Collect() {
	for _, s := range m.stats {
		s.Collect()
	}
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package health

import (
	"net/http"
	"strconv"
	"testing"
	"time"

	"github.com/issue9/assert/v2"
	"github.com/issue9/cache/memory"
	"github.com/issue9/web"
	"github.com/issue9/web/server"
	"github.com/issue9/web/server/servertest"
)

var _ server.Middleware = &Health{}

func TestHealth(t *testing.T) {
	a := assert.New(t, false)
	s := servertest.NewTester(a, &web.Options{Cache: memory.New(1 * time.Minute), Port: ":8080"})

	h := New(NewCacheStore(s.Server(), "health_"))
	r := s.NewRouter(h)
	r.Get("/", func(ctx *web.Context) web.Responser {
		status, err := strconv.Atoi(ctx.Request().FormValue("status"))
		if err != nil {
			panic(err)
		}
		return ctx.Status(status)
	})
	r.Post("/", func(ctx *web.Context) web.Responser {
		return nil
	})
	r.Delete("/users", func(ctx *web.Context) web.Responser {
		status, err := strconv.Atoi(ctx.Request().FormValue("status"))
		if err != nil {
			panic(err)
		}
		return ctx.Status(status)
	})

	s.GoServe()

	mem := h.store
	state := mem.Get(http.MethodGet, "/")
	a.Equal(0, state.Count)

	// 第一次访问 GET /
	s.Get("/").Query("status", "200").Do(nil).Status(200)
	time.Sleep(time.Microsecond * 500)
	state = mem.Get(http.MethodGet, "/")
	a.Equal(1, state.Count).True(state.Min > 0)

	// 第二次访问 GET /
	s.Get("/").Query("status", "500").Do(nil)
	time.Sleep(time.Microsecond * 500)
	state = mem.Get(http.MethodGet, "/")
	a.Equal(2, state.Count).
		Equal(1, state.ServerErrors).
		Equal(0, state.UserErrors)

	// 第一次访问 POST /
	s.NewRequest(http.MethodPost, "/", nil).Query("status", "201").Do(nil)
	time.Sleep(time.Microsecond * 500)
	state = mem.Get(http.MethodPost, "/")
	a.Equal(1, state.Count)

	// 第一次访问 DELETE /users
	s.Delete("/users").Query("status", "401").Do(nil)
	time.Sleep(time.Microsecond * 500)
	state = mem.Get(http.MethodDelete, "/users")
	a.Equal(1, state.Count).
		Equal(0, state.ServerErrors).
		Equal(1, state.UserErrors)

	all := h.States()
	a.Equal(3, len(all))

	s.Close(0)
	s.Wait()
}
<commit_msg>test(health): 修正测试用例无法记录用时的错误<commit_after>// SPDX-License-Identifier: MIT

package health

import (
	"math/rand"
	"net/http"
	"strconv"
	"testing"
	"time"

	"github.com/issue9/assert/v2"
	"github.com/issue9/cache/memory"
	"github.com/issue9/web"
	"github.com/issue9/web/server"
	"github.com/issue9/web/server/servertest"
)

var _ server.Middleware = &Health{}

func TestHealth(t *testing.T) {
	a := assert.New(t, false)
	s := servertest.NewTester(a, &web.Options{Cache: memory.New(1 * time.Minute), Port: ":8080"})

	h := New(NewCacheStore(s.Server(), "health_"))
	r := s.NewRouter(h)
	r.Get("/", func(ctx *web.Context) web.Responser {
		status, err := strconv.Atoi(ctx.Request().FormValue("status"))
		if err != nil {
			panic(err)
		}
		time.Sleep(time.Microsecond * time.Duration(rand.Int63n(100))) // 防止过快，无法记录用时。
		return ctx.Status(status)
	})
	r.Post("/", func(ctx *web.Context) web.Responser {
		time.Sleep(time.Microsecond * time.Duration(rand.Int63n(100))) // 防止过快，无法记录用时。
		return nil
	})
	r.Delete("/users", func(ctx *web.Context) web.Responser {
		status, err := strconv.Atoi(ctx.Request().FormValue("status"))
		if err != nil {
			panic(err)
		}
		time.Sleep(time.Microsecond * time.Duration(rand.Int63n(100))) // 防止过快，无法记录用时。
		return ctx.Status(status)
	})

	s.GoServe()

	mem := h.store
	state := mem.Get(http.MethodGet, "/")
	a.Equal(0, state.Count)

	// 第一次访问 GET /
	s.Get("/").Query("status", "200").Do(nil).Status(200)
	time.Sleep(time.Microsecond * 500)
	state = mem.Get(http.MethodGet, "/")
	a.Equal(1, state.Count).True(state.Min > 0)

	// 第二次访问 GET /
	s.Get("/").Query("status", "500").Do(nil)
	time.Sleep(time.Microsecond * 500)
	state = mem.Get(http.MethodGet, "/")
	a.Equal(2, state.Count).
		Equal(1, state.ServerErrors).
		Equal(0, state.UserErrors)

	// 第一次访问 POST /
	s.NewRequest(http.MethodPost, "/", nil).Query("status", "201").Do(nil)
	time.Sleep(time.Microsecond * 500)
	state = mem.Get(http.MethodPost, "/")
	a.Equal(1, state.Count)

	// 第一次访问 DELETE /users
	s.Delete("/users").Query("status", "401").Do(nil)
	time.Sleep(time.Microsecond * 500)
	state = mem.Get(http.MethodDelete, "/users")
	a.Equal(1, state.Count).
		Equal(0, state.ServerErrors).
		Equal(1, state.UserErrors)

	all := h.States()
	a.Equal(3, len(all))

	s.Close(0)
	s.Wait()
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2013 Gijs Kunze
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
//

package policy

import (
	"fmt"
)

func ExampleNewPolicy() {
	p := NewPolicy()
	p.SetId("policy-id")
	stmt := p.AddStatement()
	stmt.SetSid("statement-id")
	stmt.Effect = Allow
	stmt.AddPrincipal("*")
	stmt.AddAction("Describe*")
	stmt.Resource = "*"
	stmt.AddCondition(ConditionArnEquals, VarSourceIp, "10.0.0.0/8")

	fmt.Println(p)
	// Output:
	// {
	//     "Version": "2012-10-17",
	//     "Id": "policy-id",
	//     "Statement": [
	//         {
	//             "Sid": "statement-id",
	//             "Effect": "Allow",
	//             "Principal": {
	//                 "AWS": "*"
	//             },
	//             "Action": [
	//                 "Describe*"
	//             ],
	//             "Resource": "*",
	//             "Condition": {
	//                 "ArnEquals": {
	//                     "aws:SourceIp": [
	//                         "10.0.0.0/8"
	//                     ]
	//                 }
	//             }
	//         }
	//     ]
	// }
}
<commit_msg>Fixed example output<commit_after>//
// Copyright (c) 2013 Gijs Kunze
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
//

package policy

import (
	"fmt"
)

func ExampleNewPolicy() {
	p := NewPolicy()
	p.SetId("policy-id")
	stmt := p.AddStatement()
	stmt.SetSid("statement-id")
	stmt.Effect = Allow
	stmt.AddPrincipal("*")
	stmt.AddAction("Describe*")
	stmt.Resource = "*"
	stmt.AddCondition(ConditionArnEquals, VarSourceIp, "10.0.0.0/8")

	fmt.Println(p)
	// Output:
	// {
	//     "Version": "2012-10-17",
	//     "Id": "policy-id",
	//     "Statement": [
	//         {
	//             "Sid": "statement-id",
	//             "Effect": "Allow",
	//             "Principal": {
	//                 "AWS": [
	//                     "*"
	//                 ]
	//             },
	//             "Action": [
	//                 "Describe*"
	//             ],
	//             "Resource": "*",
	//             "Condition": {
	//                 "ArnEquals": {
	//                     "aws:SourceIp": [
	//                         "10.0.0.0/8"
	//                     ]
	//                 }
	//             }
	//         }
	//     ]
	// }
}
<|endoftext|>


<commit_before>package progress

import (
	"fmt"
	"io"
	"os"
	"time"
)

// Bar - a simple, unobtrusive progress indicator
// Usage:
//   bar := New()
//   bar.SetTotal(3)
//   defer bar.Done()
//   bar.Increment("Step 1")
//   bar.Increment("Step 2")
//   bar.Increment("Step 3")
//   bar.Complete("Done")
//
// Instantiation creates a progress bar consisiting of an optional spinner
// prefix followed by an indicator of the current step, the total steps, and a
// trailing message The behavior differs substantially if verbosity is enbled
// or not. When in verbose mode, it is expected that the process is otherwise
// printing status information to output, in which case this bar only prints a
// single status update to output when the current step is updated,
// interspersing the status line with standard program output.  When verbosity
// is not enabled (the default), it is expected that in general output is left
// clean. In this case the status is continually written to standard output
// along with an animated spinner. The status is written one line above current,
// such that any errors from the rest of the process begin at the beinning of a
// line.  Finally, the bar respects headless mode, omitting any printing unless
// explicitly configured to do so, and even then not doing the continual write
// method but the single status update in the same manner as when verbosity is
// enabled.
// Format:
//   [spinner] i/n t
type Bar struct {
	out   io.Writer
	index int    // Current step index
	total int    // Total steps
	text  string // Current display text

	// verbose mode disables progress spinner and line overwrites, instead
	// printing single, full line updates.
	Verbose bool

	// print verbose-style updates even when not attached to an interactive terminal.
	printWhileHeadless bool

	// print N/M step counter with messages
	printWithStepCounter bool

	// Ticker for animated progress when non-verbose, interactive terminal.
	ticker *time.Ticker
}

type Option func(*Bar)

func WithOutput(w io.Writer) Option {
	return func(b *Bar) {
		b.out = w
	}
}

// WithPrintHeadless allows for overriding the default behavior of
// squelching all output when the terminal is detected as headless
// (noninteractive)
func WithPrintWhileHeadless(p bool) Option {
	return func(b *Bar) {
		b.printWhileHeadless = p
	}
}

func WithPrintStepCounter(s bool) Option {
	return func(b *Bar) {
		b.printWithStepCounter = s
	}
}

func New(options ...Option) *Bar {
	b := &Bar{
		out: os.Stdout,
	}
	for _, o := range options {
		o(b)
	}
	return b
}

// SetTotal number of steps.
func (b *Bar) SetTotal(n int) {
	b.total = n
}

// Increment the currenly active step, including beginning printing on first call.
func (b *Bar) Increment(text string) {
	if b.index < b.total {
		b.index++
	}
	b.text = text

	// If this is not an interactive terminal, only print if explicitly set to
	// print while headless, and even then, only a simple line write.
	if !interactiveTerminal() {
		if b.printWhileHeadless {
			b.write()
		}
		return
	}

	// If we're in verbose mode, do a simple write
	if b.Verbose {
		b.write()
		return
	}

	// Start the spinner if not already started
	if b.ticker == nil {
		fmt.Println()
		b.ticker = time.NewTicker(100 * time.Millisecond)
		go b.spin(b.ticker.C)
	}

	// Otherwise we are in non-verbose, interactive mode.  Do a line-overwrite.
	b.overwrite("   ") // Write with space for the spinner
}

// Complete the spinner by advancing to the last step, printing the final text and stopping the write loop.
func (b *Bar) Complete(text string) {
	if !interactiveTerminal() && !b.printWhileHeadless {
		return
	}
	b.index = b.total // Skip to last step
	b.text = text

	// If this is not an interactive terminal, only print if explicitly set to
	// print while headless, and even then, only a simple line write.
	if !interactiveTerminal() {
		if b.printWhileHeadless {
			b.write()
		}
		return
	}

	// If we're interactive, but in verbose mode do a simple write
	if b.Verbose {
		b.write()
		return
	}

	// Otherwise we are in non-verbose mode with an interactive terminal.
	// We should stop the spinner and write an unindented line.
	b.Done() // stop spinner
}

// Stopping indicates the process is stopping, such as having received a context
// cancellation.
func (b *Bar) Stopping() {
	// currently stopping is equivalent in effect to Done
	b.Done()
}

// Done cancels the write loop if being used.
// Call in a defer statement after creation to ensure that the spinner stops
func (b *Bar) Done() {
	if b.ticker != nil {
		b.ticker.Stop()
		b.ticker = nil
	}
}

// Write a simple line status update.
func (b *Bar) write() {
	fmt.Fprintln(b.out, b)
}

// interactiveTerminal returns whether or not the currently attached process
// terminal is interactive.  Used for determining whether or not to
// interactively prompt the user to confirm default choices, etc.
func interactiveTerminal() bool {
	fi, err := os.Stdin.Stat()
	return err == nil && ((fi.Mode() & os.ModeCharDevice) != 0)
}

const (
	up    = "\033[1A"
	down  = "\033[1B"
	clear = "\033[K"
)

// overwrite the line prior with the bar text, optional prefix, creating an empty
// current line for potential errors/warnings.
func (b *Bar) overwrite(prefix string) {
	//  1 Move to the front of the current line
	//  2 Move up one line
	//  3 Clear to the end of the line
	//  4 Print status text with optional prefix (spinner)
	//  5 Print linebreak such that subsequent messages print correctly.
	fmt.Fprintf(b.out, "\r%v%v%v%v\n", up, clear, prefix, b)
}

func (b *Bar) String() string {
	if b.printWithStepCounter {
		return fmt.Sprintf("%v/%v %v", b.index, b.total, b.text)
	}
	return b.text
}

// Write a spinner at the beginning of the previous line.
func (b *Bar) spin(ch <-chan time.Time) {
	if b.Verbose {
		return
	}
	// Various options for spinners.
	// spinner := []string{"|", "/", "-", "\\"}
	// spinner := []string{"◢", "◣", "◤", "◥"}
	spinner := []string{
		"🕛 ",
		"🕐 ",
		"🕑 ",
		"🕒 ",
		"🕓 ",
		"🕔 ",
		"🕕 ",
		"🕖 ",
		"🕗 ",
		"🕘 ",
		"🕙 ",
		"🕚 ",
	}
	idx := 0
	for range ch {
		// Writes the spinner frame at the beginning of the previous line, moving
		// the cursor back to the beginning of the current line for any errors or
		// informative messages.
		fmt.Fprintf(b.out, "\r%v%v%v\r", up, spinner[idx], down)
		idx = (idx + 1) % len(spinner)
	}
}
<commit_msg>fix: use ascii chars in progress indicator on win (#459)<commit_after>package progress

import (
	"fmt"
	"io"
	"os"
	"runtime"
	"time"
)

// Bar - a simple, unobtrusive progress indicator
// Usage:
//   bar := New()
//   bar.SetTotal(3)
//   defer bar.Done()
//   bar.Increment("Step 1")
//   bar.Increment("Step 2")
//   bar.Increment("Step 3")
//   bar.Complete("Done")
//
// Instantiation creates a progress bar consisiting of an optional spinner
// prefix followed by an indicator of the current step, the total steps, and a
// trailing message The behavior differs substantially if verbosity is enbled
// or not. When in verbose mode, it is expected that the process is otherwise
// printing status information to output, in which case this bar only prints a
// single status update to output when the current step is updated,
// interspersing the status line with standard program output.  When verbosity
// is not enabled (the default), it is expected that in general output is left
// clean. In this case the status is continually written to standard output
// along with an animated spinner. The status is written one line above current,
// such that any errors from the rest of the process begin at the beinning of a
// line.  Finally, the bar respects headless mode, omitting any printing unless
// explicitly configured to do so, and even then not doing the continual write
// method but the single status update in the same manner as when verbosity is
// enabled.
// Format:
//   [spinner] i/n t
type Bar struct {
	out   io.Writer
	index int    // Current step index
	total int    // Total steps
	text  string // Current display text

	// verbose mode disables progress spinner and line overwrites, instead
	// printing single, full line updates.
	Verbose bool

	// print verbose-style updates even when not attached to an interactive terminal.
	printWhileHeadless bool

	// print N/M step counter with messages
	printWithStepCounter bool

	// Ticker for animated progress when non-verbose, interactive terminal.
	ticker *time.Ticker
}

type Option func(*Bar)

func WithOutput(w io.Writer) Option {
	return func(b *Bar) {
		b.out = w
	}
}

// WithPrintHeadless allows for overriding the default behavior of
// squelching all output when the terminal is detected as headless
// (noninteractive)
func WithPrintWhileHeadless(p bool) Option {
	return func(b *Bar) {
		b.printWhileHeadless = p
	}
}

func WithPrintStepCounter(s bool) Option {
	return func(b *Bar) {
		b.printWithStepCounter = s
	}
}

func New(options ...Option) *Bar {
	b := &Bar{
		out: os.Stdout,
	}
	for _, o := range options {
		o(b)
	}
	return b
}

// SetTotal number of steps.
func (b *Bar) SetTotal(n int) {
	b.total = n
}

// Increment the currenly active step, including beginning printing on first call.
func (b *Bar) Increment(text string) {
	if b.index < b.total {
		b.index++
	}
	b.text = text

	// If this is not an interactive terminal, only print if explicitly set to
	// print while headless, and even then, only a simple line write.
	if !interactiveTerminal() {
		if b.printWhileHeadless {
			b.write()
		}
		return
	}

	// If we're in verbose mode, do a simple write
	if b.Verbose {
		b.write()
		return
	}

	// Start the spinner if not already started
	if b.ticker == nil {
		fmt.Println()
		b.ticker = time.NewTicker(100 * time.Millisecond)
		go b.spin(b.ticker.C)
	}

	// Otherwise we are in non-verbose, interactive mode.  Do a line-overwrite.
	b.overwrite("   ") // Write with space for the spinner
}

// Complete the spinner by advancing to the last step, printing the final text and stopping the write loop.
func (b *Bar) Complete(text string) {
	if !interactiveTerminal() && !b.printWhileHeadless {
		return
	}
	b.index = b.total // Skip to last step
	b.text = text

	// If this is not an interactive terminal, only print if explicitly set to
	// print while headless, and even then, only a simple line write.
	if !interactiveTerminal() {
		if b.printWhileHeadless {
			b.write()
		}
		return
	}

	// If we're interactive, but in verbose mode do a simple write
	if b.Verbose {
		b.write()
		return
	}

	// Otherwise we are in non-verbose mode with an interactive terminal.
	// We should stop the spinner and write an unindented line.
	b.Done() // stop spinner
}

// Stopping indicates the process is stopping, such as having received a context
// cancellation.
func (b *Bar) Stopping() {
	// currently stopping is equivalent in effect to Done
	b.Done()
}

// Done cancels the write loop if being used.
// Call in a defer statement after creation to ensure that the spinner stops
func (b *Bar) Done() {
	if b.ticker != nil {
		b.ticker.Stop()
		b.ticker = nil
	}
}

// Write a simple line status update.
func (b *Bar) write() {
	fmt.Fprintln(b.out, b)
}

// interactiveTerminal returns whether or not the currently attached process
// terminal is interactive.  Used for determining whether or not to
// interactively prompt the user to confirm default choices, etc.
func interactiveTerminal() bool {
	fi, err := os.Stdin.Stat()
	return err == nil && ((fi.Mode() & os.ModeCharDevice) != 0)
}

const (
	up    = "\033[1A"
	down  = "\033[1B"
	clear = "\033[K"
)

// overwrite the line prior with the bar text, optional prefix, creating an empty
// current line for potential errors/warnings.
func (b *Bar) overwrite(prefix string) {
	//  1 Move to the front of the current line
	//  2 Move up one line
	//  3 Clear to the end of the line
	//  4 Print status text with optional prefix (spinner)
	//  5 Print linebreak such that subsequent messages print correctly.
	fmt.Fprintf(b.out, "\r%v%v%v%v\n", up, clear, prefix, b)
}

func (b *Bar) String() string {
	if b.printWithStepCounter {
		return fmt.Sprintf("%v/%v %v", b.index, b.total, b.text)
	}
	return b.text
}

// Write a spinner at the beginning of the previous line.
func (b *Bar) spin(ch <-chan time.Time) {
	if b.Verbose {
		return
	}
	// Various options for spinners.
	// spinner := []string{"|", "/", "-", "\\"}
	// spinner := []string{"◢", "◣", "◤", "◥"}
	spinner := []string{
		"🕛 ",
		"🕐 ",
		"🕑 ",
		"🕒 ",
		"🕓 ",
		"🕔 ",
		"🕕 ",
		"🕖 ",
		"🕗 ",
		"🕘 ",
		"🕙 ",
		"🕚 ",
	}
	if runtime.GOOS == "windows" {
		spinner = []string{"|", "/", "-", "\\"}
	}
	idx := 0
	for range ch {
		// Writes the spinner frame at the beginning of the previous line, moving
		// the cursor back to the beginning of the current line for any errors or
		// informative messages.
		fmt.Fprintf(b.out, "\r%v%v%v\r", up, spinner[idx], down)
		idx = (idx + 1) % len(spinner)
	}
}
<|endoftext|>


<commit_before>package resource

import (
	"fmt"
	"strings"
	"sync"
	"time"
)

const UniqueIdPrefix = `terraform-`

// idCounter is a monotonic counter for generating ordered unique ids.
var idMutex sync.Mutex
var idCounter uint32

// Helper for a resource to generate a unique identifier w/ default prefix
func UniqueId() string {
	return PrefixedUniqueId(UniqueIdPrefix)
}

// Helper for a resource to generate a unique identifier w/ given prefix
//
// After the prefix, the ID consists of an incrementing 26 digit value (to match
// previous timestamp output).  After the prefix, the ID consists of a timestamp
// and an incrementing 8 hex digit value The timestamp means that multiple IDs
// created with the same prefix will sort in the order of their creation, even
// across multiple terraform executions, as long as the clock is not turned back
// between calls, and as long as any given terraform execution generates fewer
// than 4 billion IDs.
func PrefixedUniqueId(prefix string) string {
	// Be precise to 4 digits of fractional seconds, but remove the dot before the
	// fractional seconds.
	timestamp := strings.Replace(
		time.Now().UTC().Format("20060102150405.0000"), ".", "", 1)

	idMutex.Lock()
	defer idMutex.Unlock()
	idCounter++
	return fmt.Sprintf("%s%s%08x", prefix, timestamp, idCounter)
}
<commit_msg>Export a const for validation methods<commit_after>package resource

import (
	"fmt"
	"strings"
	"sync"
	"time"
)

const UniqueIdPrefix = `terraform-`

// idCounter is a monotonic counter for generating ordered unique ids.
var idMutex sync.Mutex
var idCounter uint32

// Helper for a resource to generate a unique identifier w/ default prefix
func UniqueId() string {
	return PrefixedUniqueId(UniqueIdPrefix)
}

// UniqueIDSuffixLength is the string length of the suffix generated by
// PrefixedUniqueId. This can be used by length validation functions to
// ensure prefixes are the correct length for the target field.
const UniqueIDSuffixLength = 26

// Helper for a resource to generate a unique identifier w/ given prefix
//
// After the prefix, the ID consists of an incrementing 26 digit value (to match
// previous timestamp output).  After the prefix, the ID consists of a timestamp
// and an incrementing 8 hex digit value The timestamp means that multiple IDs
// created with the same prefix will sort in the order of their creation, even
// across multiple terraform executions, as long as the clock is not turned back
// between calls, and as long as any given terraform execution generates fewer
// than 4 billion IDs.
func PrefixedUniqueId(prefix string) string {
	// Be precise to 4 digits of fractional seconds, but remove the dot before the
	// fractional seconds.
	timestamp := strings.Replace(
		time.Now().UTC().Format("20060102150405.0000"), ".", "", 1)

	idMutex.Lock()
	defer idMutex.Unlock()
	idCounter++
	return fmt.Sprintf("%s%s%08x", prefix, timestamp, idCounter)
}
<|endoftext|>


<commit_before>package osxkeychain

/*
#cgo CFLAGS: -x objective-c
#cgo LDFLAGS: -framework Security -framework Foundation

#include "osxkeychain_darwin.h"
#include <stdlib.h>
*/
import "C"
import (
	"errors"
	"github.com/docker/docker-credential-helpers/credentials"
	"net/url"
	"strconv"
	"strings"
	"unsafe"
)

// errCredentialsNotFound is the specific error message returned by OS X
// when the credentials are not in the keychain.
const errCredentialsNotFound = "The specified item could not be found in the keychain."

// Osxkeychain handles secrets using the OS X Keychain as store.
type Osxkeychain struct{}

// Add adds new credentials to the keychain.
func (h Osxkeychain) Add(creds *credentials.Credentials) error {
	s, err := splitServer(creds.ServerURL)
	if err != nil {
		return err
	}
	defer freeServer(s)

	username := C.CString(creds.Username)
	defer C.free(unsafe.Pointer(username))
	secret := C.CString(creds.Secret)
	defer C.free(unsafe.Pointer(secret))

	errMsg := C.keychain_add(s, username, secret)
	if errMsg != nil {
		defer C.free(unsafe.Pointer(errMsg))
		return errors.New(C.GoString(errMsg))
	}

	return nil
}

// Delete removes credentials from the keychain.
func (h Osxkeychain) Delete(serverURL string) error {
	s, err := splitServer(serverURL)
	if err != nil {
		return err
	}
	defer freeServer(s)

	errMsg := C.keychain_delete(s)
	if errMsg != nil {
		defer C.free(unsafe.Pointer(errMsg))
		return errors.New(C.GoString(errMsg))
	}

	return nil
}

// Get returns the username and secret to use for a given registry server URL.
func (h Osxkeychain) Get(serverURL string) (string, string, error) {
	s, err := splitServer(serverURL)
	if err != nil {
		return "", "", err
	}
	defer freeServer(s)

	var usernameLen C.uint
	var username *C.char
	var secretLen C.uint
	var secret *C.char
	defer C.free(unsafe.Pointer(username))
	defer C.free(unsafe.Pointer(secret))

	errMsg := C.keychain_get(s, &usernameLen, &username, &secretLen, &secret)
	if errMsg != nil {
		defer C.free(unsafe.Pointer(errMsg))
		goMsg := C.GoString(errMsg)
		if goMsg == errCredentialsNotFound {
			return "", "", credentials.NewErrCredentialsNotFound()
		}

		return "", "", errors.New(goMsg)
	}

	user := C.GoStringN(username, C.int(usernameLen))
	pass := C.GoStringN(secret, C.int(secretLen))
	return user, pass, nil
}

// List returns the stored URLs and corresponding usernames.
func (h Osxkeychain) List() (map[string]string, error) {
	var pathsC **C.char
	defer C.free(unsafe.Pointer(pathsC))
	var acctsC **C.char
	defer C.free(unsafe.Pointer(acctsC))
	var listLenC C.uint
	errMsg := C.keychain_list(&pathsC, &acctsC, &listLenC)
	if errMsg != nil {
		defer C.free(unsafe.Pointer(errMsg))
		goMsg := C.GoString(errMsg)
		return nil, errors.New(goMsg)
	}

	defer C.freeListData(&pathsC, listLenC)
	defer C.freeListData(&acctsC, listLenC)

	var listLen int
	listLen = int(listLenC)
	pathTmp := (*[1 << 30]*C.char)(unsafe.Pointer(pathsC))[:listLen:listLen]
	acctTmp := (*[1 << 30]*C.char)(unsafe.Pointer(acctsC))[:listLen:listLen]
	//taking the array of c strings into go while ignoring all the stuff irrelevant to credentials-helper
	resp := make(map[string]string)
	for i := 0; i < listLen; i++ {
		if C.GoString(pathTmp[i]) == "0" {
			continue
		}
		resp[C.GoString(pathTmp[i])] = C.GoString(acctTmp[i])
	}
	return resp, nil
}

func splitServer(serverURL string) (*C.struct_Server, error) {
	u, err := url.Parse(serverURL)
	if err != nil {
		return nil, err
	}

	hostAndPort := strings.Split(u.Host, ":")
	host := hostAndPort[0]
	var port int
	if len(hostAndPort) == 2 {
		p, err := strconv.Atoi(hostAndPort[1])
		if err != nil {
			return nil, err
		}
		port = p
	}

	proto := C.kSecProtocolTypeHTTPS
	if u.Scheme != "https" {
		proto = C.kSecProtocolTypeHTTP
	}

	return &C.struct_Server{
		proto: C.SecProtocolType(proto),
		host:  C.CString(host),
		port:  C.uint(port),
		path:  C.CString(u.Path),
	}, nil
}

func freeServer(s *C.struct_Server) {
	C.free(unsafe.Pointer(s.host))
	C.free(unsafe.Pointer(s.path))
}
<commit_msg>mac: delete credentials before adding them to avoid already exist error (fixes #37)<commit_after>package osxkeychain

/*
#cgo CFLAGS: -x objective-c
#cgo LDFLAGS: -framework Security -framework Foundation

#include "osxkeychain_darwin.h"
#include <stdlib.h>
*/
import "C"
import (
	"errors"
	"net/url"
	"strconv"
	"strings"
	"unsafe"

	"github.com/docker/docker-credential-helpers/credentials"
)

// errCredentialsNotFound is the specific error message returned by OS X
// when the credentials are not in the keychain.
const errCredentialsNotFound = "The specified item could not be found in the keychain."

// Osxkeychain handles secrets using the OS X Keychain as store.
type Osxkeychain struct{}

// Add adds new credentials to the keychain.
func (h Osxkeychain) Add(creds *credentials.Credentials) error {
	h.Delete(creds.ServerURL)

	s, err := splitServer(creds.ServerURL)
	if err != nil {
		return err
	}
	defer freeServer(s)

	username := C.CString(creds.Username)
	defer C.free(unsafe.Pointer(username))
	secret := C.CString(creds.Secret)
	defer C.free(unsafe.Pointer(secret))

	errMsg := C.keychain_add(s, username, secret)
	if errMsg != nil {
		defer C.free(unsafe.Pointer(errMsg))
		return errors.New(C.GoString(errMsg))
	}

	return nil
}

// Delete removes credentials from the keychain.
func (h Osxkeychain) Delete(serverURL string) error {
	s, err := splitServer(serverURL)
	if err != nil {
		return err
	}
	defer freeServer(s)

	errMsg := C.keychain_delete(s)
	if errMsg != nil {
		defer C.free(unsafe.Pointer(errMsg))
		return errors.New(C.GoString(errMsg))
	}

	return nil
}

// Get returns the username and secret to use for a given registry server URL.
func (h Osxkeychain) Get(serverURL string) (string, string, error) {
	s, err := splitServer(serverURL)
	if err != nil {
		return "", "", err
	}
	defer freeServer(s)

	var usernameLen C.uint
	var username *C.char
	var secretLen C.uint
	var secret *C.char
	defer C.free(unsafe.Pointer(username))
	defer C.free(unsafe.Pointer(secret))

	errMsg := C.keychain_get(s, &usernameLen, &username, &secretLen, &secret)
	if errMsg != nil {
		defer C.free(unsafe.Pointer(errMsg))
		goMsg := C.GoString(errMsg)
		if goMsg == errCredentialsNotFound {
			return "", "", credentials.NewErrCredentialsNotFound()
		}

		return "", "", errors.New(goMsg)
	}

	user := C.GoStringN(username, C.int(usernameLen))
	pass := C.GoStringN(secret, C.int(secretLen))
	return user, pass, nil
}

// List returns the stored URLs and corresponding usernames.
func (h Osxkeychain) List() (map[string]string, error) {
	var pathsC **C.char
	defer C.free(unsafe.Pointer(pathsC))
	var acctsC **C.char
	defer C.free(unsafe.Pointer(acctsC))
	var listLenC C.uint
	errMsg := C.keychain_list(&pathsC, &acctsC, &listLenC)
	if errMsg != nil {
		defer C.free(unsafe.Pointer(errMsg))
		goMsg := C.GoString(errMsg)
		return nil, errors.New(goMsg)
	}

	defer C.freeListData(&pathsC, listLenC)
	defer C.freeListData(&acctsC, listLenC)

	var listLen int
	listLen = int(listLenC)
	pathTmp := (*[1 << 30]*C.char)(unsafe.Pointer(pathsC))[:listLen:listLen]
	acctTmp := (*[1 << 30]*C.char)(unsafe.Pointer(acctsC))[:listLen:listLen]
	//taking the array of c strings into go while ignoring all the stuff irrelevant to credentials-helper
	resp := make(map[string]string)
	for i := 0; i < listLen; i++ {
		if C.GoString(pathTmp[i]) == "0" {
			continue
		}
		resp[C.GoString(pathTmp[i])] = C.GoString(acctTmp[i])
	}
	return resp, nil
}

func splitServer(serverURL string) (*C.struct_Server, error) {
	u, err := url.Parse(serverURL)
	if err != nil {
		return nil, err
	}

	hostAndPort := strings.Split(u.Host, ":")
	host := hostAndPort[0]
	var port int
	if len(hostAndPort) == 2 {
		p, err := strconv.Atoi(hostAndPort[1])
		if err != nil {
			return nil, err
		}
		port = p
	}

	proto := C.kSecProtocolTypeHTTPS
	if u.Scheme != "https" {
		proto = C.kSecProtocolTypeHTTP
	}

	return &C.struct_Server{
		proto: C.SecProtocolType(proto),
		host:  C.CString(host),
		port:  C.uint(port),
		path:  C.CString(u.Path),
	}, nil
}

func freeServer(s *C.struct_Server) {
	C.free(unsafe.Pointer(s.host))
	C.free(unsafe.Pointer(s.path))
}
<|endoftext|>


<commit_before><commit_msg>Allow any ssh-user while cloning from git<commit_after><|endoftext|>