<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package systests

// Test various RPCs that are used mainly in other clients but not by the CLI.

import (
	"testing"
	"time"

	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/service"
	context "golang.org/x/net/context"
)

func TestAccountDeadlock(t *testing.T) {
	// t.Skip("Broken test; see CORE-5356")
	tc := setupTest(t, "resolve2")
	tc2 := cloneContext(tc)

	libkb.G.LocalDb = nil

	defer tc.Cleanup()

	stopCh := make(chan error)
	svc := service.NewService(tc.G, false)
	startCh := svc.GetStartChannel()
	go func() {
		err := svc.Run()
		if err != nil {
			t.Logf("Running the service produced an error: %v", err)
		}
		stopCh <- err
	}()

	<-startCh

	signupDoneCh := make(chan struct{})

	go func() {
		issueSignup(t, tc2.G)
		signupDoneCh <- struct{}{}
	}()

	currentStatusLoop(t, tc2.G, signupDoneCh)

	if err := client.CtlServiceStop(tc2.G); err != nil {
		t.Fatal(err)
	}

	// If the server failed, it's also an error
	if err := <-stopCh; err != nil {
		t.Fatal(err)
	}
}

func issueSignup(t *testing.T, g *libkb.GlobalContext) {
	cli, err := client.GetSignupClient(g)
	if err != nil {
		t.Fatalf("failed to get new identifyclient: %v", err)
	}

	id, err := libkb.RandString("", 5)
	if err != nil {
		t.Fatalf("Failed to get a random string: %s", err)
	}

	arg := keybase1.SignupArg{
		Email:      "test+" + id + "@keyba.se",
		Passphrase: "strong-password",
		Username:   "t_" + id,
		DeviceName: "dev0",
		InviteCode: "202020202020202020202020",
		DeviceType: keybase1.DeviceType_DESKTOP,
	}

	if _, err := cli.Signup(context.TODO(), arg); err != nil {
		t.Fatalf("signup failed: %s", err)
	}
}

func currentStatusLoop(t *testing.T, g *libkb.GlobalContext, stopCh chan struct{}) {
	cli, err := client.GetSessionClient(g)
	if err != nil {
		t.Fatal(err)
	}
	for {
		select {
		case <-stopCh:
			return
		case <-time.After(50 * time.Millisecond):
			_, err := cli.CurrentSession(context.TODO(), 0)
			if err != nil {
				t.Fatal(err)
			}
		}
	}
}
<commit_msg>Remove skip comment<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package systests

// Test various RPCs that are used mainly in other clients but not by the CLI.

import (
	"testing"
	"time"

	"github.com/keybase/client/go/client"
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	"github.com/keybase/client/go/service"
	context "golang.org/x/net/context"
)

func TestAccountDeadlock(t *testing.T) {
	tc := setupTest(t, "deadlock")
	tc2 := cloneContext(tc)

	libkb.G.LocalDb = nil

	defer tc.Cleanup()

	stopCh := make(chan error)
	svc := service.NewService(tc.G, false)
	startCh := svc.GetStartChannel()
	go func() {
		err := svc.Run()
		if err != nil {
			t.Logf("Running the service produced an error: %v", err)
		}
		stopCh <- err
	}()

	<-startCh

	signupDoneCh := make(chan struct{})

	go func() {
		issueSignup(t, tc2.G)
		signupDoneCh <- struct{}{}
	}()

	currentStatusLoop(t, tc2.G, signupDoneCh)

	if err := client.CtlServiceStop(tc2.G); err != nil {
		t.Fatal(err)
	}

	// If the server failed, it's also an error
	if err := <-stopCh; err != nil {
		t.Fatal(err)
	}
}

func issueSignup(t *testing.T, g *libkb.GlobalContext) {
	cli, err := client.GetSignupClient(g)
	if err != nil {
		t.Fatalf("failed to get new identifyclient: %v", err)
	}

	id, err := libkb.RandString("", 5)
	if err != nil {
		t.Fatalf("Failed to get a random string: %s", err)
	}

	arg := keybase1.SignupArg{
		Email:      "test+" + id + "@keyba.se",
		Passphrase: "strong-password",
		Username:   "t_" + id,
		DeviceName: "dev0",
		InviteCode: "202020202020202020202020",
		DeviceType: keybase1.DeviceType_DESKTOP,
	}

	if _, err := cli.Signup(context.TODO(), arg); err != nil {
		t.Fatalf("signup failed: %s", err)
	}
}

func currentStatusLoop(t *testing.T, g *libkb.GlobalContext, stopCh chan struct{}) {
	cli, err := client.GetSessionClient(g)
	if err != nil {
		t.Fatal(err)
	}
	for {
		select {
		case <-stopCh:
			return
		case <-time.After(50 * time.Millisecond):
			_, err := cli.CurrentSession(context.TODO(), 0)
			if err != nil {
				t.Fatal(err)
			}
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess 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 s3backupstorage implements the BackupStorage interface for AWS S3.
//
// AWS access credentials are configured via standard AWS means, such as:
// - AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables
// - credentials file at ~/.aws/credentials
// - if running on an EC2 instance, an IAM role
// See details at http://blogs.aws.amazon.com/security/post/Tx3D6U6WSFGOK2H/A-New-and-Standardized-Way-to-Manage-Credentials-in-the-AWS-SDKs
package s3backupstorage

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"math"
	"net/http"
	"sort"
	"strings"
	"sync"

	"vitess.io/vitess/go/vt/log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"golang.org/x/net/context"

	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/mysqlctl/backupstorage"
)

var (
	// AWS API region
	region = flag.String("s3_backup_aws_region", "us-east-1", "AWS region to use")

	// AWS request retries
	retryCount = flag.Int("s3_backup_aws_retries", -1, "AWS request retries")

	// AWS endpoint, defaults to amazonaws.com but appliances may use a different location
	endpoint = flag.String("s3_backup_aws_endpoint", "", "endpoint of the S3 backend (region must be provided)")

	// bucket is where the backups will go.
	bucket = flag.String("s3_backup_storage_bucket", "", "S3 bucket to use for backups")

	// root is a prefix added to all object names.
	root = flag.String("s3_backup_storage_root", "", "root prefix for all backup-related object names")

	// forcePath is used to ensure that the certificate and path used match the endpoint + region
	forcePath = flag.Bool("s3_backup_force_path_style", false, "force the s3 path style")

	tlsSkipVerifyCert = flag.Bool("s3_backup_tls_skip_verify_cert", false, "skip the 'certificate is valid' check for SSL connections")

	// verboseLogging provides more verbose logging of AWS actions
	requiredLogLevel = flag.String("s3_backup_log_level", "LogOff", "determine the S3 loglevel to use from LogOff, LogDebug, LogDebugWithSigning, LogDebugWithHTTPBody, LogDebugWithRequestRetries, LogDebugWithRequestErrors")

	// sse is the server-side encryption algorithm used when storing this object in S3
	sse = flag.String("s3_backup_server_side_encryption", "", "server-side encryption algorithm (e.g., AES256, aws:kms)")

	// path component delimiter
	delimiter = "/"
)

type logNameToLogLevel map[string]aws.LogLevelType

var logNameMap logNameToLogLevel

// S3BackupHandle implements the backupstorage.BackupHandle interface.
type S3BackupHandle struct {
	client    s3iface.S3API
	bs        *S3BackupStorage
	dir       string
	name      string
	readOnly  bool
	errors    concurrency.AllErrorRecorder
	waitGroup sync.WaitGroup
}

// Directory is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) Directory() string {
	return bh.dir
}

// Name is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) Name() string {
	return bh.name
}

// RecordError is part of the concurrency.ErrorRecorder interface.
func (bh *S3BackupHandle) RecordError(err error) {
	bh.errors.RecordError(err)
}

// HasErrors is part of the concurrency.ErrorRecorder interface.
func (bh *S3BackupHandle) HasErrors() bool {
	return bh.errors.HasErrors()
}

// Error is part of the concurrency.ErrorRecorder interface.
func (bh *S3BackupHandle) Error() error {
	return bh.errors.Error()
}

// AddFile is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) AddFile(ctx context.Context, filename string, filesize int64) (io.WriteCloser, error) {
	if bh.readOnly {
		return nil, fmt.Errorf("AddFile cannot be called on read-only backup")
	}

	// Calculate s3 upload part size using the source filesize
	partSizeBytes := s3manager.DefaultUploadPartSize
	if filesize > 0 {
		minimumPartSize := float64(filesize) / float64(s3manager.MaxUploadParts)
		// Round up to ensure large enough partsize
		calculatedPartSizeBytes := int64(math.Ceil(minimumPartSize))
		if calculatedPartSizeBytes > partSizeBytes {
			partSizeBytes = calculatedPartSizeBytes
		}
	}

	reader, writer := io.Pipe()
	bh.waitGroup.Add(1)

	go func() {
		defer bh.waitGroup.Done()
		uploader := s3manager.NewUploaderWithClient(bh.client, func(u *s3manager.Uploader) {
			u.PartSize = partSizeBytes
		})
		object := objName(bh.dir, bh.name, filename)

		var sseOption *string
		if *sse != "" {
			sseOption = sse
		}
		_, err := uploader.Upload(&s3manager.UploadInput{
			Bucket:               bucket,
			Key:                  object,
			Body:                 reader,
			ServerSideEncryption: sseOption,
		})
		if err != nil {
			reader.CloseWithError(err)
			bh.RecordError(err)
		}
	}()

	return writer, nil
}

// EndBackup is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) EndBackup(ctx context.Context) error {
	if bh.readOnly {
		return fmt.Errorf("EndBackup cannot be called on read-only backup")
	}
	bh.waitGroup.Wait()
	return bh.Error()
}

// AbortBackup is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) AbortBackup(ctx context.Context) error {
	if bh.readOnly {
		return fmt.Errorf("AbortBackup cannot be called on read-only backup")
	}
	return bh.bs.RemoveBackup(ctx, bh.dir, bh.name)
}

// ReadFile is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) ReadFile(ctx context.Context, filename string) (io.ReadCloser, error) {
	if !bh.readOnly {
		return nil, fmt.Errorf("ReadFile cannot be called on read-write backup")
	}
	object := objName(bh.dir, bh.name, filename)
	out, err := bh.client.GetObject(&s3.GetObjectInput{
		Bucket: bucket,
		Key:    object,
	})
	if err != nil {
		return nil, err
	}
	return out.Body, nil
}

var _ backupstorage.BackupHandle = (*S3BackupHandle)(nil)

// S3BackupStorage implements the backupstorage.BackupStorage interface.
type S3BackupStorage struct {
	_client *s3.S3
	mu      sync.Mutex
}

// ListBackups is part of the backupstorage.BackupStorage interface.
func (bs *S3BackupStorage) ListBackups(ctx context.Context, dir string) ([]backupstorage.BackupHandle, error) {
	log.Infof("ListBackups: [s3] dir: %v, bucket: %v", dir, *bucket)
	c, err := bs.client()
	if err != nil {
		return nil, err
	}

	searchPrefix := objName(dir, "")
	query := &s3.ListObjectsV2Input{
		Bucket:    bucket,
		Delimiter: &delimiter,
		Prefix:    searchPrefix,
	}

	var subdirs []string
	for {
		objs, err := c.ListObjectsV2(query)
		if err != nil {
			return nil, err
		}
		for _, prefix := range objs.CommonPrefixes {
			subdir := strings.TrimPrefix(*prefix.Prefix, *searchPrefix)
			subdir = strings.TrimSuffix(subdir, delimiter)
			subdirs = append(subdirs, subdir)
		}

		if objs.NextContinuationToken == nil {
			break
		}
		query.ContinuationToken = objs.NextContinuationToken
	}

	// Backups must be returned in order, oldest first.
	sort.Strings(subdirs)

	result := make([]backupstorage.BackupHandle, 0, len(subdirs))
	for _, subdir := range subdirs {
		result = append(result, &S3BackupHandle{
			client:   c,
			bs:       bs,
			dir:      dir,
			name:     subdir,
			readOnly: true,
		})
	}
	return result, nil
}

// StartBackup is part of the backupstorage.BackupStorage interface.
func (bs *S3BackupStorage) StartBackup(ctx context.Context, dir, name string) (backupstorage.BackupHandle, error) {
	log.Infof("StartBackup: [s3] dir: %v, name: %v, bucket: %v", dir, name, *bucket)
	c, err := bs.client()
	if err != nil {
		return nil, err
	}

	return &S3BackupHandle{
		client:   c,
		bs:       bs,
		dir:      dir,
		name:     name,
		readOnly: false,
	}, nil
}

// RemoveBackup is part of the backupstorage.BackupStorage interface.
func (bs *S3BackupStorage) RemoveBackup(ctx context.Context, dir, name string) error {
	log.Infof("RemoveBackup: [s3] dir: %v, name: %v, bucket: %v", dir, name, *bucket)

	c, err := bs.client()
	if err != nil {
		return err
	}

	query := &s3.ListObjectsV2Input{
		Bucket: bucket,
		Prefix: objName(dir, name),
	}

	for {
		objs, err := c.ListObjectsV2(query)
		if err != nil {
			return err
		}

		objIds := make([]*s3.ObjectIdentifier, 0, len(objs.Contents))
		for _, obj := range objs.Contents {
			objIds = append(objIds, &s3.ObjectIdentifier{
				Key: obj.Key,
			})
		}

		quiet := true // return less in the Delete response
		out, err := c.DeleteObjects(&s3.DeleteObjectsInput{
			Bucket: bucket,
			Delete: &s3.Delete{
				Objects: objIds,
				Quiet:   &quiet,
			},
		})

		if err != nil {
			return err
		}

		for _, objError := range out.Errors {
			return fmt.Errorf(objError.String())
		}

		if objs.NextContinuationToken == nil {
			break
		}

		query.ContinuationToken = objs.NextContinuationToken
	}

	return nil
}

// Close is part of the backupstorage.BackupStorage interface.
func (bs *S3BackupStorage) Close() error {
	bs.mu.Lock()
	defer bs.mu.Unlock()
	bs._client = nil
	return nil
}

var _ backupstorage.BackupStorage = (*S3BackupStorage)(nil)

// getLogLevel converts the string loglevel to an aws.LogLevelType
func getLogLevel() *aws.LogLevelType {
	l := new(aws.LogLevelType)
	*l = aws.LogOff // default setting
	if level, found := logNameMap[*requiredLogLevel]; found {
		*l = level // adjust as required
	}
	return l
}

func (bs *S3BackupStorage) client() (*s3.S3, error) {
	bs.mu.Lock()
	defer bs.mu.Unlock()
	if bs._client == nil {
		logLevel := getLogLevel()

		tlsClientConf := &tls.Config{InsecureSkipVerify: *tlsSkipVerifyCert}
		httpTransport := &http.Transport{TLSClientConfig: tlsClientConf}
		httpClient := &http.Client{Transport: httpTransport}

		session, err := session.NewSession()
		if err != nil {
			return nil, err
		}

		awsConfig := aws.Config{
			HTTPClient:       httpClient,
			LogLevel:         logLevel,
			Endpoint:         aws.String(*endpoint),
			Region:           aws.String(*region),
			S3ForcePathStyle: aws.Bool(*forcePath),
		}

		if *retryCount >= 0 {
			awsConfig = *request.WithRetryer(&awsConfig, &ClosedConnectionRetryer{
				awsRetryer: &client.DefaultRetryer{
					NumMaxRetries: *retryCount,
				},
			})
		}

		bs._client = s3.New(session, &awsConfig)

		if len(*bucket) == 0 {
			return nil, fmt.Errorf("-s3_backup_storage_bucket required")
		}

		if _, err := bs._client.HeadBucket(&s3.HeadBucketInput{Bucket: bucket}); err != nil {
			return nil, err
		}
	}
	return bs._client, nil
}

func objName(parts ...string) *string {
	res := ""
	if *root != "" {
		res += *root + delimiter
	}
	res += strings.Join(parts, delimiter)
	return &res
}

func init() {
	backupstorage.BackupStorageMap["s3"] = &S3BackupStorage{}

	logNameMap = logNameToLogLevel{
		"LogOff":                     aws.LogOff,
		"LogDebug":                   aws.LogDebug,
		"LogDebugWithSigning":        aws.LogDebugWithSigning,
		"LogDebugWithHTTPBody":       aws.LogDebugWithHTTPBody,
		"LogDebugWithRequestRetries": aws.LogDebugWithRequestRetries,
		"LogDebugWithRequestErrors":  aws.LogDebugWithRequestErrors,
	}
}
<commit_msg>fix s3 ListAllBackups() to list backups from the root directory<commit_after>/*
Copyright 2019 The Vitess 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 s3backupstorage implements the BackupStorage interface for AWS S3.
//
// AWS access credentials are configured via standard AWS means, such as:
// - AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables
// - credentials file at ~/.aws/credentials
// - if running on an EC2 instance, an IAM role
// See details at http://blogs.aws.amazon.com/security/post/Tx3D6U6WSFGOK2H/A-New-and-Standardized-Way-to-Manage-Credentials-in-the-AWS-SDKs
package s3backupstorage

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"math"
	"net/http"
	"sort"
	"strings"
	"sync"

	"vitess.io/vitess/go/vt/log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/client"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"golang.org/x/net/context"

	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/mysqlctl/backupstorage"
)

var (
	// AWS API region
	region = flag.String("s3_backup_aws_region", "us-east-1", "AWS region to use")

	// AWS request retries
	retryCount = flag.Int("s3_backup_aws_retries", -1, "AWS request retries")

	// AWS endpoint, defaults to amazonaws.com but appliances may use a different location
	endpoint = flag.String("s3_backup_aws_endpoint", "", "endpoint of the S3 backend (region must be provided)")

	// bucket is where the backups will go.
	bucket = flag.String("s3_backup_storage_bucket", "", "S3 bucket to use for backups")

	// root is a prefix added to all object names.
	root = flag.String("s3_backup_storage_root", "", "root prefix for all backup-related object names")

	// forcePath is used to ensure that the certificate and path used match the endpoint + region
	forcePath = flag.Bool("s3_backup_force_path_style", false, "force the s3 path style")

	tlsSkipVerifyCert = flag.Bool("s3_backup_tls_skip_verify_cert", false, "skip the 'certificate is valid' check for SSL connections")

	// verboseLogging provides more verbose logging of AWS actions
	requiredLogLevel = flag.String("s3_backup_log_level", "LogOff", "determine the S3 loglevel to use from LogOff, LogDebug, LogDebugWithSigning, LogDebugWithHTTPBody, LogDebugWithRequestRetries, LogDebugWithRequestErrors")

	// sse is the server-side encryption algorithm used when storing this object in S3
	sse = flag.String("s3_backup_server_side_encryption", "", "server-side encryption algorithm (e.g., AES256, aws:kms)")

	// path component delimiter
	delimiter = "/"
)

type logNameToLogLevel map[string]aws.LogLevelType

var logNameMap logNameToLogLevel

// S3BackupHandle implements the backupstorage.BackupHandle interface.
type S3BackupHandle struct {
	client    s3iface.S3API
	bs        *S3BackupStorage
	dir       string
	name      string
	readOnly  bool
	errors    concurrency.AllErrorRecorder
	waitGroup sync.WaitGroup
}

// Directory is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) Directory() string {
	return bh.dir
}

// Name is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) Name() string {
	return bh.name
}

// RecordError is part of the concurrency.ErrorRecorder interface.
func (bh *S3BackupHandle) RecordError(err error) {
	bh.errors.RecordError(err)
}

// HasErrors is part of the concurrency.ErrorRecorder interface.
func (bh *S3BackupHandle) HasErrors() bool {
	return bh.errors.HasErrors()
}

// Error is part of the concurrency.ErrorRecorder interface.
func (bh *S3BackupHandle) Error() error {
	return bh.errors.Error()
}

// AddFile is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) AddFile(ctx context.Context, filename string, filesize int64) (io.WriteCloser, error) {
	if bh.readOnly {
		return nil, fmt.Errorf("AddFile cannot be called on read-only backup")
	}

	// Calculate s3 upload part size using the source filesize
	partSizeBytes := s3manager.DefaultUploadPartSize
	if filesize > 0 {
		minimumPartSize := float64(filesize) / float64(s3manager.MaxUploadParts)
		// Round up to ensure large enough partsize
		calculatedPartSizeBytes := int64(math.Ceil(minimumPartSize))
		if calculatedPartSizeBytes > partSizeBytes {
			partSizeBytes = calculatedPartSizeBytes
		}
	}

	reader, writer := io.Pipe()
	bh.waitGroup.Add(1)

	go func() {
		defer bh.waitGroup.Done()
		uploader := s3manager.NewUploaderWithClient(bh.client, func(u *s3manager.Uploader) {
			u.PartSize = partSizeBytes
		})
		object := objName(bh.dir, bh.name, filename)

		var sseOption *string
		if *sse != "" {
			sseOption = sse
		}
		_, err := uploader.Upload(&s3manager.UploadInput{
			Bucket:               bucket,
			Key:                  object,
			Body:                 reader,
			ServerSideEncryption: sseOption,
		})
		if err != nil {
			reader.CloseWithError(err)
			bh.RecordError(err)
		}
	}()

	return writer, nil
}

// EndBackup is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) EndBackup(ctx context.Context) error {
	if bh.readOnly {
		return fmt.Errorf("EndBackup cannot be called on read-only backup")
	}
	bh.waitGroup.Wait()
	return bh.Error()
}

// AbortBackup is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) AbortBackup(ctx context.Context) error {
	if bh.readOnly {
		return fmt.Errorf("AbortBackup cannot be called on read-only backup")
	}
	return bh.bs.RemoveBackup(ctx, bh.dir, bh.name)
}

// ReadFile is part of the backupstorage.BackupHandle interface.
func (bh *S3BackupHandle) ReadFile(ctx context.Context, filename string) (io.ReadCloser, error) {
	if !bh.readOnly {
		return nil, fmt.Errorf("ReadFile cannot be called on read-write backup")
	}
	object := objName(bh.dir, bh.name, filename)
	out, err := bh.client.GetObject(&s3.GetObjectInput{
		Bucket: bucket,
		Key:    object,
	})
	if err != nil {
		return nil, err
	}
	return out.Body, nil
}

var _ backupstorage.BackupHandle = (*S3BackupHandle)(nil)

// S3BackupStorage implements the backupstorage.BackupStorage interface.
type S3BackupStorage struct {
	_client *s3.S3
	mu      sync.Mutex
}

// ListBackups is part of the backupstorage.BackupStorage interface.
func (bs *S3BackupStorage) ListBackups(ctx context.Context, dir string) ([]backupstorage.BackupHandle, error) {
	log.Infof("ListBackups: [s3] dir: %v, bucket: %v", dir, *bucket)
	c, err := bs.client()
	if err != nil {
		return nil, err
	}

	var searchPrefix *string
	if dir == "/" {
		searchPrefix = objName("")
	} else {
		searchPrefix = objName(dir, "")
	}
	log.Infof("objName: %v", searchPrefix)

	query := &s3.ListObjectsV2Input{
		Bucket:    bucket,
		Delimiter: &delimiter,
		Prefix:    searchPrefix,
	}

	var subdirs []string
	for {
		objs, err := c.ListObjectsV2(query)
		if err != nil {
			return nil, err
		}
		for _, prefix := range objs.CommonPrefixes {
			subdir := strings.TrimPrefix(*prefix.Prefix, *searchPrefix)
			subdir = strings.TrimSuffix(subdir, delimiter)
			subdirs = append(subdirs, subdir)
		}

		if objs.NextContinuationToken == nil {
			break
		}
		query.ContinuationToken = objs.NextContinuationToken
	}

	// Backups must be returned in order, oldest first.
	sort.Strings(subdirs)

	result := make([]backupstorage.BackupHandle, 0, len(subdirs))
	for _, subdir := range subdirs {
		result = append(result, &S3BackupHandle{
			client:   c,
			bs:       bs,
			dir:      dir,
			name:     subdir,
			readOnly: true,
		})
	}
	return result, nil
}

// StartBackup is part of the backupstorage.BackupStorage interface.
func (bs *S3BackupStorage) StartBackup(ctx context.Context, dir, name string) (backupstorage.BackupHandle, error) {
	log.Infof("StartBackup: [s3] dir: %v, name: %v, bucket: %v", dir, name, *bucket)
	c, err := bs.client()
	if err != nil {
		return nil, err
	}

	return &S3BackupHandle{
		client:   c,
		bs:       bs,
		dir:      dir,
		name:     name,
		readOnly: false,
	}, nil
}

// RemoveBackup is part of the backupstorage.BackupStorage interface.
func (bs *S3BackupStorage) RemoveBackup(ctx context.Context, dir, name string) error {
	log.Infof("RemoveBackup: [s3] dir: %v, name: %v, bucket: %v", dir, name, *bucket)

	c, err := bs.client()
	if err != nil {
		return err
	}

	query := &s3.ListObjectsV2Input{
		Bucket: bucket,
		Prefix: objName(dir, name),
	}

	for {
		objs, err := c.ListObjectsV2(query)
		if err != nil {
			return err
		}

		objIds := make([]*s3.ObjectIdentifier, 0, len(objs.Contents))
		for _, obj := range objs.Contents {
			objIds = append(objIds, &s3.ObjectIdentifier{
				Key: obj.Key,
			})
		}

		quiet := true // return less in the Delete response
		out, err := c.DeleteObjects(&s3.DeleteObjectsInput{
			Bucket: bucket,
			Delete: &s3.Delete{
				Objects: objIds,
				Quiet:   &quiet,
			},
		})

		if err != nil {
			return err
		}

		for _, objError := range out.Errors {
			return fmt.Errorf(objError.String())
		}

		if objs.NextContinuationToken == nil {
			break
		}

		query.ContinuationToken = objs.NextContinuationToken
	}

	return nil
}

// Close is part of the backupstorage.BackupStorage interface.
func (bs *S3BackupStorage) Close() error {
	bs.mu.Lock()
	defer bs.mu.Unlock()
	bs._client = nil
	return nil
}

var _ backupstorage.BackupStorage = (*S3BackupStorage)(nil)

// getLogLevel converts the string loglevel to an aws.LogLevelType
func getLogLevel() *aws.LogLevelType {
	l := new(aws.LogLevelType)
	*l = aws.LogOff // default setting
	if level, found := logNameMap[*requiredLogLevel]; found {
		*l = level // adjust as required
	}
	return l
}

func (bs *S3BackupStorage) client() (*s3.S3, error) {
	bs.mu.Lock()
	defer bs.mu.Unlock()
	if bs._client == nil {
		logLevel := getLogLevel()

		tlsClientConf := &tls.Config{InsecureSkipVerify: *tlsSkipVerifyCert}
		httpTransport := &http.Transport{TLSClientConfig: tlsClientConf}
		httpClient := &http.Client{Transport: httpTransport}

		session, err := session.NewSession()
		if err != nil {
			return nil, err
		}

		awsConfig := aws.Config{
			HTTPClient:       httpClient,
			LogLevel:         logLevel,
			Endpoint:         aws.String(*endpoint),
			Region:           aws.String(*region),
			S3ForcePathStyle: aws.Bool(*forcePath),
		}

		if *retryCount >= 0 {
			awsConfig = *request.WithRetryer(&awsConfig, &ClosedConnectionRetryer{
				awsRetryer: &client.DefaultRetryer{
					NumMaxRetries: *retryCount,
				},
			})
		}

		bs._client = s3.New(session, &awsConfig)

		if len(*bucket) == 0 {
			return nil, fmt.Errorf("-s3_backup_storage_bucket required")
		}

		if _, err := bs._client.HeadBucket(&s3.HeadBucketInput{Bucket: bucket}); err != nil {
			return nil, err
		}
	}
	return bs._client, nil
}

func objName(parts ...string) *string {
	res := ""
	if *root != "" {
		res += *root + delimiter
	}
	res += strings.Join(parts, delimiter)
	return &res
}

func init() {
	backupstorage.BackupStorageMap["s3"] = &S3BackupStorage{}

	logNameMap = logNameToLogLevel{
		"LogOff":                     aws.LogOff,
		"LogDebug":                   aws.LogDebug,
		"LogDebugWithSigning":        aws.LogDebugWithSigning,
		"LogDebugWithHTTPBody":       aws.LogDebugWithHTTPBody,
		"LogDebugWithRequestRetries": aws.LogDebugWithRequestRetries,
		"LogDebugWithRequestErrors":  aws.LogDebugWithRequestErrors,
	}
}
<|endoftext|>


<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package sqlannotation provides functions
// for annotating DML statements with keyspace-id
// comments and parsing them. These annotations
// are used during filtered-replication to route
// the DML statement to the correct shard.
// TOOD(erez): Move the code for the "_stream" annotations
// from vttablet to here.
package sqlannotation

import (
	"encoding/hex"
	"fmt"
	"log"
	"strings"

	"github.com/youtube/vitess/go/stats"
)

const (
	filteredReplicationUnfriendlyComment = "/* vtgate:: filtered_replication_unfriendly */"
	insertDML                            = "insert"
	updateDML                            = "update"
	deleteDML                            = "delete"
)

var (
	filteredReplicationUnfriendlyCountStat = stats.NewInt("ReplicationUnfriendlyStatementsCount")
)

// AddIfDML annotates 'sql' based on
// the keyspaceIDs found in 'keyspaceIDs':
//
// If 'sql' is not a DML statement no annotation is added.
// If 'sql' is a DML statement and contains exactly one element
// it is used to annotate 'sql'; otherwise 'sql' is annotated
// as replication-unfriendly.
func AddIfDML(sql string, keyspaceIDs [][]byte) string {
	if len(keyspaceIDs) == 1 {
		return AddKeyspaceIDIfDML(sql, keyspaceIDs[0])
	}
	return AddFilteredReplicationUnfriendlyIfDML(sql)
}

// AddKeyspaceIDIfDML returns a copy of 'sql' annotated
// with the given keyspace id if 'sql' is a DML statement;
// otherwise, it returns a copy of 'sql'.
func AddKeyspaceIDIfDML(sql string, keyspaceID []byte) string {
	if isDml(sql) {
		return AddKeyspaceID(sql, keyspaceID)
	}
	return sql
}

// AddKeyspaceID returns a copy of 'sql' annotated
// with the given keyspace id.
func AddKeyspaceID(sql string, keyspaceID []byte) string {
	return fmt.Sprintf("%s /* vtgate:: keyspace_id:%s */",
		sql, hex.EncodeToString(keyspaceID))
}

// AddFilteredReplicationUnfriendlyIfDML annotates the given 'sql'
// query as filtered-replication-unfriendly if its a DML statement
// (filtered-replication should halt if it encounters such an annotation).
// Does nothing if the statement is not a DML statement.
func AddFilteredReplicationUnfriendlyIfDML(sql string) string {
	if isDml(sql) {
		return AddFilteredReplicationUnfriendly(sql)
	}
	return sql
}

// AddFilteredReplicationUnfriendly annotates the given 'sql'
// query as filtered-replication-unfriendly.
func AddFilteredReplicationUnfriendly(sql string) string {
	replicationUnfriendlyCountStat.Add(1)
	return sql + filteredReplicationUnfriendlyComment
}

// Copied from go/vt/vtgate/resolver.go
// TODO(erez): Refactor this.

// Returns true if 'querySQL' is an INSERT, UPDATE or DELETE
// statement.
func isDml(querySQL string) bool {
	var sqlKW string
	if i := strings.Index(querySQL, " "); i >= 0 {
		sqlKW = querySQL[:i]
	}
	sqlKW = strings.ToLower(sqlKW)
	if sqlKW == insertDML || sqlKW == updateDML || sqlKW == deleteDML {
		return true
	}
	return false
}

// ExtractKeySpaceID parses the annotation of the given statement and tries
// to extract the keyspace id.
// If a keyspace-id comment exists 'keyspaceID' is set to the parsed keyspace id
// and err is set to nil; otherwise, if a filtered-replication-unfriendly comment exists
// or some other parsing error occured, keyspaceID is set to nil and err is set to a non-nil
// error value.
func ExtractKeySpaceID(sql string) (keyspaceID []byte, err error) {
	keyspaceIDString, hasKeySpaceID := extractStringBetween(sql, "/* vtgate:: keyspace_id:", " ")
	hasUnfriendlyAnnotation := (strings.Index(sql, filteredReplicationUnfriendlyComment) != -1)
	err = nil
	if hasKeySpaceID {
		if hasUnfriendlyAnnotation {
			keyspaceID = nil
			err = &ExtractKeySpaceIDError{
				Kind:    ExtractKeySpaceIDParseError,
				Message: fmt.Sprintf("Conflicting annotations in statement '%v'", sql),
			}
			return
		}
		keyspaceID, err = hex.DecodeString(keyspaceIDString)
		if err != nil {
			keyspaceID = nil
			err = &ExtractKeySpaceIDError{
				Kind: ExtractKeySpaceIDParseError,
				Message: fmt.Sprintf(
					"Error parsing keyspace id value in statement: %v (%v)", sql, err),
			}
		}
		return
	}

	if hasUnfriendlyAnnotation {
		err = &ExtractKeySpaceIDError{
			Kind:    ExtractKeySpaceIDReplicationUnfriendlyError,
			Message: fmt.Sprintf("Statement: %v", sql),
		}
		keyspaceID = nil
		return
	}

	// No annotations.
	keyspaceID = nil
	err = &ExtractKeySpaceIDError{
		Kind:    ExtractKeySpaceIDParseError,
		Message: fmt.Sprintf("No annotation found in '%v'", sql),
	}
	return
}

// Extracts the string from source contained between the leftmost instance of
// 'leftDelim' and the next instance of 'rightDelim'. If there is no next instance
// of 'rightDelim', returns the string contained between the end of the leftmost instance
// of 'leftDelim' to the end of 'source'. If 'leftDelim' does not appear in 'source',
// sets 'found' to false and 'match' to the empty string, otherwise 'found' is set to true
// and 'match' is set to the extracted string.
func extractStringBetween(source string, leftDelim string, rightDelim string) (match string, found bool) {
	leftDelimStart := strings.Index(source, leftDelim)
	if leftDelimStart == -1 {
		found = false
		match = ""
		return
	}
	found = true
	matchStart := leftDelimStart + len(leftDelim)
	matchEnd := strings.Index(source[matchStart:], rightDelim)
	if matchEnd != -1 {
		match = source[matchStart : matchStart+matchEnd]
		return
	}
	match = source[matchStart:]
	return
}

// ExtractKeySpaceIDError is the error type returned
// from ExtractKeySpaceID
// Kind is a numeric code for the error (see constants below)
// and Message is an error message string.
type ExtractKeySpaceIDError struct {
	Kind    int
	Message string
}

// Possible values for ExtractKeySpaceIDError.Kind
const (
	ExtractKeySpaceIDParseError                 = iota
	ExtractKeySpaceIDReplicationUnfriendlyError = iota
)

func (err ExtractKeySpaceIDError) Error() string {
	switch err.Kind {
	case ExtractKeySpaceIDParseError:
		return fmt.Sprintf("Parse-Error. %v", err.Message)
	case ExtractKeySpaceIDReplicationUnfriendlyError:
		return fmt.Sprintf(
			"Statement is filtered-replication-unfriendly. %v", err.Message)
	default:
		log.Fatalf("Unknown error type: %v", err)
		return "" // Unreachable.
	}
}
<commit_msg>Renamed go variable at call-site as well.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package sqlannotation provides functions
// for annotating DML statements with keyspace-id
// comments and parsing them. These annotations
// are used during filtered-replication to route
// the DML statement to the correct shard.
// TOOD(erez): Move the code for the "_stream" annotations
// from vttablet to here.
package sqlannotation

import (
	"encoding/hex"
	"fmt"
	"log"
	"strings"

	"github.com/youtube/vitess/go/stats"
)

const (
	filteredReplicationUnfriendlyComment = "/* vtgate:: filtered_replication_unfriendly */"
	insertDML                            = "insert"
	updateDML                            = "update"
	deleteDML                            = "delete"
)

var (
	filteredReplicationUnfriendlyCountStat = stats.NewInt("ReplicationUnfriendlyStatementsCount")
)

// AddIfDML annotates 'sql' based on
// the keyspaceIDs found in 'keyspaceIDs':
//
// If 'sql' is not a DML statement no annotation is added.
// If 'sql' is a DML statement and contains exactly one element
// it is used to annotate 'sql'; otherwise 'sql' is annotated
// as replication-unfriendly.
func AddIfDML(sql string, keyspaceIDs [][]byte) string {
	if len(keyspaceIDs) == 1 {
		return AddKeyspaceIDIfDML(sql, keyspaceIDs[0])
	}
	return AddFilteredReplicationUnfriendlyIfDML(sql)
}

// AddKeyspaceIDIfDML returns a copy of 'sql' annotated
// with the given keyspace id if 'sql' is a DML statement;
// otherwise, it returns a copy of 'sql'.
func AddKeyspaceIDIfDML(sql string, keyspaceID []byte) string {
	if isDml(sql) {
		return AddKeyspaceID(sql, keyspaceID)
	}
	return sql
}

// AddKeyspaceID returns a copy of 'sql' annotated
// with the given keyspace id.
func AddKeyspaceID(sql string, keyspaceID []byte) string {
	return fmt.Sprintf("%s /* vtgate:: keyspace_id:%s */",
		sql, hex.EncodeToString(keyspaceID))
}

// AddFilteredReplicationUnfriendlyIfDML annotates the given 'sql'
// query as filtered-replication-unfriendly if its a DML statement
// (filtered-replication should halt if it encounters such an annotation).
// Does nothing if the statement is not a DML statement.
func AddFilteredReplicationUnfriendlyIfDML(sql string) string {
	if isDml(sql) {
		return AddFilteredReplicationUnfriendly(sql)
	}
	return sql
}

// AddFilteredReplicationUnfriendly annotates the given 'sql'
// query as filtered-replication-unfriendly.
func AddFilteredReplicationUnfriendly(sql string) string {
	filteredReplicationUnfriendlyCountStat.Add(1)
	return sql + filteredReplicationUnfriendlyComment
}

// Copied from go/vt/vtgate/resolver.go
// TODO(erez): Refactor this.

// Returns true if 'querySQL' is an INSERT, UPDATE or DELETE
// statement.
func isDml(querySQL string) bool {
	var sqlKW string
	if i := strings.Index(querySQL, " "); i >= 0 {
		sqlKW = querySQL[:i]
	}
	sqlKW = strings.ToLower(sqlKW)
	if sqlKW == insertDML || sqlKW == updateDML || sqlKW == deleteDML {
		return true
	}
	return false
}

// ExtractKeySpaceID parses the annotation of the given statement and tries
// to extract the keyspace id.
// If a keyspace-id comment exists 'keyspaceID' is set to the parsed keyspace id
// and err is set to nil; otherwise, if a filtered-replication-unfriendly comment exists
// or some other parsing error occured, keyspaceID is set to nil and err is set to a non-nil
// error value.
func ExtractKeySpaceID(sql string) (keyspaceID []byte, err error) {
	keyspaceIDString, hasKeySpaceID := extractStringBetween(sql, "/* vtgate:: keyspace_id:", " ")
	hasUnfriendlyAnnotation := (strings.Index(sql, filteredReplicationUnfriendlyComment) != -1)
	err = nil
	if hasKeySpaceID {
		if hasUnfriendlyAnnotation {
			keyspaceID = nil
			err = &ExtractKeySpaceIDError{
				Kind:    ExtractKeySpaceIDParseError,
				Message: fmt.Sprintf("Conflicting annotations in statement '%v'", sql),
			}
			return
		}
		keyspaceID, err = hex.DecodeString(keyspaceIDString)
		if err != nil {
			keyspaceID = nil
			err = &ExtractKeySpaceIDError{
				Kind: ExtractKeySpaceIDParseError,
				Message: fmt.Sprintf(
					"Error parsing keyspace id value in statement: %v (%v)", sql, err),
			}
		}
		return
	}

	if hasUnfriendlyAnnotation {
		err = &ExtractKeySpaceIDError{
			Kind:    ExtractKeySpaceIDReplicationUnfriendlyError,
			Message: fmt.Sprintf("Statement: %v", sql),
		}
		keyspaceID = nil
		return
	}

	// No annotations.
	keyspaceID = nil
	err = &ExtractKeySpaceIDError{
		Kind:    ExtractKeySpaceIDParseError,
		Message: fmt.Sprintf("No annotation found in '%v'", sql),
	}
	return
}

// Extracts the string from source contained between the leftmost instance of
// 'leftDelim' and the next instance of 'rightDelim'. If there is no next instance
// of 'rightDelim', returns the string contained between the end of the leftmost instance
// of 'leftDelim' to the end of 'source'. If 'leftDelim' does not appear in 'source',
// sets 'found' to false and 'match' to the empty string, otherwise 'found' is set to true
// and 'match' is set to the extracted string.
func extractStringBetween(source string, leftDelim string, rightDelim string) (match string, found bool) {
	leftDelimStart := strings.Index(source, leftDelim)
	if leftDelimStart == -1 {
		found = false
		match = ""
		return
	}
	found = true
	matchStart := leftDelimStart + len(leftDelim)
	matchEnd := strings.Index(source[matchStart:], rightDelim)
	if matchEnd != -1 {
		match = source[matchStart : matchStart+matchEnd]
		return
	}
	match = source[matchStart:]
	return
}

// ExtractKeySpaceIDError is the error type returned
// from ExtractKeySpaceID
// Kind is a numeric code for the error (see constants below)
// and Message is an error message string.
type ExtractKeySpaceIDError struct {
	Kind    int
	Message string
}

// Possible values for ExtractKeySpaceIDError.Kind
const (
	ExtractKeySpaceIDParseError                 = iota
	ExtractKeySpaceIDReplicationUnfriendlyError = iota
)

func (err ExtractKeySpaceIDError) Error() string {
	switch err.Kind {
	case ExtractKeySpaceIDParseError:
		return fmt.Sprintf("Parse-Error. %v", err.Message)
	case ExtractKeySpaceIDReplicationUnfriendlyError:
		return fmt.Sprintf(
			"Statement is filtered-replication-unfriendly. %v", err.Message)
	default:
		log.Fatalf("Unknown error type: %v", err)
		return "" // Unreachable.
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess 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 vtgate

import (
	"testing"
	"time"

	"context"

	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/discovery"
	querypb "vitess.io/vitess/go/vt/proto/query"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	_ "vitess.io/vitess/go/vt/vtgate/vindexes"
	"vitess.io/vitess/go/vt/vttablet/sandboxconn"
)

func TestStreamSQLUnsharded(t *testing.T) {
	executor, _, _, _ := createLegacyExecutorEnv()
	logChan := QueryLogger.Subscribe("Test")
	defer QueryLogger.Unsubscribe(logChan)

	sql := "stream * from user_msgs"
	result, err := executorStreamMessages(executor, sql)
	require.NoError(t, err)
	wantResult := sandboxconn.StreamRowResult
	if !result.Equal(wantResult) {
		t.Errorf("result: %+v, want %+v", result, wantResult)
	}
}

func TestStreamSQLSharded(t *testing.T) {
	// Special setup: Don't use createLegacyExecutorEnv.
	cell := "aa"
	hc := discovery.NewFakeLegacyHealthCheck()
	s := createSandbox("TestExecutor")
	s.VSchema = executorVSchema
	getSandbox(KsTestUnsharded).VSchema = unshardedVSchema
	serv := new(sandboxTopo)
	resolver := newTestLegacyResolver(hc, serv, cell)
	shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"}
	for _, shard := range shards {
		_ = hc.AddTestTablet(cell, shard, 1, "TestExecutor", shard, topodatapb.TabletType_MASTER, true, 1, nil)
	}
	executor := NewExecutor(context.Background(), serv, cell, resolver, false, testBufferSize, testCacheSize)

	sql := "stream * from sharded_user_msgs"
	result, err := executorStreamMessages(executor, sql)
	require.NoError(t, err)
	wantResult := &sqltypes.Result{
		Fields: sandboxconn.SingleRowResult.Fields,
		Rows: [][]sqltypes.Value{
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
		},
	}
	if !result.Equal(wantResult) {
		t.Errorf("result: %+v, want %+v", result, wantResult)
	}
}

func executorStreamMessages(executor *Executor, sql string) (qr *sqltypes.Result, err error) {
	results := make(chan *sqltypes.Result, 100)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()
	err = executor.StreamExecute(
		ctx,
		"TestExecuteStream",
		NewSafeSession(masterSession),
		sql,
		nil,
		querypb.Target{
			TabletType: topodatapb.TabletType_MASTER,
		},
		func(qr *sqltypes.Result) error {
			results <- qr
			return nil
		},
	)
	close(results)
	if err != nil {
		return nil, err
	}
	first := true
	for r := range results {
		if first {
			qr = &sqltypes.Result{Fields: r.Fields}
			first = false
		}
		qr.Rows = append(qr.Rows, r.Rows...)
	}
	return qr, nil
}
<commit_msg>added executor stream unit test<commit_after>/*
Copyright 2019 The Vitess 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 vtgate

import (
	"testing"
	"time"

	"context"

	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/discovery"
	querypb "vitess.io/vitess/go/vt/proto/query"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	_ "vitess.io/vitess/go/vt/vtgate/vindexes"
	"vitess.io/vitess/go/vt/vttablet/sandboxconn"
)

func TestStreamSQLUnsharded(t *testing.T) {
	executor, _, _, _ := createLegacyExecutorEnv()
	logChan := QueryLogger.Subscribe("Test")
	defer QueryLogger.Unsubscribe(logChan)

	sql := "stream * from user_msgs"
	result, err := executorStreamMessages(executor, sql)
	require.NoError(t, err)
	wantResult := sandboxconn.StreamRowResult
	if !result.Equal(wantResult) {
		t.Errorf("result: %+v, want %+v", result, wantResult)
	}
}

func TestStreamSQLSharded(t *testing.T) {
	// Special setup: Don't use createLegacyExecutorEnv.
	cell := "aa"
	hc := discovery.NewFakeLegacyHealthCheck()
	s := createSandbox("TestExecutor")
	s.VSchema = executorVSchema
	getSandbox(KsTestUnsharded).VSchema = unshardedVSchema
	serv := new(sandboxTopo)
	resolver := newTestLegacyResolver(hc, serv, cell)
	shards := []string{"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-"}
	for _, shard := range shards {
		_ = hc.AddTestTablet(cell, shard, 1, "TestExecutor", shard, topodatapb.TabletType_MASTER, true, 1, nil)
	}
	executor := NewExecutor(context.Background(), serv, cell, resolver, false, testBufferSize, testCacheSize)

	sql := "stream * from sharded_user_msgs"
	result, err := executorStreamMessages(executor, sql)
	require.NoError(t, err)
	wantResult := &sqltypes.Result{
		Fields: sandboxconn.SingleRowResult.Fields,
		Rows: [][]sqltypes.Value{
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
			sandboxconn.StreamRowResult.Rows[0],
		},
	}
	if !result.Equal(wantResult) {
		t.Errorf("result: %+v, want %+v", result, wantResult)
	}
}

func TestStreamError(t *testing.T) {
	executor, _, _, _ := createLegacyExecutorEnv()
	logChan := QueryLogger.Subscribe("TestStreamError")
	defer QueryLogger.Unsubscribe(logChan)

	queries := []string{
		"start transaction",
		"begin",
		"rollback",
		"commit",
	}
	for _, query := range queries {
		t.Run(query, func(t *testing.T) {
			_, err := executorStreamMessages(executor, query)
			require.Error(t, err)
			require.Contains(t, err.Error(), "unsupported statement type for OLAP")
		})
	}
}

func executorStreamMessages(executor *Executor, sql string) (qr *sqltypes.Result, err error) {
	results := make(chan *sqltypes.Result, 100)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()
	err = executor.StreamExecute(
		ctx,
		"TestExecuteStream",
		NewSafeSession(masterSession),
		sql,
		nil,
		querypb.Target{
			TabletType: topodatapb.TabletType_MASTER,
		},
		func(qr *sqltypes.Result) error {
			results <- qr
			return nil
		},
	)
	close(results)
	if err != nil {
		return nil, err
	}
	first := true
	for r := range results {
		if first {
			qr = &sqltypes.Result{Fields: r.Fields}
			first = false
		}
		qr.Rows = append(qr.Rows, r.Rows...)
	}
	return qr, nil
}
<|endoftext|>


<commit_before>package mo

import (
	"runtime"
	"time"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/mem"

	"runtime/pprof"

	n "github.com/shirou/gopsutil/net"
)

//Snapshot Snapshot
type Snapshot struct {
	CPUPercent        float64
	MemPercent        float64
	AverageNetIn      uint64
	AverageNetOut     uint64
	Thread            uint64
	Goroutine         uint64
	TotalConn         uint64
	TotalTopic        uint64
	TotalConsumer     uint64
	MaxConsumer       uint64
	TotalMessageIn    uint64
	AverageMessageIn  uint64
	TotalMessageOut   uint64
	AverageMessageOut uint64
}

//DataSource DataSource
type DataSource interface {
	GetTotalConsumer() map[string]uint64
	GetTotalConn() uint64
	GetTotalTopic() uint64
	GetMessageIn() uint64
	GetMessageOut() uint64
}

//Monitor Monitor
type Monitor struct {
	dataSource DataSource
}

//NewMonitor NewMonitor
func NewMonitor(dataSource DataSource) *Monitor {
	return &Monitor{dataSource: dataSource}
}

//Monitor Monitor
func (mo *Monitor) Monitor(duration time.Duration) *Snapshot {
	snap := Snapshot{}
	netIn1, netOut1 := mo.netMo()
	messageIn := mo.dataSource.GetMessageIn()
	messageOut := mo.dataSource.GetMessageOut()
	snap.CPUPercent = mo.cpuMo(duration)
	netIn2, netOut2 := mo.netMo()
	snap.MemPercent = mo.memMo()
	snap.AverageNetIn = uint64(float64(netIn2-netIn1) / duration.Seconds())
	snap.AverageNetOut = uint64(float64(netOut2-netOut1) / duration.Seconds())
	snap.TotalConn = mo.dataSource.GetTotalConn()
	snap.MaxConsumer, snap.TotalConsumer = mo.consumerMo()
	snap.TotalTopic = mo.dataSource.GetTotalTopic()
	snap.TotalMessageIn = mo.dataSource.GetMessageIn()
	snap.AverageMessageIn = uint64(float64(snap.TotalMessageIn-messageIn) / duration.Seconds())
	snap.TotalMessageOut = mo.dataSource.GetMessageOut()
	snap.AverageMessageOut = uint64(float64(snap.TotalMessageOut-messageOut) / duration.Seconds())
	snap.Goroutine = uint64(runtime.NumGoroutine())
	snap.Thread = uint64(pprof.Lookup("threadcreate").Count())
	return &snap
}

func (mo *Monitor) cpuMo(duration time.Duration) (percent float64) {
	v, err := cpu.Percent(duration, false)
	if nil == err && len(v) > 0 {
		return v[0]
	}
	return 0
	// return v[0]
	// <-time.After(duration)
	// return 0
}

func (mo *Monitor) consumerMo() (maximum uint64, total uint64) {
	c := mo.dataSource.GetTotalConsumer()
	for _, consumer := range c {
		total += consumer
		if consumer > maximum {
			maximum = consumer
		}
	}
	return maximum, total
}

func (mo *Monitor) netMo() (in uint64, out uint64) {
	c, err := n.IOCounters(false)
	if nil == err && len(c) > 0 {
		return c[0].BytesRecv, c[0].BytesSent
	}
	return 0, 0
}

func (mo *Monitor) memMo() (usedPercent float64) {
	v, _ := mem.VirtualMemory()
	return v.UsedPercent
}
<commit_msg>no message<commit_after>package mo

import (
	"log"
	"runtime"
	"time"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/mem"

	"runtime/pprof"

	n "github.com/shirou/gopsutil/net"
)

//Snapshot Snapshot
type Snapshot struct {
	CPUPercent        float64
	MemPercent        float64
	AverageNetIn      uint64
	AverageNetOut     uint64
	Thread            uint64
	Goroutine         uint64
	TotalConn         uint64
	TotalTopic        uint64
	TotalConsumer     uint64
	MaxConsumer       uint64
	TotalMessageIn    uint64
	AverageMessageIn  uint64
	TotalMessageOut   uint64
	AverageMessageOut uint64
}

//DataSource DataSource
type DataSource interface {
	GetTotalConsumer() map[string]uint64
	GetTotalConn() uint64
	GetTotalTopic() uint64
	GetMessageIn() uint64
	GetMessageOut() uint64
}

//Monitor Monitor
type Monitor struct {
	dataSource DataSource
}

//NewMonitor NewMonitor
func NewMonitor(dataSource DataSource) *Monitor {
	return &Monitor{dataSource: dataSource}
}

//Monitor Monitor
func (mo *Monitor) Monitor(duration time.Duration) *Snapshot {
	snap := Snapshot{}
	netIn1, netOut1 := mo.netMo()
	messageIn := mo.dataSource.GetMessageIn()
	messageOut := mo.dataSource.GetMessageOut()
	snap.CPUPercent = mo.cpuMo(duration)
	netIn2, netOut2 := mo.netMo()
	snap.MemPercent = mo.memMo()
	snap.AverageNetIn = uint64(float64(netIn2-netIn1) / duration.Seconds())
	snap.AverageNetOut = uint64(float64(netOut2-netOut1) / duration.Seconds())
	snap.TotalConn = mo.dataSource.GetTotalConn()
	snap.MaxConsumer, snap.TotalConsumer = mo.consumerMo()
	snap.TotalTopic = mo.dataSource.GetTotalTopic()
	snap.TotalMessageIn = mo.dataSource.GetMessageIn()
	snap.AverageMessageIn = uint64(float64(snap.TotalMessageIn-messageIn) / duration.Seconds())
	snap.TotalMessageOut = mo.dataSource.GetMessageOut()
	snap.AverageMessageOut = uint64(float64(snap.TotalMessageOut-messageOut) / duration.Seconds())
	snap.Goroutine = uint64(runtime.NumGoroutine())
	snap.Thread = uint64(pprof.Lookup("threadcreate").Count())
	return &snap
}

func (mo *Monitor) cpuMo(duration time.Duration) (percent float64) {
	v, err := cpu.Percent(duration, false)
	if nil != err {
		log.Println("cpuMo:", err)
		return 0
	}
	if len(v) == 0 {
		return 0
	}
	return v[0]
}

func (mo *Monitor) consumerMo() (maximum uint64, total uint64) {
	c := mo.dataSource.GetTotalConsumer()
	for _, consumer := range c {
		total += consumer
		if consumer > maximum {
			maximum = consumer
		}
	}
	return maximum, total
}

func (mo *Monitor) netMo() (in uint64, out uint64) {
	c, err := n.IOCounters(false)
	if nil != err {
		log.Println("netMo:", err)
	}
	if len(c) == 0 {
		return 0, 0
	}
	return c[0].BytesRecv, c[0].BytesSent
}

func (mo *Monitor) memMo() (usedPercent float64) {
	v, err := mem.VirtualMemory()
	if nil != err {
		log.Println("memMo:", err)
		return 0
	}
	if nil == v {
		return 0
	}
	return v.UsedPercent
}
<|endoftext|>


<commit_before>package chaos_cluster

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"reflect"
	"testing"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/grafana/metrictank/logger"
	"github.com/grafana/metrictank/stacktest/docker"
	"github.com/grafana/metrictank/stacktest/fakemetrics"
	"github.com/grafana/metrictank/stacktest/grafana"
	"github.com/grafana/metrictank/stacktest/graphite"
	"github.com/grafana/metrictank/stacktest/track"
	"github.com/grafana/metrictank/test"
	log "github.com/sirupsen/logrus"
)

// TODO: cleanup when ctrl-C go test (teardown all containers)

const numPartitions = 12

var tracker *track.Tracker
var fm *fakemetrics.FakeMetrics

func init() {
	formatter := &logger.TextFormatter{}
	formatter.TimestampFormat = "2006-01-02 15:04:05.000"
	log.SetFormatter(formatter)
	log.SetLevel(log.InfoLevel)
}
func TestMain(m *testing.M) {
	flag.Parse()
	if testing.Short() {
		fmt.Println("skipping chaos cluster test in short mode")
		return
	}
	ctx, cancelFunc := context.WithCancel(context.Background())

	fmt.Println("stopping docker-chaos stack should it be running...")
	cmd := exec.CommandContext(ctx, "docker-compose", "down")
	cmd.Dir = test.Path("docker/docker-chaos")
	err := cmd.Start()
	if err != nil {
		log.Fatal(err.Error())
	}
	err = cmd.Wait()
	if err != nil {
		log.Fatal(err.Error())
	}

	fmt.Println("launching docker-chaos stack...")
	cmd = exec.CommandContext(ctx, "docker-compose", "up", "--force-recreate", "-V")
	cmd.Dir = test.Path("docker/docker-chaos")
	cmd.Env = append(cmd.Env, "MT_CLUSTER_MIN_AVAILABLE_SHARDS=12")

	tracker, err = track.NewTracker(cmd, false, false, "launch-stdout", "launch-stderr")
	if err != nil {
		log.Fatal(err.Error())
	}

	err = cmd.Start()
	if err != nil {
		log.Fatal(err.Error())
	}

	retcode := m.Run()
	fm.Close()

	fmt.Println("stopping docker-compose stack...")
	cancelFunc()
	if err := cmd.Wait(); err != nil {
		log.Printf("ERROR: could not cleanly shutdown running docker-compose command: %s", err)
		retcode = 1
	}

	os.Exit(retcode)
}

func TestClusterStartup(t *testing.T) {
	matchers := []track.Matcher{
		{Str: "metrictank0_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank1_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank2_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank3_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank4_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank5_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "grafana.*HTTP Server Listen.*3000"},
	}
	select {
	case <-tracker.Match(matchers):
		fmt.Println("stack now running.")
		fmt.Println("Go to http://localhost:3000 (and login as admin:admin) to see what's going on")
	case <-time.After(time.Second * 40):
		grafana.PostAnnotation("TestClusterStartup:FAIL")
		t.Fatal("timed out while waiting for all metrictank instances to come up")
	}
}

func TestClusterBaseIngestWorkload(t *testing.T) {
	grafana.PostAnnotation("TestClusterBaseIngestWorkload:begin")

	fm = fakemetrics.NewKafka(numPartitions)

	req := graphite.RequestForLocalTestingGraphite("perSecond(metrictank.stats.docker-cluster.*.input.kafka-mdm.metrics_received.counter32)", "-8s")
	resp, ok := graphite.Retry(req, 18, func(resp graphite.Response) bool {
		exp := []string{
			"perSecond(metrictank.stats.docker-cluster.metrictank0.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank1.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank2.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank3.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank4.input.kafka-mdm.metrics_received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank5.input.kafka-mdm.metrics_received.counter32)",
		}
		// avg rate must be 4 (metrics ingested per second by each instance)
		return graphite.ValidateTargets(exp)(resp) && graphite.ValidatorAvgWindowed(8, graphite.Eq(4))(resp)
	})
	if !ok {
		grafana.PostAnnotation("TestClusterBaseIngestWorkload:FAIL")
		t.Fatalf("cluster did not reach a state where each MT instance receives 4 points per second. last response was: %s", spew.Sdump(resp))
	}

	req = graphite.RequestForLocalTestingMT("sum(some.id.of.a.metric.*)", "-16s")
	resp, ok = graphite.Retry(req, 20, graphite.ValidateCorrect(12))
	if !ok {
		grafana.PostAnnotation("TestClusterBaseIngestWorkload:FAIL")
		t.Fatalf("could not query correct result set. sum of 12 series, each valued 1, should result in 12.  last response was: %s", spew.Sdump(resp))
	}
}

func TestQueryWorkload(t *testing.T) {
	grafana.PostAnnotation("TestQueryWorkload:begin")

	results := graphite.CheckMT([]int{6060, 6061, 6062, 6063, 6064, 6065}, "sum(some.id.of.a.metric.*)", "-14s", time.Minute, 6000, graphite.ValidateCorrect(12))

	exp := graphite.CheckResults{
		Valid:   []int{6000},
		Empty:   0,
		Timeout: 0,
		Other:   0,
	}
	if !reflect.DeepEqual(exp, results) {
		grafana.PostAnnotation("TestQueryWorkload:FAIL")
		t.Fatalf("expected only correct results. got %s", spew.Sdump(results))
	}
}

// TestIsolateOneInstance tests what happens during the isolation of one instance, when min-available-shards is 12
// this should happen:
// at all times, all queries to all of the remaining nodes should be successful
// since they have at least 1 instance running for each shard.
// the isolated shard should either return correct replies, or errors (in two cases: when it marks any shards as down,
// but also before it does, but fails to get data via clustered requests from peers)
func TestIsolateOneInstance(t *testing.T) {
	grafana.PostAnnotation("TestIsolateOneInstance:begin")
	numReqMt4 := 1200

	mt4ResultsChan := make(chan graphite.CheckResults, 1)
	otherResultsChan := make(chan graphite.CheckResults, 1)

	go func() {
		mt4ResultsChan <- graphite.CheckMT([]int{6064}, "sum(some.id.of.a.metric.*)", "-15s", time.Minute, numReqMt4, graphite.ValidateCorrect(12), graphite.ValidateCode(503))
	}()
	go func() {
		otherResultsChan <- graphite.CheckMT([]int{6060, 6061, 6062, 6063, 6065}, "sum(some.id.of.a.metric.*)", "-15s", time.Minute, 6000, graphite.ValidateCorrect(12))
	}()

	// now go ahead and isolate for 30s
	docker.Isolate([]string{"metrictank4"}, []string{"metrictank0", "metrictank1", "metrictank2", "metrictank3", "metrictank5"}, "30s")

	// collect results of the minute long experiment
	mt4Results := <-mt4ResultsChan
	otherResults := <-otherResultsChan

	// validate results of isolated node
	if mt4Results.Valid[0]+mt4Results.Valid[1] != numReqMt4 {
		t.Fatalf("expected mt4 to return either correct or erroring responses (total %d). got %s", numReqMt4, spew.Sdump(mt4Results))
	}
	if mt4Results.Valid[1] < numReqMt4*30/100 {
		// the instance is completely down for 30s of the 60s experiment run, but we allow some slack
		t.Fatalf("expected at least 30%% of all mt4 results to succeed. did %d queries. got %s", numReqMt4, spew.Sdump(mt4Results))
	}

	// validate results of other cluster nodes
	exp := graphite.CheckResults{
		Valid:   []int{6000},
		Empty:   0,
		Timeout: 0,
		Other:   0,
	}
	if !reflect.DeepEqual(exp, otherResults) {
		grafana.PostAnnotation("TestIsolateOneInstance:FAIL")
		t.Fatalf("expected only correct results for all cluster nodes. got %s", spew.Sdump(otherResults))
	}
}

func TestHang(t *testing.T) {
	grafana.PostAnnotation("TestHang:begin")
	t.Log("whatever happens, keep hanging for now, so that we can query grafana dashboards still")
	var ch chan struct{}
	<-ch
}

// maybe useful in the future, test also clean exit and rejoin like so:
//stop("metrictank4")
//time.AfterFunc(30*time.Second, func() {
//	start("metrictank4")
//})
<commit_msg>bugfix: query for per-type ingest metrics<commit_after>package chaos_cluster

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"reflect"
	"testing"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/grafana/metrictank/logger"
	"github.com/grafana/metrictank/stacktest/docker"
	"github.com/grafana/metrictank/stacktest/fakemetrics"
	"github.com/grafana/metrictank/stacktest/grafana"
	"github.com/grafana/metrictank/stacktest/graphite"
	"github.com/grafana/metrictank/stacktest/track"
	"github.com/grafana/metrictank/test"
	log "github.com/sirupsen/logrus"
)

// TODO: cleanup when ctrl-C go test (teardown all containers)

const numPartitions = 12

var tracker *track.Tracker
var fm *fakemetrics.FakeMetrics

func init() {
	formatter := &logger.TextFormatter{}
	formatter.TimestampFormat = "2006-01-02 15:04:05.000"
	log.SetFormatter(formatter)
	log.SetLevel(log.InfoLevel)
}
func TestMain(m *testing.M) {
	flag.Parse()
	if testing.Short() {
		fmt.Println("skipping chaos cluster test in short mode")
		return
	}
	ctx, cancelFunc := context.WithCancel(context.Background())

	fmt.Println("stopping docker-chaos stack should it be running...")
	cmd := exec.CommandContext(ctx, "docker-compose", "down")
	cmd.Dir = test.Path("docker/docker-chaos")
	err := cmd.Start()
	if err != nil {
		log.Fatal(err.Error())
	}
	err = cmd.Wait()
	if err != nil {
		log.Fatal(err.Error())
	}

	fmt.Println("launching docker-chaos stack...")
	cmd = exec.CommandContext(ctx, "docker-compose", "up", "--force-recreate", "-V")
	cmd.Dir = test.Path("docker/docker-chaos")
	cmd.Env = append(cmd.Env, "MT_CLUSTER_MIN_AVAILABLE_SHARDS=12")

	tracker, err = track.NewTracker(cmd, false, false, "launch-stdout", "launch-stderr")
	if err != nil {
		log.Fatal(err.Error())
	}

	err = cmd.Start()
	if err != nil {
		log.Fatal(err.Error())
	}

	retcode := m.Run()
	fm.Close()

	fmt.Println("stopping docker-compose stack...")
	cancelFunc()
	if err := cmd.Wait(); err != nil {
		log.Printf("ERROR: could not cleanly shutdown running docker-compose command: %s", err)
		retcode = 1
	}

	os.Exit(retcode)
}

func TestClusterStartup(t *testing.T) {
	matchers := []track.Matcher{
		{Str: "metrictank0_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank1_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank2_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank3_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank4_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank5_1.*metricIndex initialized.*starting data consumption$"},
		{Str: "grafana.*HTTP Server Listen.*3000"},
	}
	select {
	case <-tracker.Match(matchers):
		fmt.Println("stack now running.")
		fmt.Println("Go to http://localhost:3000 (and login as admin:admin) to see what's going on")
	case <-time.After(time.Second * 40):
		grafana.PostAnnotation("TestClusterStartup:FAIL")
		t.Fatal("timed out while waiting for all metrictank instances to come up")
	}
}

func TestClusterBaseIngestWorkload(t *testing.T) {
	grafana.PostAnnotation("TestClusterBaseIngestWorkload:begin")

	// generate exactly numPartitions metrics, numbered 0..numPartitions where each metric goes to the partition of the same number
	// each partition is consumed by 2 instances, and each instance consumes 4 partitions thus 4 metrics/s
	fm = fakemetrics.NewKafka(numPartitions)

	req := graphite.RequestForLocalTestingGraphite("perSecond(metrictank.stats.docker-cluster.*.input.kafka-mdm.metricdata.received.counter32)", "-8s")
	resp, ok := graphite.Retry(req, 18, func(resp graphite.Response) bool {
		exp := []string{
			"perSecond(metrictank.stats.docker-cluster.metrictank0.input.kafka-mdm.metricdata.received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank1.input.kafka-mdm.metricdata.received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank2.input.kafka-mdm.metricdata.received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank3.input.kafka-mdm.metricdata.received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank4.input.kafka-mdm.metricdata.received.counter32)",
			"perSecond(metrictank.stats.docker-cluster.metrictank5.input.kafka-mdm.metricdata.received.counter32)",
		}
		// avg rate must be 4 (metrics ingested per second by each instance)
		return graphite.ValidateTargets(exp)(resp) && graphite.ValidatorAvgWindowed(8, graphite.Eq(4))(resp)
	})
	if !ok {
		grafana.PostAnnotation("TestClusterBaseIngestWorkload:FAIL")
		t.Fatalf("cluster did not reach a state where each MT instance receives 4 points per second. last response was: %s", spew.Sdump(resp))
	}

	req = graphite.RequestForLocalTestingMT("sum(some.id.of.a.metric.*)", "-16s")
	resp, ok = graphite.Retry(req, 20, graphite.ValidateCorrect(12))
	if !ok {
		grafana.PostAnnotation("TestClusterBaseIngestWorkload:FAIL")
		t.Fatalf("could not query correct result set. sum of 12 series, each valued 1, should result in 12.  last response was: %s", spew.Sdump(resp))
	}
}

func TestQueryWorkload(t *testing.T) {
	grafana.PostAnnotation("TestQueryWorkload:begin")

	results := graphite.CheckMT([]int{6060, 6061, 6062, 6063, 6064, 6065}, "sum(some.id.of.a.metric.*)", "-14s", time.Minute, 6000, graphite.ValidateCorrect(12))

	exp := graphite.CheckResults{
		Valid:   []int{6000},
		Empty:   0,
		Timeout: 0,
		Other:   0,
	}
	if !reflect.DeepEqual(exp, results) {
		grafana.PostAnnotation("TestQueryWorkload:FAIL")
		t.Fatalf("expected only correct results. got %s", spew.Sdump(results))
	}
}

// TestIsolateOneInstance tests what happens during the isolation of one instance, when min-available-shards is 12
// this should happen:
// at all times, all queries to all of the remaining nodes should be successful
// since they have at least 1 instance running for each shard.
// the isolated shard should either return correct replies, or errors (in two cases: when it marks any shards as down,
// but also before it does, but fails to get data via clustered requests from peers)
func TestIsolateOneInstance(t *testing.T) {
	grafana.PostAnnotation("TestIsolateOneInstance:begin")
	numReqMt4 := 1200

	mt4ResultsChan := make(chan graphite.CheckResults, 1)
	otherResultsChan := make(chan graphite.CheckResults, 1)

	go func() {
		mt4ResultsChan <- graphite.CheckMT([]int{6064}, "sum(some.id.of.a.metric.*)", "-15s", time.Minute, numReqMt4, graphite.ValidateCorrect(12), graphite.ValidateCode(503))
	}()
	go func() {
		otherResultsChan <- graphite.CheckMT([]int{6060, 6061, 6062, 6063, 6065}, "sum(some.id.of.a.metric.*)", "-15s", time.Minute, 6000, graphite.ValidateCorrect(12))
	}()

	// now go ahead and isolate for 30s
	docker.Isolate([]string{"metrictank4"}, []string{"metrictank0", "metrictank1", "metrictank2", "metrictank3", "metrictank5"}, "30s")

	// collect results of the minute long experiment
	mt4Results := <-mt4ResultsChan
	otherResults := <-otherResultsChan

	// validate results of isolated node
	if mt4Results.Valid[0]+mt4Results.Valid[1] != numReqMt4 {
		t.Fatalf("expected mt4 to return either correct or erroring responses (total %d). got %s", numReqMt4, spew.Sdump(mt4Results))
	}
	if mt4Results.Valid[1] < numReqMt4*30/100 {
		// the instance is completely down for 30s of the 60s experiment run, but we allow some slack
		t.Fatalf("expected at least 30%% of all mt4 results to succeed. did %d queries. got %s", numReqMt4, spew.Sdump(mt4Results))
	}

	// validate results of other cluster nodes
	exp := graphite.CheckResults{
		Valid:   []int{6000},
		Empty:   0,
		Timeout: 0,
		Other:   0,
	}
	if !reflect.DeepEqual(exp, otherResults) {
		grafana.PostAnnotation("TestIsolateOneInstance:FAIL")
		t.Fatalf("expected only correct results for all cluster nodes. got %s", spew.Sdump(otherResults))
	}
}

func TestHang(t *testing.T) {
	grafana.PostAnnotation("TestHang:begin")
	t.Log("whatever happens, keep hanging for now, so that we can query grafana dashboards still")
	var ch chan struct{}
	<-ch
}

// maybe useful in the future, test also clean exit and rejoin like so:
//stop("metrictank4")
//time.AfterFunc(30*time.Second, func() {
//	start("metrictank4")
//})
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package flag

import (
	goflag "flag"
	"strings"

	"github.com/golang/glog"
	"github.com/spf13/pflag"
)

// WordSepNormalizeFunc changes all flags that contain "_" separators
func WordSepNormalizeFunc(f *pflag.FlagSet, name string) pflag.NormalizedName {
	if strings.Contains(name, "_") {
		return pflag.NormalizedName(strings.Replace(name, "_", "-", -1))
	}
	return pflag.NormalizedName(name)
}

// WarnWordSepNormalizeFunc changes and warns for flags that contain "_" separators
func WarnWordSepNormalizeFunc(f *pflag.FlagSet, name string) pflag.NormalizedName {
	if strings.Contains(name, "_") {
		nname := strings.Replace(name, "_", "-", -1)
		glog.Warningf("%s is DEPRECATED and will be removed in a future version. Use %s instead.", name, nname)

		return pflag.NormalizedName(nname)
	}
	return pflag.NormalizedName(name)
}

// InitFlags normalizes, parses, then logs the command line flags
func InitFlags() {
	pflag.CommandLine.SetNormalizeFunc(WordSepNormalizeFunc)
	pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
	pflag.Parse()
	pflag.VisitAll(func(flag *pflag.Flag) {
		glog.Infof("FLAG: --%s=%q", flag.Name, flag.Value)
	})
}
<commit_msg>add level for print flags<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package flag

import (
	goflag "flag"
	"strings"

	"github.com/golang/glog"
	"github.com/spf13/pflag"
)

// WordSepNormalizeFunc changes all flags that contain "_" separators
func WordSepNormalizeFunc(f *pflag.FlagSet, name string) pflag.NormalizedName {
	if strings.Contains(name, "_") {
		return pflag.NormalizedName(strings.Replace(name, "_", "-", -1))
	}
	return pflag.NormalizedName(name)
}

// WarnWordSepNormalizeFunc changes and warns for flags that contain "_" separators
func WarnWordSepNormalizeFunc(f *pflag.FlagSet, name string) pflag.NormalizedName {
	if strings.Contains(name, "_") {
		nname := strings.Replace(name, "_", "-", -1)
		glog.Warningf("%s is DEPRECATED and will be removed in a future version. Use %s instead.", name, nname)

		return pflag.NormalizedName(nname)
	}
	return pflag.NormalizedName(name)
}

// InitFlags normalizes, parses, then logs the command line flags
func InitFlags() {
	pflag.CommandLine.SetNormalizeFunc(WordSepNormalizeFunc)
	pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
	pflag.Parse()
	pflag.VisitAll(func(flag *pflag.Flag) {
		glog.V(4).Infof("FLAG: --%s=%q", flag.Name, flag.Value)
	})
}
<|endoftext|>


<commit_before>// Copyright 2013 Prometheus Team
// 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 retrieval

import (
	"encoding/json"
	"fmt"
	"github.com/matttproud/golang_instrumentation/metrics"
	"github.com/matttproud/prometheus/model"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"
)

// The state of the given Target.
type TargetState int

const (
	// The Target has not been seen; we know nothing about it, except that it is
	// on our docket for examination.
	UNKNOWN TargetState = iota
	// The Target has been found and successfully queried.
	ALIVE
	// The Target was either historically found or not found and then determined
	// to be unhealthy by either not responding or disappearing.
	UNREACHABLE
)

// A healthReporter is a type that can provide insight into its health state.
//
// It mainly exists for testability reasons to decouple the scheduler behaviors
// from fully-fledged Target and other types.
type healthReporter interface {
	// Report the last-known health state for this target.
	State() TargetState
}

// A Target represents an endpoint that should be interrogated for metrics.
//
// The protocol described by this type will likely change in future iterations,
// as it offers no good support for aggregated targets and fan out.  Thusly,
// it is likely that the current Target and target uses will be
// wrapped with some resolver type.
//
// For the future, the Target protocol will abstract away the exact means that
// metrics are retrieved and deserialized from the given instance to which it
// refers.
type Target interface {
	// Retrieve values from this target.
	//
	// earliest refers to the soonest available opportunity to reschedule the
	// target for a future retrieval.  It is up to the underlying scheduler type,
	// alluded to in the scheduledFor function, to use this as it wants to.  The
	// current use case is to create a common batching time for scraping multiple
	// Targets in the future through the TargetPool.
	Scrape(earliest time.Time, results chan Result) error
	// Fulfill the healthReporter interface.
	State() TargetState
	// Report the soonest time at which this Target may be scheduled for
	// retrieval.  This value needn't convey that the operation occurs at this
	// time, but it should occur no sooner than it.
	//
	// Right now, this is used as the sorting key in TargetPool.
	scheduledFor() time.Time
	// The address to which the Target corresponds.  Out of all of the available
	// points in this interface, this one is the best candidate to change given
	// the ways to express the endpoint.
	Address() string
	// How frequently queries occur.
	Interval() time.Duration
}

// target is a Target that refers to a singular HTTP or HTTPS endpoint.
type target struct {
	// scheduler provides the scheduling strategy that is used to formulate what
	// is returned in Target.scheduledFor.
	scheduler scheduler
	state     TargetState

	address string
	// What is the deadline for the HTTP or HTTPS against this endpoint.
	Deadline time.Duration
	// Any base labels that are added to this target and its metrics.
	BaseLabels model.LabelSet

	// XXX: Move this to a field with the target manager initialization instead of here.
	interval time.Duration
}

// Furnish a reasonably configured target for querying.
func NewTarget(address string, interval, deadline time.Duration, baseLabels model.LabelSet) Target {
	target := &target{
		address:    address,
		Deadline:   deadline,
		interval:   interval,
		BaseLabels: baseLabels,
	}

	scheduler := &healthScheduler{
		target: target,
	}
	target.scheduler = scheduler

	return target
}

type Result struct {
	Err     error
	Samples []model.Sample
	Target  Target
}

func (t *target) Scrape(earliest time.Time, results chan Result) (err error) {
	result := Result{}

	defer func() {
		futureState := t.state

		switch err {
		case nil:
			futureState = ALIVE
		default:
			futureState = UNREACHABLE
		}

		t.scheduler.Reschedule(earliest, futureState)

		result.Err = err
		results <- result
	}()

	done := make(chan bool)

	request := func() {
		ti := time.Now()
		resp, err := http.Get(t.Address())
		if err != nil {
			return
		}

		defer resp.Body.Close()

		raw, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return
		}

		intermediate := make(map[string]interface{})
		err = json.Unmarshal(raw, &intermediate)
		if err != nil {
			return
		}

		baseLabels := map[string]string{"instance": t.Address()}

		for name, v := range intermediate {
			asMap, ok := v.(map[string]interface{})

			if !ok {
				continue
			}

			switch asMap["type"] {
			case "counter":
				m := model.Metric{}
				m["name"] = model.LabelValue(name)
				asFloat, ok := asMap["value"].(float64)
				if !ok {
					continue
				}

				s := model.Sample{
					Metric:    m,
					Value:     model.SampleValue(asFloat),
					Timestamp: ti,
				}

				for baseK, baseV := range baseLabels {
					m[model.LabelName(baseK)] = model.LabelValue(baseV)
				}

				result.Samples = append(result.Samples, s)
			case "histogram":
				values, ok := asMap["value"].(map[string]interface{})
				if !ok {
					continue
				}

				for p, pValue := range values {
					asString, ok := pValue.(string)
					if !ok {
						continue
					}

					float, err := strconv.ParseFloat(asString, 64)
					if err != nil {
						continue
					}

					m := model.Metric{}
					m["name"] = model.LabelValue(name)
					m["percentile"] = model.LabelValue(p)

					s := model.Sample{
						Metric:    m,
						Value:     model.SampleValue(float),
						Timestamp: ti,
					}

					for baseK, baseV := range baseLabels {
						m[model.LabelName(baseK)] = model.LabelValue(baseV)
					}

					result.Samples = append(result.Samples, s)
				}
			}
		}

		done <- true
	}

	accumulator := func(d time.Duration) {
		ms := float64(d) / float64(time.Millisecond)
		if err == nil {
			scrapeLatencyHealthy.Add(ms)
		} else {
			scrapeLatencyUnhealthy.Add(ms)
		}
	}

	go metrics.InstrumentCall(request, accumulator)

	select {
	case <-done:
		break
	case <-time.After(t.Deadline):
		err = fmt.Errorf("Target %s exceeded %s deadline.", t, t.Deadline)
	}

	return
}

func (t target) State() TargetState {
	return t.state
}

func (t target) scheduledFor() time.Time {
	return t.scheduler.ScheduledFor()
}

func (t target) Address() string {
	return t.address
}

func (t target) Interval() time.Duration {
	return t.interval
}
<commit_msg>Add support for configured job base labels.<commit_after>// Copyright 2013 Prometheus Team
// 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 retrieval

import (
	"encoding/json"
	"fmt"
	"github.com/matttproud/golang_instrumentation/metrics"
	"github.com/matttproud/prometheus/model"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"
)

// The state of the given Target.
type TargetState int

const (
	// The Target has not been seen; we know nothing about it, except that it is
	// on our docket for examination.
	UNKNOWN TargetState = iota
	// The Target has been found and successfully queried.
	ALIVE
	// The Target was either historically found or not found and then determined
	// to be unhealthy by either not responding or disappearing.
	UNREACHABLE
)

// A healthReporter is a type that can provide insight into its health state.
//
// It mainly exists for testability reasons to decouple the scheduler behaviors
// from fully-fledged Target and other types.
type healthReporter interface {
	// Report the last-known health state for this target.
	State() TargetState
}

// A Target represents an endpoint that should be interrogated for metrics.
//
// The protocol described by this type will likely change in future iterations,
// as it offers no good support for aggregated targets and fan out.  Thusly,
// it is likely that the current Target and target uses will be
// wrapped with some resolver type.
//
// For the future, the Target protocol will abstract away the exact means that
// metrics are retrieved and deserialized from the given instance to which it
// refers.
type Target interface {
	// Retrieve values from this target.
	//
	// earliest refers to the soonest available opportunity to reschedule the
	// target for a future retrieval.  It is up to the underlying scheduler type,
	// alluded to in the scheduledFor function, to use this as it wants to.  The
	// current use case is to create a common batching time for scraping multiple
	// Targets in the future through the TargetPool.
	Scrape(earliest time.Time, results chan Result) error
	// Fulfill the healthReporter interface.
	State() TargetState
	// Report the soonest time at which this Target may be scheduled for
	// retrieval.  This value needn't convey that the operation occurs at this
	// time, but it should occur no sooner than it.
	//
	// Right now, this is used as the sorting key in TargetPool.
	scheduledFor() time.Time
	// The address to which the Target corresponds.  Out of all of the available
	// points in this interface, this one is the best candidate to change given
	// the ways to express the endpoint.
	Address() string
	// How frequently queries occur.
	Interval() time.Duration
}

// target is a Target that refers to a singular HTTP or HTTPS endpoint.
type target struct {
	// scheduler provides the scheduling strategy that is used to formulate what
	// is returned in Target.scheduledFor.
	scheduler scheduler
	state     TargetState

	address string
	// What is the deadline for the HTTP or HTTPS against this endpoint.
	Deadline time.Duration
	// Any base labels that are added to this target and its metrics.
	BaseLabels model.LabelSet

	// XXX: Move this to a field with the target manager initialization instead of here.
	interval time.Duration
}

// Furnish a reasonably configured target for querying.
func NewTarget(address string, interval, deadline time.Duration, baseLabels model.LabelSet) Target {
	target := &target{
		address:    address,
		Deadline:   deadline,
		interval:   interval,
		BaseLabels: baseLabels,
	}

	scheduler := &healthScheduler{
		target: target,
	}
	target.scheduler = scheduler

	return target
}

type Result struct {
	Err     error
	Samples []model.Sample
	Target  Target
}

func (t *target) Scrape(earliest time.Time, results chan Result) (err error) {
	result := Result{}

	defer func() {
		futureState := t.state

		switch err {
		case nil:
			futureState = ALIVE
		default:
			futureState = UNREACHABLE
		}

		t.scheduler.Reschedule(earliest, futureState)

		result.Err = err
		results <- result
	}()

	done := make(chan bool)

	request := func() {
		ti := time.Now()
		resp, err := http.Get(t.Address())
		if err != nil {
			return
		}

		defer resp.Body.Close()

		raw, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return
		}

		intermediate := make(map[string]interface{})
		err = json.Unmarshal(raw, &intermediate)
		if err != nil {
			return
		}

		baseLabels := model.LabelSet{"instance": model.LabelValue(t.Address())}
		for baseK, baseV := range t.BaseLabels {
			baseLabels[baseK] = baseV
		}

		for name, v := range intermediate {
			asMap, ok := v.(map[string]interface{})

			if !ok {
				continue
			}

			switch asMap["type"] {
			case "counter":
				m := model.Metric{}
				m["name"] = model.LabelValue(name)
				asFloat, ok := asMap["value"].(float64)
				if !ok {
					continue
				}

				s := model.Sample{
					Metric:    m,
					Value:     model.SampleValue(asFloat),
					Timestamp: ti,
				}

				for baseK, baseV := range baseLabels {
					m[baseK] = baseV
				}

				result.Samples = append(result.Samples, s)
			case "histogram":
				values, ok := asMap["value"].(map[string]interface{})
				if !ok {
					continue
				}

				for p, pValue := range values {
					asString, ok := pValue.(string)
					if !ok {
						continue
					}

					float, err := strconv.ParseFloat(asString, 64)
					if err != nil {
						continue
					}

					m := model.Metric{}
					m["name"] = model.LabelValue(name)
					m["percentile"] = model.LabelValue(p)

					s := model.Sample{
						Metric:    m,
						Value:     model.SampleValue(float),
						Timestamp: ti,
					}

					for baseK, baseV := range baseLabels {
						m[baseK] = baseV
					}

					result.Samples = append(result.Samples, s)
				}
			}
		}

		done <- true
	}

	accumulator := func(d time.Duration) {
		ms := float64(d) / float64(time.Millisecond)
		if err == nil {
			scrapeLatencyHealthy.Add(ms)
		} else {
			scrapeLatencyUnhealthy.Add(ms)
		}
	}

	go metrics.InstrumentCall(request, accumulator)

	select {
	case <-done:
		break
	case <-time.After(t.Deadline):
		err = fmt.Errorf("Target %s exceeded %s deadline.", t, t.Deadline)
	}

	return
}

func (t target) State() TargetState {
	return t.state
}

func (t target) scheduledFor() time.Time {
	return t.scheduler.ScheduledFor()
}

func (t target) Address() string {
	return t.address
}

func (t target) Interval() time.Duration {
	return t.interval
}
<|endoftext|>