<commit_before>package localfs

import (
	"context"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/pkg/errors"

	"github.com/kopia/kopia/fs"
)

const (
	numEntriesToRead   = 100 // number of directory entries to read in one shot
	dirListingPrefetch = 200 // number of directory items to os.Lstat() in advance
)

type filesystemEntry struct {
	name       string
	size       int64
	mtimeNanos int64
	mode       os.FileMode
	owner      fs.OwnerInfo
	device     fs.DeviceInfo

	parentDir string
}

func (e *filesystemEntry) Name() string {
	return e.name
}

func (e *filesystemEntry) IsDir() bool {
	return e.mode.IsDir()
}

func (e *filesystemEntry) Mode() os.FileMode {
	return e.mode
}

func (e *filesystemEntry) Size() int64 {
	return e.size
}

func (e *filesystemEntry) ModTime() time.Time {
	return time.Unix(0, e.mtimeNanos)
}

func (e *filesystemEntry) Sys() interface{} {
	return nil
}

func (e *filesystemEntry) fullPath() string {
	return filepath.Join(e.parentDir, e.Name())
}

func (e *filesystemEntry) Owner() fs.OwnerInfo {
	return e.owner
}

func (e *filesystemEntry) Device() fs.DeviceInfo {
	return e.device
}

func (e *filesystemEntry) LocalFilesystemPath() string {
	return e.fullPath()
}

var _ os.FileInfo = (*filesystemEntry)(nil)

func newEntry(fi os.FileInfo, parentDir string) filesystemEntry {
	return filesystemEntry{
		TrimShallowSuffix(fi.Name()),
		fi.Size(),
		fi.ModTime().UnixNano(),
		fi.Mode(),
		platformSpecificOwnerInfo(fi),
		platformSpecificDeviceInfo(fi),
		parentDir,
	}
}

type filesystemDirectory struct {
	filesystemEntry
}

type filesystemSymlink struct {
	filesystemEntry
}

type filesystemFile struct {
	filesystemEntry
}

type filesystemErrorEntry struct {
	filesystemEntry
	err error
}

func (fsd *filesystemDirectory) Size() int64 {
	// force directory size to always be zero
	return 0
}

func (fsd *filesystemDirectory) Child(ctx context.Context, name string) (fs.Entry, error) {
	fullPath := fsd.fullPath()

	st, err := os.Lstat(filepath.Join(fullPath, name))
	if err != nil {
		if os.IsNotExist(err) {
			return nil, fs.ErrEntryNotFound
		}

		return nil, errors.Wrap(err, "unable to get child")
	}

	return entryFromChildFileInfo(st, fullPath), nil
}

type entryWithError struct {
	entry fs.Entry
	err   error
}

func (fsd *filesystemDirectory) Readdir(ctx context.Context) (fs.Entries, error) {
	fullPath := fsd.fullPath()

	f, direrr := os.Open(fullPath) //nolint:gosec
	if direrr != nil {
		return nil, errors.Wrap(direrr, "unable to read directory")
	}
	defer f.Close() //nolint:errcheck,gosec

	// start feeding directory entry names to namesCh
	namesCh := make(chan string, dirListingPrefetch)

	var readDirErr error

	go func() {
		defer close(namesCh)

		for {
			names, err := f.Readdirnames(numEntriesToRead)
			for _, name := range names {
				namesCh <- name
			}

			if err == nil {
				continue
			}

			if errors.Is(err, io.EOF) {
				break
			}

			readDirErr = err

			break
		}
	}()

	entriesCh := make(chan entryWithError, dirListingPrefetch)

	var workersWG sync.WaitGroup

	// launch N workers to os.Lstat() each name in parallel and push to entriesCh
	workers := 16
	for i := 0; i < workers; i++ {
		workersWG.Add(1)

		go func() {
			defer workersWG.Done()

			for n := range namesCh {
				fi, staterr := os.Lstat(fullPath + "/" + n)

				switch {
				case os.IsNotExist(staterr):
					// lost the race - ignore.
					continue
				case staterr != nil:
					entriesCh <- entryWithError{err: errors.Errorf("unable to stat directory entry %q: %v", n, staterr)}
					continue
				}

				entriesCh <- entryWithError{entry: entryFromChildFileInfo(fi, fullPath)}
			}
		}()
	}

	// close entriesCh channel when all workers terminate
	go func() {
		workersWG.Wait()
		close(entriesCh)
	}()

	// drain the entriesCh into a slice and sort it
	var entries fs.Entries

	for e := range entriesCh {
		if e.err != nil {
			// only return the first error
			if readDirErr == nil {
				readDirErr = e.err
			}

			continue
		}

		entries = append(entries, e.entry)
	}

	entries.Sort()

	// return any error encountered when listing or reading the directory
	return entries, readDirErr
}

type fileWithMetadata struct {
	*os.File
}

func (f *fileWithMetadata) Entry() (fs.Entry, error) {
	fi, err := f.Stat()
	if err != nil {
		return nil, errors.Wrap(err, "unable to stat() local file")
	}

	return &filesystemFile{newEntry(fi, filepath.Dir(f.Name()))}, nil
}

func (fsf *filesystemFile) Open(ctx context.Context) (fs.Reader, error) {
	f, err := os.Open(fsf.fullPath())
	if err != nil {
		return nil, errors.Wrap(err, "unable to open local file")
	}

	return &fileWithMetadata{f}, nil
}

func (fsl *filesystemSymlink) Readlink(ctx context.Context) (string, error) {
	// nolint:wrapcheck
	return os.Readlink(fsl.fullPath())
}

func (e *filesystemErrorEntry) ErrorInfo() error {
	return e.err
}

// NewEntry returns fs.Entry for the specified path, the result will be one of supported entry types: fs.File, fs.Directory, fs.Symlink
// or fs.UnsupportedEntry.
func NewEntry(path string) (fs.Entry, error) {
	fi, err := os.Lstat(path)
	if err != nil {
		return nil, errors.Wrap(err, "unable to determine entry type")
	}

	return entryFromChildFileInfo(fi, filepath.Dir(path)), nil
}

// Directory returns fs.Directory for the specified path.
func Directory(path string) (fs.Directory, error) {
	e, err := NewEntry(path)
	if err != nil {
		return nil, err
	}

	if d, ok := e.(fs.Directory); ok {
		return d, nil
	}

	return nil, errors.Errorf("not a directory: %v", path)
}

func entryFromChildFileInfo(fi os.FileInfo, parentDir string) fs.Entry {
	isplaceholder := strings.HasSuffix(fi.Name(), ShallowEntrySuffix)
	maskedmode := fi.Mode() & os.ModeType

	switch {
	case maskedmode == os.ModeDir && !isplaceholder:
		return &filesystemDirectory{newEntry(fi, parentDir)}

	case maskedmode == os.ModeDir && isplaceholder:
		return &shallowFilesystemDirectory{newEntry(fi, parentDir)}

	case maskedmode == os.ModeSymlink && !isplaceholder:
		return &filesystemSymlink{newEntry(fi, parentDir)}

	case maskedmode == 0 && !isplaceholder:
		return &filesystemFile{newEntry(fi, parentDir)}

	case maskedmode == 0 && isplaceholder:
		return &shallowFilesystemFile{newEntry(fi, parentDir)}

	default:
		return &filesystemErrorEntry{newEntry(fi, parentDir), fs.ErrUnknown}
	}
}

var (
	_ fs.Directory  = &filesystemDirectory{}
	_ fs.File       = &filesystemFile{}
	_ fs.Symlink    = &filesystemSymlink{}
	_ fs.ErrorEntry = &filesystemErrorEntry{}
)
<commit_msg>localfs: reduce memory usage when scanning short directories (#1343)<commit_after>package localfs

import (
	"context"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/pkg/errors"

	"github.com/kopia/kopia/fs"
)

const (
	numEntriesToReadFirst    = 100 // number of directory entries to read in the first batch before parallelism kicks in.
	numEntriesToRead         = 100 // number of directory entries to read in one shot
	dirListingPrefetch       = 200 // number of directory items to os.Lstat() in advance
	paralellelStatGoroutines = 4   // how many goroutines to use when Lstat() on large directory
)

type filesystemEntry struct {
	name       string
	size       int64
	mtimeNanos int64
	mode       os.FileMode
	owner      fs.OwnerInfo
	device     fs.DeviceInfo

	parentDir string
}

func (e *filesystemEntry) Name() string {
	return e.name
}

func (e *filesystemEntry) IsDir() bool {
	return e.mode.IsDir()
}

func (e *filesystemEntry) Mode() os.FileMode {
	return e.mode
}

func (e *filesystemEntry) Size() int64 {
	return e.size
}

func (e *filesystemEntry) ModTime() time.Time {
	return time.Unix(0, e.mtimeNanos)
}

func (e *filesystemEntry) Sys() interface{} {
	return nil
}

func (e *filesystemEntry) fullPath() string {
	return filepath.Join(e.parentDir, e.Name())
}

func (e *filesystemEntry) Owner() fs.OwnerInfo {
	return e.owner
}

func (e *filesystemEntry) Device() fs.DeviceInfo {
	return e.device
}

func (e *filesystemEntry) LocalFilesystemPath() string {
	return e.fullPath()
}

var _ os.FileInfo = (*filesystemEntry)(nil)

func newEntry(fi os.FileInfo, parentDir string) filesystemEntry {
	return filesystemEntry{
		TrimShallowSuffix(fi.Name()),
		fi.Size(),
		fi.ModTime().UnixNano(),
		fi.Mode(),
		platformSpecificOwnerInfo(fi),
		platformSpecificDeviceInfo(fi),
		parentDir,
	}
}

type filesystemDirectory struct {
	filesystemEntry
}

type filesystemSymlink struct {
	filesystemEntry
}

type filesystemFile struct {
	filesystemEntry
}

type filesystemErrorEntry struct {
	filesystemEntry
	err error
}

func (fsd *filesystemDirectory) Size() int64 {
	// force directory size to always be zero
	return 0
}

func (fsd *filesystemDirectory) Child(ctx context.Context, name string) (fs.Entry, error) {
	fullPath := fsd.fullPath()

	st, err := os.Lstat(filepath.Join(fullPath, name))
	if err != nil {
		if os.IsNotExist(err) {
			return nil, fs.ErrEntryNotFound
		}

		return nil, errors.Wrap(err, "unable to get child")
	}

	return entryFromDirEntry(st, fullPath), nil
}

type entryWithError struct {
	entry fs.Entry
	err   error
}

func toDirEntryOrNil(basename, dirPath string) (fs.Entry, error) {
	fi, err := os.Lstat(dirPath + "/" + basename)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, errors.Wrap(err, "error reading directory")
	}

	return entryFromDirEntry(fi, dirPath), nil
}

func (fsd *filesystemDirectory) Readdir(ctx context.Context) (fs.Entries, error) {
	fullPath := fsd.fullPath()

	f, direrr := os.Open(fullPath) //nolint:gosec
	if direrr != nil {
		return nil, errors.Wrap(direrr, "unable to read directory")
	}
	defer f.Close() //nolint:errcheck,gosec

	var entries fs.Entries

	// read first batch of directory entries using Readdir() before parallelization.
	firstBatch, firstBatchErr := f.Readdirnames(numEntriesToReadFirst)
	if firstBatchErr != nil && !errors.Is(firstBatchErr, io.EOF) {
		return nil, errors.Wrap(firstBatchErr, "unable to read directory entries")
	}

	for _, de := range firstBatch {
		e, err := toDirEntryOrNil(de, fullPath)
		if err != nil {
			return nil, errors.Wrap(err, "error reading entry")
		}

		if e != nil {
			entries = append(entries, e)
		}
	}

	// first batch was complete with EOF, we're done here.
	if errors.Is(firstBatchErr, io.EOF) {
		entries.Sort()

		return entries, nil
	}

	// first batch was shorter than expected, perform another read to make sure we get EOF.
	if len(firstBatch) < numEntriesToRead {
		secondBatch, secondBatchErr := f.Readdirnames(numEntriesToRead)
		if secondBatchErr != nil && !errors.Is(secondBatchErr, io.EOF) {
			return nil, errors.Wrap(secondBatchErr, "unable to read directory entries")
		}

		// process results in case it's not EOF.
		for _, de := range secondBatch {
			e, err := toDirEntryOrNil(de, fullPath)
			if err != nil {
				return nil, errors.Wrap(err, "error reading entry")
			}

			if e != nil {
				entries = append(entries, e)
			}
		}

		// if we got EOF at this point, return.
		if errors.Is(secondBatchErr, io.EOF) {
			entries.Sort()

			return entries, nil
		}
	}

	return fsd.readRemainingDirEntriesInParallel(fullPath, entries, f)
}

func (fsd *filesystemDirectory) readRemainingDirEntriesInParallel(fullPath string, entries fs.Entries, f *os.File) (fs.Entries, error) {
	// start feeding directory entries to dirEntryCh
	dirEntryCh := make(chan string, dirListingPrefetch)

	var readDirErr error

	go func() {
		defer close(dirEntryCh)

		for {
			des, err := f.Readdirnames(numEntriesToRead)
			for _, de := range des {
				dirEntryCh <- de
			}

			if err == nil {
				continue
			}

			if errors.Is(err, io.EOF) {
				break
			}

			readDirErr = err

			break
		}
	}()

	entriesCh := make(chan entryWithError, dirListingPrefetch)

	var workersWG sync.WaitGroup

	for i := 0; i < paralellelStatGoroutines; i++ {
		workersWG.Add(1)

		go func() {
			defer workersWG.Done()

			for de := range dirEntryCh {
				e, err := toDirEntryOrNil(de, fullPath)
				if err != nil {
					entriesCh <- entryWithError{err: errors.Errorf("unable to stat directory entry %q: %v", de, err)}
					continue
				}

				if e != nil {
					entriesCh <- entryWithError{entry: e}
				}
			}
		}()
	}

	// close entriesCh channel when all goroutines terminate
	go func() {
		workersWG.Wait()
		close(entriesCh)
	}()

	// drain the entriesCh into a slice and sort it

	for e := range entriesCh {
		if e.err != nil {
			// only return the first error
			if readDirErr == nil {
				readDirErr = e.err
			}

			continue
		}

		entries = append(entries, e.entry)
	}

	entries.Sort()

	// return any error encountered when listing or reading the directory
	return entries, readDirErr
}

type fileWithMetadata struct {
	*os.File
}

func (f *fileWithMetadata) Entry() (fs.Entry, error) {
	fi, err := f.Stat()
	if err != nil {
		return nil, errors.Wrap(err, "unable to stat() local file")
	}

	return &filesystemFile{newEntry(fi, filepath.Dir(f.Name()))}, nil
}

func (fsf *filesystemFile) Open(ctx context.Context) (fs.Reader, error) {
	f, err := os.Open(fsf.fullPath())
	if err != nil {
		return nil, errors.Wrap(err, "unable to open local file")
	}

	return &fileWithMetadata{f}, nil
}

func (fsl *filesystemSymlink) Readlink(ctx context.Context) (string, error) {
	// nolint:wrapcheck
	return os.Readlink(fsl.fullPath())
}

func (e *filesystemErrorEntry) ErrorInfo() error {
	return e.err
}

// NewEntry returns fs.Entry for the specified path, the result will be one of supported entry types: fs.File, fs.Directory, fs.Symlink
// or fs.UnsupportedEntry.
func NewEntry(path string) (fs.Entry, error) {
	fi, err := os.Lstat(path)
	if err != nil {
		return nil, errors.Wrap(err, "unable to determine entry type")
	}

	return entryFromDirEntry(fi, filepath.Dir(path)), nil
}

// Directory returns fs.Directory for the specified path.
func Directory(path string) (fs.Directory, error) {
	e, err := NewEntry(path)
	if err != nil {
		return nil, err
	}

	if d, ok := e.(fs.Directory); ok {
		return d, nil
	}

	return nil, errors.Errorf("not a directory: %v", path)
}

func entryFromDirEntry(fi os.FileInfo, parentDir string) fs.Entry {
	isplaceholder := strings.HasSuffix(fi.Name(), ShallowEntrySuffix)
	maskedmode := fi.Mode() & os.ModeType

	switch {
	case maskedmode == os.ModeDir && !isplaceholder:
		return &filesystemDirectory{newEntry(fi, parentDir)}

	case maskedmode == os.ModeDir && isplaceholder:
		return &shallowFilesystemDirectory{newEntry(fi, parentDir)}

	case maskedmode == os.ModeSymlink && !isplaceholder:
		return &filesystemSymlink{newEntry(fi, parentDir)}

	case maskedmode == 0 && !isplaceholder:
		return &filesystemFile{newEntry(fi, parentDir)}

	case maskedmode == 0 && isplaceholder:
		return &shallowFilesystemFile{newEntry(fi, parentDir)}

	default:
		return &filesystemErrorEntry{newEntry(fi, parentDir), fs.ErrUnknown}
	}
}

var (
	_ fs.Directory  = &filesystemDirectory{}
	_ fs.File       = &filesystemFile{}
	_ fs.Symlink    = &filesystemSymlink{}
	_ fs.ErrorEntry = &filesystemErrorEntry{}
)
<|endoftext|>


<commit_before>package http

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"path"

	"github.com/mediocregopher/gobdns/config"
	"github.com/mediocregopher/gobdns/ips"
)

var usage = `
	GET    /                     Gives you this page

	GET    /api/domains/all      Gives you a space separated mapping of domains
	                             to ips

	POST   /api/domains/<domain> Maps the given domain to the given ip, which
	                             will be the body data for the request

	PUT    /api/domains/<domain> Same as POST'ing

	DELETE /api/domains/<domain> Removes the domain->ip mapping for the given
	                             domain

`

func init() {
	if config.APIAddr == "" {
		return
	}

	go func() {
		log.Printf("API Listening on %s", config.APIAddr)
		http.HandleFunc("/api/domains/all", getAll)
		http.HandleFunc("/api/domains/", putDelete)
		http.HandleFunc("/", root)
		http.ListenAndServe(config.APIAddr, nil)
	}()
}

func getAll(w http.ResponseWriter, r *http.Request) {
	for domain, ip := range ips.GetAll() {
		fmt.Fprintf(w, "%s %s\n", domain, ip)
	}
}

func putDelete(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/api/domains/" {
		w.WriteHeader(404)
		return
	}

	domain := path.Base(r.URL.Path)

	switch r.Method {
	case "PUT", "POST":
		ipB, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Println(err)
			w.WriteHeader(500)
			return
		}
		ip := string(ipB)
		if ip == "" {
			w.WriteHeader(400)
			return
		}
		ips.Set(domain, ip)

	case "DELETE":
		ips.Unset(domain)
	}

	r.Body.Close()
}

func root(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, usage)
}
<commit_msg>add some parsing checking in rest api<commit_after>package http

import (
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"path"
	"strings"

	"github.com/mediocregopher/gobdns/config"
	"github.com/mediocregopher/gobdns/ips"
)

var usage = `
	GET    /                     Gives you this page

	GET    /api/domains/all      Gives you a space separated mapping of domains
	                             to ips

	POST   /api/domains/<domain> Maps the given domain to the given ip, which
	                             will be the body data for the request

	PUT    /api/domains/<domain> Same as POST'ing

	DELETE /api/domains/<domain> Removes the domain->ip mapping for the given
	                             domain

`

func init() {
	if config.APIAddr == "" {
		return
	}

	go func() {
		log.Printf("API Listening on %s", config.APIAddr)
		http.HandleFunc("/api/domains/all", getAll)
		http.HandleFunc("/api/domains/", putDelete)
		http.HandleFunc("/", root)
		http.ListenAndServe(config.APIAddr, nil)
	}()
}

func getAll(w http.ResponseWriter, r *http.Request) {
	for domain, ip := range ips.GetAll() {
		fmt.Fprintf(w, "%s %s\n", domain, ip)
	}
}

func putDelete(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/api/domains/" {
		w.WriteHeader(404)
		return
	}

	domain := path.Base(r.URL.Path)

	switch r.Method {
	case "PUT", "POST":
		ipB, err := ioutil.ReadAll(r.Body)
		if err != nil {
			log.Println(err)
			w.WriteHeader(500)
			return
		}
		ip := strings.TrimSpace(string(ipB))
		if ip == "" {
			w.WriteHeader(400)
			fmt.Fprintf(w, "no ip given in request body")
			return
		}
		if net.ParseIP(ip) == nil {
			w.WriteHeader(400)
			fmt.Fprintf(w, "invalid ip given in request body")
			return
		}
		ips.Set(domain, ip)

	case "DELETE":
		ips.Unset(domain)
	}

	r.Body.Close()
}

func root(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, usage)
}
<|endoftext|>


<commit_before>package benchmarkbbs_test

import (
	"fmt"
	"math/rand"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"code.cloudfoundry.org/bbs/models"
	"code.cloudfoundry.org/benchmarkbbs/reporter"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/operationq"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

const (
	RepBulkFetching                   = "RepBulkFetching"
	RepBulkLoop                       = "RepBulkLoop"
	RepClaimActualLRP                 = "RepClaimActualLRP"
	RepStartActualLRP                 = "RepStartActualLRP"
	NsyncBulkerFetching               = "NsyncBulkerFetching"
	ConvergenceGathering              = "ConvergenceGathering"
	FetchActualLRPsAndSchedulingInfos = "FetchActualLRPsAndSchedulingInfos"
)

var bulkCycle = 30 * time.Second
var eventCount int32 = 0
var claimCount int32 = 0

var BenchmarkTests = func(numReps, numTrials int, localRouteEmitters bool) {
	Describe("main benchmark test", func() {
		eventCountRunner := func(signals <-chan os.Signal, ready chan<- struct{}) error {
			eventSource, err := bbsClient.SubscribeToEvents(logger)
			Expect(err).NotTo(HaveOccurred())
			close(ready)

			eventChan := make(chan models.Event)
			go func() {
				for {
					event, err := eventSource.Next()
					if err != nil {
						logger.Error("error-getting-next-event", err)
						return
					}
					if event != nil {
						eventChan <- event
					}
				}
			}()

			for {
				select {
				case <-eventChan:
					atomic.AddInt32(&eventCount, 1)

				case <-signals:
					if eventSource != nil {
						err := eventSource.Close()
						if err != nil {
							logger.Error("failed-closing-event-source", err)
						}
					}
					return nil
				}
			}
		}

		var process ifrit.Process

		BeforeEach(func() {
			process = ifrit.Invoke(ifrit.RunFunc(eventCountRunner))
		})

		AfterEach(func() {
			ginkgomon.Kill(process)
		})

		Measure("data for benchmarks", func(b Benchmarker) {
			wg := sync.WaitGroup{}

			// start nsync
			go func() {
				defer GinkgoRecover()
				logger.Info("start-nsync-bulker-loop")
				defer logger.Info("finish-nsync-bulker-loop")
				wg.Add(1)
				defer wg.Done()
				for i := 0; i < numTrials; i++ {
					sleepDuration := getSleepDuration(i, bulkCycle)
					time.Sleep(sleepDuration)
					b.Time("fetch all desired LRP scheduling info", func() {
						desireds, err := bbsClient.DesiredLRPSchedulingInfos(logger, models.DesiredLRPFilter{})
						Expect(err).NotTo(HaveOccurred())
						Expect(len(desireds)).To(BeNumerically("~", expectedLRPCount, expectedLRPVariation), "Number of DesiredLRPs retrieved in Nsync Bulk Loop")
					}, reporter.ReporterInfo{
						MetricName: NsyncBulkerFetching,
					})
				}
			}()

			// start convergence
			go func() {
				defer GinkgoRecover()
				logger.Info("start-lrp-convergence-loop")
				defer logger.Info("finish-lrp-convergence-loop")
				wg.Add(1)
				defer wg.Done()
				for i := 0; i < numTrials; i++ {
					sleepDuration := getSleepDuration(i, bulkCycle)
					time.Sleep(sleepDuration)
					cellSet := models.NewCellSet()
					for i := 0; i < numReps; i++ {
						cellID := fmt.Sprintf("cell-%d", i)
						presence := models.NewCellPresence(cellID, "earth", "http://planet-earth", "north", models.CellCapacity{}, nil, nil, nil, nil)
						cellSet.Add(&presence)
					}

					b.Time("BBS' internal gathering of LRPs", func() {
						activeDB.ConvergeLRPs(logger, cellSet)
					}, reporter.ReporterInfo{
						MetricName: ConvergenceGathering,
					})
				}
			}()

			// we need to make sure we don't run out of ports so limit amount of
			// active http requests to 25000
			semaphore := make(chan struct{}, 25000)

			numRouteEmitters := 1

			if localRouteEmitters {
				numRouteEmitters = numReps
			}

			for i := 0; i < numRouteEmitters; i++ {
				cellID := fmt.Sprintf("cell-%d", i)
				wg.Add(1)

				// start route-emitter
				go func() {
					defer GinkgoRecover()

					lagerData := lager.Data{}
					if localRouteEmitters {
						lagerData = lager.Data{"cell-id": cellID}
					}
					logger := logger.WithData(lagerData)
					logger.Info("start-route-emitter-loop")
					defer logger.Info("finish-route-emitter-loop")

					defer wg.Done()

					for j := 0; j < numTrials; j++ {
						sleepDuration := getSleepDuration(j, bulkCycle)
						time.Sleep(sleepDuration)
						b.Time("fetch all actualLRPs and schedulingInfos", func() {
							semaphore <- struct{}{}
							actuals, err := bbsClient.ActualLRPGroups(logger, models.ActualLRPFilter{})
							<-semaphore
							Expect(err).NotTo(HaveOccurred())
							Expect(len(actuals)).To(BeNumerically("~", expectedLRPCount, expectedLRPVariation), "Number of ActualLRPs retrieved in router-emitter")

							semaphore <- struct{}{}
							desireds, err := bbsClient.DesiredLRPSchedulingInfos(logger, models.DesiredLRPFilter{})
							<-semaphore
							Expect(err).NotTo(HaveOccurred())
							Expect(len(desireds)).To(BeNumerically("~", expectedLRPCount, expectedLRPVariation), "Number of DesiredLRPs retrieved in route-emitter")
						}, reporter.ReporterInfo{
							MetricName: FetchActualLRPsAndSchedulingInfos,
						})
					}
				}()
			}

			totalRan := int32(0)
			totalQueued := int32(0)
			var err error
			queue := operationq.NewSlidingQueue(numTrials)

			for i := 0; i < numReps; i++ {
				cellID := fmt.Sprintf("cell-%d", i)
				wg.Add(1)

				go func(cellID string) {
					defer GinkgoRecover()
					defer wg.Done()

					for j := 0; j < numTrials; j++ {
						sleepDuration := getSleepDuration(j, bulkCycle)
						time.Sleep(sleepDuration)

						b.Time("rep bulk loop", func() {
							defer GinkgoRecover()
							var actuals []*models.ActualLRPGroup
							b.Time("rep bulk fetch", func() {
								semaphore <- struct{}{}
								actuals, err = bbsClient.ActualLRPGroups(logger, models.ActualLRPFilter{CellID: cellID})
								<-semaphore
								Expect(err).NotTo(HaveOccurred())
							}, reporter.ReporterInfo{
								MetricName: RepBulkFetching,
							})

							expectedActualLRPCount, ok := expectedActualLRPCounts[cellID]
							Expect(ok).To(BeTrue())

							expectedActualLRPVariation, ok := expectedActualLRPVariations[cellID]
							Expect(ok).To(BeTrue())

							Expect(len(actuals)).To(BeNumerically("~", expectedActualLRPCount, expectedActualLRPVariation), "Number of ActualLRPs retrieved by cell %s in rep bulk loop", cellID)

							numActuals := len(actuals)
							for k := 0; k < numActuals; k++ {
								actualLRP, _ := actuals[k].Resolve()
								atomic.AddInt32(&totalQueued, 1)
								queue.Push(&lrpOperation{actualLRP, config.PercentWrites, b, &totalRan, &claimCount, semaphore})
							}
						}, reporter.ReporterInfo{
							MetricName: RepBulkLoop,
						})
					}
				}(cellID)
			}

			wg.Wait()

			eventTolerance := float64(claimCount) * config.ErrorTolerance
			Eventually(func() int32 { return atomic.LoadInt32(&eventCount) }, 2*time.Minute).Should(BeNumerically("~", claimCount, eventTolerance), "events received")
			Eventually(func() int32 { return atomic.LoadInt32(&totalRan) }, 2*time.Minute).Should(Equal(totalQueued), "should have run the same number of queued LRP operations")
		}, 1)
	})
}

type lrpOperation struct {
	actualLRP        *models.ActualLRP
	percentWrites    float64
	b                Benchmarker
	globalCount      *int32
	globalClaimCount *int32
	semaphore        chan struct{}
}

func (lo *lrpOperation) Key() string {
	return lo.actualLRP.ProcessGuid
}

func (lo *lrpOperation) Execute() {
	defer GinkgoRecover()
	defer atomic.AddInt32(lo.globalCount, 1)
	var err error
	randomNum := rand.Float64() * 100.0

	// divided by 2 because the start following the claim cause two writes.
	isClaiming := randomNum < (lo.percentWrites / 2)
	actualLRP := lo.actualLRP

	lo.b.Time("start actual LRP", func() {
		netInfo := models.NewActualLRPNetInfo("1.2.3.4", models.NewPortMapping(61999, 8080))
		lo.semaphore <- struct{}{}
		err = bbsClient.StartActualLRP(logger, &actualLRP.ActualLRPKey, &actualLRP.ActualLRPInstanceKey, &netInfo)
		<-lo.semaphore
		Expect(err).NotTo(HaveOccurred())
		if actualLRP.State == models.ActualLRPStateClaimed {
			defer atomic.AddInt32(lo.globalClaimCount, 1)
		}
	}, reporter.ReporterInfo{
		MetricName: RepStartActualLRP,
	})

	if isClaiming {
		lo.b.Time("claim actual LRP", func() {
			index := int(actualLRP.ActualLRPKey.Index)
			lo.semaphore <- struct{}{}
			err = bbsClient.ClaimActualLRP(logger, actualLRP.ActualLRPKey.ProcessGuid, index, &actualLRP.ActualLRPInstanceKey)
			<-lo.semaphore
			Expect(err).NotTo(HaveOccurred())
			defer atomic.AddInt32(lo.globalClaimCount, 1)
		}, reporter.ReporterInfo{
			MetricName: RepClaimActualLRP,
		})
	}
}

func getSleepDuration(loopCounter int, cycleTime time.Duration) time.Duration {
	sleepDuration := cycleTime
	if loopCounter == 0 {
		numMilli := rand.Intn(int(cycleTime.Nanoseconds() / 1000000))
		sleepDuration = time.Duration(numMilli) * time.Millisecond
	}
	return sleepDuration
}
<commit_msg>route emitters subscribe to event stream<commit_after>package benchmarkbbs_test

import (
	"fmt"
	"math/rand"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"code.cloudfoundry.org/bbs/models"
	"code.cloudfoundry.org/benchmarkbbs/reporter"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/operationq"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

const (
	RepBulkFetching                   = "RepBulkFetching"
	RepBulkLoop                       = "RepBulkLoop"
	RepClaimActualLRP                 = "RepClaimActualLRP"
	RepStartActualLRP                 = "RepStartActualLRP"
	NsyncBulkerFetching               = "NsyncBulkerFetching"
	ConvergenceGathering              = "ConvergenceGathering"
	FetchActualLRPsAndSchedulingInfos = "FetchActualLRPsAndSchedulingInfos"
)

var bulkCycle = 30 * time.Second
var eventCount int32 = 0
var expectedEventCount int32 = 0

func eventCountRunner(counter *int32) func(signals <-chan os.Signal, ready chan<- struct{}) error {
	return func(signals <-chan os.Signal, ready chan<- struct{}) error {
		eventSource, err := bbsClient.SubscribeToEvents(logger)
		Expect(err).NotTo(HaveOccurred())
		close(ready)

		eventChan := make(chan models.Event)
		go func() {
			for {
				event, err := eventSource.Next()
				if err != nil {
					logger.Error("error-getting-next-event", err)
					return
				}
				if event != nil {
					eventChan <- event
				}
			}
		}()

		for {
			select {
			case <-eventChan:
				atomic.AddInt32(counter, 1)

			case <-signals:
				if eventSource != nil {
					err := eventSource.Close()
					if err != nil {
						logger.Error("failed-closing-event-source", err)
					}
				}
				return nil
			}
		}
	}
}

var BenchmarkTests = func(numReps, numTrials int, localRouteEmitters bool) {
	Describe("main benchmark test", func() {
		var process ifrit.Process

		BeforeEach(func() {
			process = ifrit.Invoke(ifrit.RunFunc(eventCountRunner(&eventCount)))
		})

		AfterEach(func() {
			ginkgomon.Kill(process)
		})

		Measure("data for benchmarks", func(b Benchmarker) {
			wg := sync.WaitGroup{}

			// start nsync
			wg.Add(1)
			go nsyncBulkerLoop(b, &wg, numTrials)

			// start convergence
			wg.Add(1)
			go convergence(b, &wg, numTrials, numReps)

			// we need to make sure we don't run out of ports so limit amount of
			// active http requests to 25000
			semaphore := make(chan struct{}, 25000)

			numRouteEmitters := 1

			if localRouteEmitters {
				numRouteEmitters = numReps
			}

			routeEmitterEventCounts := make(map[string]*int32)

			for i := 0; i < numRouteEmitters; i++ {
				cellID := fmt.Sprintf("cell-%d", i)

				routeEmitterEventCount := new(int32)
				routeEmitterEventCounts[cellID] = routeEmitterEventCount

				// start route-emitter
				wg.Add(1)
				go routeEmitter(b, &wg, localRouteEmitters, cellID, routeEmitterEventCount, semaphore, numTrials)
			}

			queue := operationq.NewSlidingQueue(numTrials)

			totalRan := int32(0)
			totalQueued := int32(0)

			for i := 0; i < numReps; i++ {
				cellID := fmt.Sprintf("cell-%d", i)
				wg.Add(1)

				go repBulker(b, &wg, cellID, numTrials, semaphore, &totalQueued, &totalRan, &expectedEventCount, queue)
			}

			wg.Wait()

			eventTolerance := float64(atomic.LoadInt32(&expectedEventCount)) * config.ErrorTolerance

			Eventually(func() int32 {
				return atomic.LoadInt32(&eventCount)
			}, 2*time.Minute).Should(BeNumerically("~", expectedEventCount, eventTolerance), "events received")

			Eventually(func() int32 {
				return atomic.LoadInt32(&totalRan)
			}, 2*time.Minute).Should(Equal(totalQueued), "should have run the same number of queued LRP operations")

			for _, v := range routeEmitterEventCounts {
				Eventually(func() int32 {
					return atomic.LoadInt32(v)
				}, 2*time.Minute).Should(BeNumerically("~", expectedEventCount, eventTolerance), "events received")
			}
		}, 1)
	})
}

func getSleepDuration(loopCounter int, cycleTime time.Duration) time.Duration {
	sleepDuration := cycleTime
	if loopCounter == 0 {
		numMilli := rand.Intn(int(cycleTime.Nanoseconds() / 1000000))
		sleepDuration = time.Duration(numMilli) * time.Millisecond
	}
	return sleepDuration
}

func nsyncBulkerLoop(b Benchmarker, wg *sync.WaitGroup, numTrials int) {
	defer GinkgoRecover()
	logger.Info("start-nsync-bulker-loop")
	defer logger.Info("finish-nsync-bulker-loop")
	defer wg.Done()

	for i := 0; i < numTrials; i++ {
		sleepDuration := getSleepDuration(i, bulkCycle)
		time.Sleep(sleepDuration)
		b.Time("fetch all desired LRP scheduling info", func() {
			desireds, err := bbsClient.DesiredLRPSchedulingInfos(logger, models.DesiredLRPFilter{})
			Expect(err).NotTo(HaveOccurred())
			Expect(len(desireds)).To(BeNumerically("~", expectedLRPCount, expectedLRPVariation), "Number of DesiredLRPs retrieved in Nsync Bulk Loop")
		}, reporter.ReporterInfo{
			MetricName: NsyncBulkerFetching,
		})
	}
}

func convergence(b Benchmarker, wg *sync.WaitGroup, numTrials, numReps int) {
	defer GinkgoRecover()
	logger.Info("start-lrp-convergence-loop")
	defer logger.Info("finish-lrp-convergence-loop")
	defer wg.Done()

	for i := 0; i < numTrials; i++ {
		sleepDuration := getSleepDuration(i, bulkCycle)
		time.Sleep(sleepDuration)
		cellSet := models.NewCellSet()
		for i := 0; i < numReps; i++ {
			cellID := fmt.Sprintf("cell-%d", i)
			presence := models.NewCellPresence(cellID, "earth", "http://planet-earth", "north", models.CellCapacity{}, nil, nil, nil, nil)
			cellSet.Add(&presence)
		}

		b.Time("BBS' internal gathering of LRPs", func() {
			activeDB.ConvergeLRPs(logger, cellSet)
		}, reporter.ReporterInfo{
			MetricName: ConvergenceGathering,
		})
	}
}

func repBulker(b Benchmarker, wg *sync.WaitGroup, cellID string, numTrials int, semaphore chan struct{}, totalQueued, totalRan, expectedEventCount *int32, queue operationq.Queue) {
	defer GinkgoRecover()
	defer wg.Done()

	var err error

	for j := 0; j < numTrials; j++ {
		sleepDuration := getSleepDuration(j, bulkCycle)
		time.Sleep(sleepDuration)

		b.Time("rep bulk loop", func() {
			defer GinkgoRecover()
			var actuals []*models.ActualLRPGroup
			b.Time("rep bulk fetch", func() {
				semaphore <- struct{}{}
				actuals, err = bbsClient.ActualLRPGroups(logger, models.ActualLRPFilter{CellID: cellID})
				<-semaphore
				Expect(err).NotTo(HaveOccurred())
			}, reporter.ReporterInfo{
				MetricName: RepBulkFetching,
			})

			expectedActualLRPCount, ok := expectedActualLRPCounts[cellID]
			Expect(ok).To(BeTrue())

			expectedActualLRPVariation, ok := expectedActualLRPVariations[cellID]
			Expect(ok).To(BeTrue())

			Expect(len(actuals)).To(BeNumerically("~", expectedActualLRPCount, expectedActualLRPVariation), "Number of ActualLRPs retrieved by cell %s in rep bulk loop", cellID)

			numActuals := len(actuals)
			for k := 0; k < numActuals; k++ {
				actualLRP, _ := actuals[k].Resolve()
				atomic.AddInt32(totalQueued, 1)
				queue.Push(&lrpOperation{actualLRP, config.PercentWrites, b, totalRan, expectedEventCount, semaphore})
			}
		}, reporter.ReporterInfo{
			MetricName: RepBulkLoop,
		})
	}
}

func routeEmitter(b Benchmarker, wg *sync.WaitGroup, localRouteEmitters bool, cellID string, routeEmitterEventCount *int32, semaphore chan struct{}, numTrials int) {
	defer GinkgoRecover()

	lagerData := lager.Data{}
	if localRouteEmitters {
		lagerData = lager.Data{"cell-id": cellID}
	}
	logger := logger.WithData(lagerData)
	logger.Info("start-route-emitter-loop")
	defer logger.Info("finish-route-emitter-loop")

	defer wg.Done()

	ifrit.Invoke(ifrit.RunFunc(eventCountRunner(routeEmitterEventCount)))

	for j := 0; j < numTrials; j++ {
		sleepDuration := getSleepDuration(j, bulkCycle)
		time.Sleep(sleepDuration)
		b.Time("fetch all actualLRPs and schedulingInfos", func() {
			semaphore <- struct{}{}
			actuals, err := bbsClient.ActualLRPGroups(logger, models.ActualLRPFilter{})
			<-semaphore
			Expect(err).NotTo(HaveOccurred())
			Expect(len(actuals)).To(BeNumerically("~", expectedLRPCount, expectedLRPVariation), "Number of ActualLRPs retrieved in router-emitter")

			semaphore <- struct{}{}
			desireds, err := bbsClient.DesiredLRPSchedulingInfos(logger, models.DesiredLRPFilter{})
			<-semaphore
			Expect(err).NotTo(HaveOccurred())
			Expect(len(desireds)).To(BeNumerically("~", expectedLRPCount, expectedLRPVariation), "Number of DesiredLRPs retrieved in route-emitter")

		}, reporter.ReporterInfo{
			MetricName: FetchActualLRPsAndSchedulingInfos,
		})
	}
}

type lrpOperation struct {
	actualLRP        *models.ActualLRP
	percentWrites    float64
	b                Benchmarker
	globalCount      *int32
	globalEventCount *int32
	semaphore        chan struct{}
}

func (lo *lrpOperation) Key() string {
	return lo.actualLRP.ProcessGuid
}

func (lo *lrpOperation) Execute() {
	defer GinkgoRecover()
	defer atomic.AddInt32(lo.globalCount, 1)
	var err error
	randomNum := rand.Float64() * 100.0

	// divided by 2 because the start following the claim cause two writes.
	isClaiming := randomNum < (lo.percentWrites / 2)
	actualLRP := lo.actualLRP

	lo.b.Time("start actual LRP", func() {
		netInfo := models.NewActualLRPNetInfo("1.2.3.4", models.NewPortMapping(61999, 8080))
		lo.semaphore <- struct{}{}
		err = bbsClient.StartActualLRP(logger, &actualLRP.ActualLRPKey, &actualLRP.ActualLRPInstanceKey, &netInfo)
		<-lo.semaphore
		Expect(err).NotTo(HaveOccurred())

		// if the actual lrp was not already started, an event will be generated
		if actualLRP.State != models.ActualLRPStateRunning {
			atomic.AddInt32(lo.globalEventCount, 1)
		}
	}, reporter.ReporterInfo{
		MetricName: RepStartActualLRP,
	})

	if isClaiming {
		lo.b.Time("claim actual LRP", func() {
			index := int(actualLRP.ActualLRPKey.Index)
			lo.semaphore <- struct{}{}
			err = bbsClient.ClaimActualLRP(logger, actualLRP.ActualLRPKey.ProcessGuid, index, &actualLRP.ActualLRPInstanceKey)
			<-lo.semaphore
			Expect(err).NotTo(HaveOccurred())
			atomic.AddInt32(lo.globalEventCount, 1)
		}, reporter.ReporterInfo{
			MetricName: RepClaimActualLRP,
		})
	}
}
<|endoftext|>


<commit_before>package program

// Version can be requested through the command line with:
//
//     c2go -v
//
// See https://github.com/elliotchance/c2go/wiki/Release-Process
const Version = "v0.21.3 Zinc 2018-01-28"
<commit_msg>Bump version: v0.21.4 Zinc 2018-01-30<commit_after>package program

// Version can be requested through the command line with:
//
//     c2go -v
//
// See https://github.com/elliotchance/c2go/wiki/Release-Process
const Version = "v0.21.4 Zinc 2018-01-30"
<|endoftext|>


<commit_before>// Copyright (C) 2017 Google 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,
// 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 client

import (
	"context"
	"fmt"
	"io"
	"net"
	"reflect"
	"time"

	"github.com/google/gapid/core/app/crash"
	"github.com/google/gapid/core/context/keys"
	"github.com/google/gapid/core/event/task"
	"github.com/google/gapid/core/java/jdbg"
	"github.com/google/gapid/core/java/jdwp"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/android/adb"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/gapidapk"
)

func expect(r io.Reader, expected []byte) error {
	got := make([]byte, len(expected))
	if _, err := io.ReadFull(r, got); err != nil {
		return err
	}
	if !reflect.DeepEqual(expected, got) {
		return fmt.Errorf("Expected %v, got %v", expected, got)
	}
	return nil
}

// waitForOnCreate waits for android.app.Application.onCreate to be called, and
// then suspends the thread.
func waitForOnCreate(ctx context.Context, conn *jdwp.Connection, wakeup jdwp.ThreadID) (*jdwp.EventMethodEntry, error) {
	app, err := conn.GetClassBySignature("Landroid/app/Application;")
	if err != nil {
		return nil, err
	}

	constructor, err := conn.GetClassMethod(app.ClassID(), "<init>", "()V")
	if err != nil {
		return nil, err
	}

	log.I(ctx, "   Waiting for Application.<init>()")
	return conn.WaitForMethodEntry(ctx, app.ClassID(), constructor.ID, wakeup)
}

// waitForVulkanLoad for android.app.ApplicationLoaders.getClassLoader to be called,
// and then suspends the thread.
// This function is what is used to tell the vulkan loader where to search for
// layers.
func waitForVulkanLoad(ctx context.Context, conn *jdwp.Connection) (*jdwp.EventMethodEntry, error) {
	loaders, err := conn.GetClassBySignature("Landroid/app/ApplicationLoaders;")
	if err != nil {
		return nil, err
	}

	getClassLoader, err := conn.GetClassMethod(loaders.ClassID(), "getClassLoader",
		"(Ljava/lang/String;IZLjava/lang/String;Ljava/lang/String;Ljava/lang/ClassLoader;)Ljava/lang/ClassLoader;")
	if err != nil {
		return nil, err
	}

	return conn.WaitForMethodEntry(ctx, loaders.ClassID(), getClassLoader.ID, 0)
}

// loadAndConnectViaJDWP connects to the application waiting for a JDWP
// connection with the specified process id, sends a number of JDWP commands to
// load the list of libraries.
func (p *Process) loadAndConnectViaJDWP(
	ctx context.Context,
	gapidAPK *gapidapk.APK,
	pid int,
	d adb.Device) error {

	const (
		reconnectAttempts = 10
		reconnectDelay    = time.Second
	)

	jdwpPort, err := adb.LocalFreeTCPPort()
	if err != nil {
		return log.Err(ctx, err, "Finding free port")
	}
	ctx = log.V{"jdwpPort": jdwpPort}.Bind(ctx)

	log.I(ctx, "Forwarding TCP port %v -> JDWP pid %v", jdwpPort, pid)
	if err := d.Forward(ctx, adb.TCPPort(jdwpPort), adb.Jdwp(pid)); err != nil {
		return log.Err(ctx, err, "Setting up JDWP port forwarding")
	}
	defer func() {
		// Clone context to ignore cancellation.
		ctx := keys.Clone(context.Background(), ctx)
		d.RemoveForward(ctx, adb.TCPPort(jdwpPort))
	}()

	ctx, stop := task.WithCancel(ctx)
	defer stop()

	log.I(ctx, "Connecting to JDWP")

	// Create a JDWP connection with the application.
	var sock net.Conn
	var conn *jdwp.Connection
	err = task.Retry(ctx, reconnectAttempts, reconnectDelay, func(ctx context.Context) (bool, error) {
		if sock, err = net.Dial("tcp", fmt.Sprintf("localhost:%v", jdwpPort)); err != nil {
			return false, err
		}
		if conn, err = jdwp.Open(ctx, sock); err != nil {
			sock.Close()
			return false, err
		}
		return true, nil
	})
	if err != nil {
		if err == io.EOF {
			return fmt.Errorf("Unable to connect to the application.\n\n" +
				"This can happen when another debugger or IDE is running " +
				"in the background, such as Android Studio.\n" +
				"Please close any running Android debuggers and try again.\n\n" +
				"See https://github.com/google/gapid/issues/911 for more " +
				"information")
		}
		return log.Err(ctx, err, "Connecting to JDWP")
	}
	defer sock.Close()

	processABI := func(j *jdbg.JDbg) (*device.ABI, error) {
		abiName := j.Class("android.os.Build").Field("CPU_ABI").Get().(string)
		abi := device.ABIByName(abiName)
		if abi == nil {
			return nil, fmt.Errorf("Unknown ABI %v", abiName)
		}

		// For NativeBridge emulated devices opt for the native ABI of the
		// emulator.
		abi = d.NativeBridgeABI(ctx, abi)

		return abi, nil
	}

	classLoaderThread := jdwp.ThreadID(0)

	log.I(ctx, "Waiting for ApplicationLoaders.getClassLoader()")
	getClassLoader, err := waitForVulkanLoad(ctx, conn)
	if err == nil {
		// If err != nil that means we could not find or break in getClassLoader
		// so we have no vulkan support.
		classLoaderThread = getClassLoader.Thread
		err = jdbg.Do(conn, getClassLoader.Thread, func(j *jdbg.JDbg) error {
			abi, err := processABI(j)
			if err != nil {
				return err
			}
			libsPath := gapidAPK.LibsPath(abi)
			newLibraryPath := j.String(":" + libsPath)
			obj := j.GetStackObject("librarySearchPath").Call("concat", newLibraryPath)
			j.SetStackObject("librarySearchPath", obj)
			return nil
		})
		if err != nil {
			return log.Err(ctx, err, "JDWP failure")
		}
	} else {
		log.W(ctx, "Couldn't break in ApplicationLoaders.getClassLoader. Vulkan will not be supported.")
	}

	// Wait for Application.onCreate to be called.
	log.I(ctx, "Waiting for Application Creation")
	onCreate, err := waitForOnCreate(ctx, conn, classLoaderThread)
	if err != nil {
		return log.Err(ctx, err, "Waiting for Application Creation")
	}

	// Attempt to get the GVR library handle.
	// Will throw an exception for non-GVR apps.
	var gvrHandle uint64
	log.I(ctx, "Installing interceptor libraries")
	loadNativeGvrLibrary, vrCoreLibraryLoader := "loadNativeGvrLibrary", "com/google/vr/cardboard/VrCoreLibraryLoader"
	gvrMajor, gvrMinor, gvrPoint := 1, 8, 1

	getGVRHandle := func(j *jdbg.JDbg, libLoader jdbg.Type) error {
		// loadNativeGvrLibrary has a couple of different signatures depending
		// on GVR release.
		for _, f := range []func() error{
			// loadNativeGvrLibrary(Context, int major, int minor, int point)
			func() error {
				gvrHandle = (uint64)(libLoader.Call(loadNativeGvrLibrary, j.This(), gvrMajor, gvrMinor, gvrPoint).Get().(int64))
				return nil
			},
			// loadNativeGvrLibrary(Context)
			func() error {
				gvrHandle = (uint64)(libLoader.Call(loadNativeGvrLibrary, j.This()).Get().(int64))
				return nil
			},
		} {
			if jdbg.Try(f) == nil {
				return nil
			}
		}
		return fmt.Errorf("Couldn't call loadNativeGvrLibrary")
	}
	for _, f := range []func(j *jdbg.JDbg) error{
		func(j *jdbg.JDbg) error {
			libLoader := j.Class(vrCoreLibraryLoader)
			getGVRHandle(j, libLoader)
			return nil
		},
		func(j *jdbg.JDbg) error {
			classLoader := j.This().Call("getClassLoader")
			libLoader := classLoader.Call("findClass", vrCoreLibraryLoader).AsType()
			getGVRHandle(j, libLoader)
			return nil
		},
	} {
		if err := jdbg.Do(conn, onCreate.Thread, f); err == nil {
			break
		}
	}
	if gvrHandle == 0 {
		log.I(ctx, "GVR library not found")
	} else {
		log.I(ctx, "GVR library found")
	}

	// Connect to GAPII.
	// This has to be done on a separate go-routine as the call to load gapii
	// will block until a connection is made.
	connErr := make(chan error)

	// Load GAPII library.
	err = jdbg.Do(conn, onCreate.Thread, func(j *jdbg.JDbg) error {
		abi, err := processABI(j)
		if err != nil {
			return err
		}

		interceptorPath := gapidAPK.LibInterceptorPath(abi)
		crash.Go(func() { connErr <- p.connect(ctx, gvrHandle, interceptorPath) })

		gapiiPath := gapidAPK.LibGAPIIPath(abi)
		ctx = log.V{"gapii.so": gapiiPath, "process abi": abi.Name}.Bind(ctx)

		// Load the library.
		log.D(ctx, "Loading GAPII library...")
		// Work around for loading libraries in the N previews. See b/29441142.
		j.Class("java.lang.Runtime").Call("getRuntime").Call("doLoad", gapiiPath, nil)
		log.D(ctx, "Library loaded")
		return nil
	})
	if err != nil {
		return log.Err(ctx, err, "loadGAPII")
	}

	return <-connErr
}
<commit_msg>If we cannot find the signature for getClassLoader, try an alternate.<commit_after>// Copyright (C) 2017 Google 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,
// 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 client

import (
	"context"
	"fmt"
	"io"
	"net"
	"reflect"
	"time"

	"github.com/google/gapid/core/app/crash"
	"github.com/google/gapid/core/context/keys"
	"github.com/google/gapid/core/event/task"
	"github.com/google/gapid/core/java/jdbg"
	"github.com/google/gapid/core/java/jdwp"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/android/adb"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/gapidapk"
)

func expect(r io.Reader, expected []byte) error {
	got := make([]byte, len(expected))
	if _, err := io.ReadFull(r, got); err != nil {
		return err
	}
	if !reflect.DeepEqual(expected, got) {
		return fmt.Errorf("Expected %v, got %v", expected, got)
	}
	return nil
}

// waitForOnCreate waits for android.app.Application.onCreate to be called, and
// then suspends the thread.
func waitForOnCreate(ctx context.Context, conn *jdwp.Connection, wakeup jdwp.ThreadID) (*jdwp.EventMethodEntry, error) {
	app, err := conn.GetClassBySignature("Landroid/app/Application;")
	if err != nil {
		return nil, err
	}

	constructor, err := conn.GetClassMethod(app.ClassID(), "<init>", "()V")
	if err != nil {
		return nil, err
	}

	log.I(ctx, "   Waiting for Application.<init>()")
	return conn.WaitForMethodEntry(ctx, app.ClassID(), constructor.ID, wakeup)
}

// waitForVulkanLoad for android.app.ApplicationLoaders.getClassLoader to be called,
// and then suspends the thread.
// This function is what is used to tell the vulkan loader where to search for
// layers.
func waitForVulkanLoad(ctx context.Context, conn *jdwp.Connection) (*jdwp.EventMethodEntry, error) {
	loaders, err := conn.GetClassBySignature("Landroid/app/ApplicationLoaders;")
	if err != nil {
		return nil, err
	}

	getClassLoader, err := conn.GetClassMethod(loaders.ClassID(), "getClassLoader",
		"(Ljava/lang/String;IZLjava/lang/String;Ljava/lang/String;Ljava/lang/ClassLoader;)Ljava/lang/ClassLoader;")
	if err != nil {
		getClassLoader, err = conn.GetClassMethod(loaders.ClassID(), "getClassLoader",
			"(Ljava/lang/String;IZLjava/lang/String;Ljava/lang/String;Ljava/lang/ClassLoader;Ljava/lang/String;)Ljava/lang/ClassLoader;")
		if err != nil {
			return nil, err
		}
	}

	return conn.WaitForMethodEntry(ctx, loaders.ClassID(), getClassLoader.ID, 0)
}

// loadAndConnectViaJDWP connects to the application waiting for a JDWP
// connection with the specified process id, sends a number of JDWP commands to
// load the list of libraries.
func (p *Process) loadAndConnectViaJDWP(
	ctx context.Context,
	gapidAPK *gapidapk.APK,
	pid int,
	d adb.Device) error {

	const (
		reconnectAttempts = 10
		reconnectDelay    = time.Second
	)

	jdwpPort, err := adb.LocalFreeTCPPort()
	if err != nil {
		return log.Err(ctx, err, "Finding free port")
	}
	ctx = log.V{"jdwpPort": jdwpPort}.Bind(ctx)

	log.I(ctx, "Forwarding TCP port %v -> JDWP pid %v", jdwpPort, pid)
	if err := d.Forward(ctx, adb.TCPPort(jdwpPort), adb.Jdwp(pid)); err != nil {
		return log.Err(ctx, err, "Setting up JDWP port forwarding")
	}
	defer func() {
		// Clone context to ignore cancellation.
		ctx := keys.Clone(context.Background(), ctx)
		d.RemoveForward(ctx, adb.TCPPort(jdwpPort))
	}()

	ctx, stop := task.WithCancel(ctx)
	defer stop()

	log.I(ctx, "Connecting to JDWP")

	// Create a JDWP connection with the application.
	var sock net.Conn
	var conn *jdwp.Connection
	err = task.Retry(ctx, reconnectAttempts, reconnectDelay, func(ctx context.Context) (bool, error) {
		if sock, err = net.Dial("tcp", fmt.Sprintf("localhost:%v", jdwpPort)); err != nil {
			return false, err
		}
		if conn, err = jdwp.Open(ctx, sock); err != nil {
			sock.Close()
			return false, err
		}
		return true, nil
	})
	if err != nil {
		if err == io.EOF {
			return fmt.Errorf("Unable to connect to the application.\n\n" +
				"This can happen when another debugger or IDE is running " +
				"in the background, such as Android Studio.\n" +
				"Please close any running Android debuggers and try again.\n\n" +
				"See https://github.com/google/gapid/issues/911 for more " +
				"information")
		}
		return log.Err(ctx, err, "Connecting to JDWP")
	}
	defer sock.Close()

	processABI := func(j *jdbg.JDbg) (*device.ABI, error) {
		abiName := j.Class("android.os.Build").Field("CPU_ABI").Get().(string)
		abi := device.ABIByName(abiName)
		if abi == nil {
			return nil, fmt.Errorf("Unknown ABI %v", abiName)
		}

		// For NativeBridge emulated devices opt for the native ABI of the
		// emulator.
		abi = d.NativeBridgeABI(ctx, abi)

		return abi, nil
	}

	classLoaderThread := jdwp.ThreadID(0)

	log.I(ctx, "Waiting for ApplicationLoaders.getClassLoader()")
	getClassLoader, err := waitForVulkanLoad(ctx, conn)
	if err == nil {
		// If err != nil that means we could not find or break in getClassLoader
		// so we have no vulkan support.
		classLoaderThread = getClassLoader.Thread
		err = jdbg.Do(conn, getClassLoader.Thread, func(j *jdbg.JDbg) error {
			abi, err := processABI(j)
			if err != nil {
				return err
			}
			libsPath := gapidAPK.LibsPath(abi)
			newLibraryPath := j.String(":" + libsPath)
			obj := j.GetStackObject("librarySearchPath").Call("concat", newLibraryPath)
			j.SetStackObject("librarySearchPath", obj)
			return nil
		})
		if err != nil {
			return log.Err(ctx, err, "JDWP failure")
		}
	} else {
		log.W(ctx, "Couldn't break in ApplicationLoaders.getClassLoader. Vulkan will not be supported.")
	}

	// Wait for Application.onCreate to be called.
	log.I(ctx, "Waiting for Application Creation")
	onCreate, err := waitForOnCreate(ctx, conn, classLoaderThread)
	if err != nil {
		return log.Err(ctx, err, "Waiting for Application Creation")
	}

	// Attempt to get the GVR library handle.
	// Will throw an exception for non-GVR apps.
	var gvrHandle uint64
	log.I(ctx, "Installing interceptor libraries")
	loadNativeGvrLibrary, vrCoreLibraryLoader := "loadNativeGvrLibrary", "com/google/vr/cardboard/VrCoreLibraryLoader"
	gvrMajor, gvrMinor, gvrPoint := 1, 8, 1

	getGVRHandle := func(j *jdbg.JDbg, libLoader jdbg.Type) error {
		// loadNativeGvrLibrary has a couple of different signatures depending
		// on GVR release.
		for _, f := range []func() error{
			// loadNativeGvrLibrary(Context, int major, int minor, int point)
			func() error {
				gvrHandle = (uint64)(libLoader.Call(loadNativeGvrLibrary, j.This(), gvrMajor, gvrMinor, gvrPoint).Get().(int64))
				return nil
			},
			// loadNativeGvrLibrary(Context)
			func() error {
				gvrHandle = (uint64)(libLoader.Call(loadNativeGvrLibrary, j.This()).Get().(int64))
				return nil
			},
		} {
			if jdbg.Try(f) == nil {
				return nil
			}
		}
		return fmt.Errorf("Couldn't call loadNativeGvrLibrary")
	}
	for _, f := range []func(j *jdbg.JDbg) error{
		func(j *jdbg.JDbg) error {
			libLoader := j.Class(vrCoreLibraryLoader)
			getGVRHandle(j, libLoader)
			return nil
		},
		func(j *jdbg.JDbg) error {
			classLoader := j.This().Call("getClassLoader")
			libLoader := classLoader.Call("findClass", vrCoreLibraryLoader).AsType()
			getGVRHandle(j, libLoader)
			return nil
		},
	} {
		if err := jdbg.Do(conn, onCreate.Thread, f); err == nil {
			break
		}
	}
	if gvrHandle == 0 {
		log.I(ctx, "GVR library not found")
	} else {
		log.I(ctx, "GVR library found")
	}

	// Connect to GAPII.
	// This has to be done on a separate go-routine as the call to load gapii
	// will block until a connection is made.
	connErr := make(chan error)

	// Load GAPII library.
	err = jdbg.Do(conn, onCreate.Thread, func(j *jdbg.JDbg) error {
		abi, err := processABI(j)
		if err != nil {
			return err
		}

		interceptorPath := gapidAPK.LibInterceptorPath(abi)
		crash.Go(func() { connErr <- p.connect(ctx, gvrHandle, interceptorPath) })

		gapiiPath := gapidAPK.LibGAPIIPath(abi)
		ctx = log.V{"gapii.so": gapiiPath, "process abi": abi.Name}.Bind(ctx)

		// Load the library.
		log.D(ctx, "Loading GAPII library...")
		// Work around for loading libraries in the N previews. See b/29441142.
		j.Class("java.lang.Runtime").Call("getRuntime").Call("doLoad", gapiiPath, nil)
		log.D(ctx, "Library loaded")
		return nil
	})
	if err != nil {
		return log.Err(ctx, err, "loadGAPII")
	}

	return <-connErr
}
<|endoftext|>


<commit_before>// Package imdb implements Title find and information using AppEngine JSON API.
package imdb

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

type Result struct {
	Id, Name string
	Year     int
}

type Title struct {
	Id, Name, Type, Rating, Duration, Description, Poster string
	Year, Year_production, Year_release                   int
	Aka, Genres, Languages, Nationalities                 []string
	Directors, Writers, Actors                            []Name
}

type Name struct {
	Id, Name string
}

func (t *Title) String() string {
	var infos []string
	name := t.Name
	if t.Year != 0 {
		name = fmt.Sprintf("%s (%d)", name, t.Year)
	}
	infos = append(infos, name)
	if len(t.Genres) > 0 {
		infos = append(infos, strings.Join(t.Genres[:3], ", "))
	}
	if len(t.Directors) > 0 {
		max := len(t.Directors)
		if max > 3 {
			max = 3
		}
		var directors []string
		for _, director := range t.Directors {
			directors = append(directors, director.String())
		}
		infos = append(infos, strings.Join(directors, ", "))
	}
	if len(t.Actors) > 0 {
		max := len(t.Actors)
		if max > 3 {
			max = 3
		}
		var actors []string
		for _, actor := range t.Actors[:max] {
			actors = append(actors, actor.String())
		}
		infos = append(infos, strings.Join(actors, ", "))
	}
	if t.Duration != "" {
		infos = append(infos, t.Duration)
	}
	if t.Rating != "" {
		infos = append(infos, t.Rating)
	}
	infos = append(infos, fmt.Sprintf("http://www.imdb.com/title/%s", t.Id))
	infos = append(infos, "tg")
	return strings.Join(infos, " - ")
}

func (n *Name) String() string {
	return n.Name
}

// NewTitle obtains a Title ID with its information and returns a Title.
func NewTitle(id string) (t Title, e error) {
	base := "http://movie-db-api.appspot.com/title"
	resp, err := http.Get(fmt.Sprintf("%s/%s", base, id))
	if err != nil {
		return t, err
	}
	defer resp.Body.Close()
	contents, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return t, err
	}
	err = json.Unmarshal(contents, &t)
	if err != nil {
		return t, err
	}
	return t, nil
}

// FindTitle searches a Title and returns a list of Result.
func FindTitle(q string) (r []Result, e error) {
	base := "http://movie-db-api.appspot.com/find"
	params := url.Values{}
	params.Set("s", "tt")
	params.Set("q", q)
	resp, err := http.Get(fmt.Sprintf("%s?%s", base, params.Encode()))
	if err != nil {
		return r, err
	}
	defer resp.Body.Close()
	contents, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return r, err
	}
	err = json.Unmarshal(contents, &r)
	if err != nil {
		return r, err
	}
	return r, nil
}
<commit_msg>imdb: cleanup<commit_after>// Package imdb implements Title find and information using AppEngine JSON API.
package imdb

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

type Result struct {
	Id, Name string
	Year     int
}

type Title struct {
	Id, Name, Type, Rating, Duration, Description, Poster string
	Year, Year_production, Year_release                   int
	Aka, Genres, Languages, Nationalities                 []string
	Directors, Writers, Actors                            []Name
}

type Name struct {
	Id, Name string
}

func (t *Title) String() string {
	var infos []string
	name := t.Name
	if t.Year != 0 {
		name = fmt.Sprintf("%s (%d)", name, t.Year)
	}
	infos = append(infos, name)
	if len(t.Genres) > 0 {
		infos = append(infos, strings.Join(t.Genres[:3], ", "))
	}
	if len(t.Directors) > 0 {
		max := len(t.Directors)
		if max > 3 {
			max = 3
		}
		var directors []string
		for _, director := range t.Directors {
			directors = append(directors, director.String())
		}
		infos = append(infos, strings.Join(directors, ", "))
	}
	if len(t.Actors) > 0 {
		max := len(t.Actors)
		if max > 3 {
			max = 3
		}
		var actors []string
		for _, actor := range t.Actors[:max] {
			actors = append(actors, actor.String())
		}
		infos = append(infos, strings.Join(actors, ", "))
	}
	if t.Duration != "" {
		infos = append(infos, t.Duration)
	}
	if t.Rating != "" {
		infos = append(infos, t.Rating)
	}
	infos = append(infos, fmt.Sprintf("http://www.imdb.com/title/%s", t.Id))
	return strings.Join(infos, " - ")
}

func (n *Name) String() string {
	return n.Name
}

// NewTitle obtains a Title ID with its information and returns a Title.
func NewTitle(id string) (t Title, e error) {
	base := "http://movie-db-api.appspot.com/title"
	resp, err := http.Get(fmt.Sprintf("%s/%s", base, id))
	if err != nil {
		return t, err
	}
	defer resp.Body.Close()
	contents, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return t, err
	}
	err = json.Unmarshal(contents, &t)
	if err != nil {
		return t, err
	}
	return t, nil
}

// FindTitle searches a Title and returns a list of Result.
func FindTitle(q string) (r []Result, e error) {
	base := "http://movie-db-api.appspot.com/find"
	params := url.Values{}
	params.Set("s", "tt")
	params.Set("q", q)
	resp, err := http.Get(fmt.Sprintf("%s?%s", base, params.Encode()))
	if err != nil {
		return r, err
	}
	defer resp.Body.Close()
	contents, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return r, err
	}
	err = json.Unmarshal(contents, &r)
	if err != nil {
		return r, err
	}
	return r, nil
}
<|endoftext|>


<commit_before>package kubernetesprovider

import (
	"context"
	"errors"

	detector "github.com/rancher/kubernetes-provider-detector"
	v3 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	v32 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/wrangler"
	"k8s.io/client-go/kubernetes"
)

type handler struct {
	ctx                context.Context
	clusters           v32.ClusterClient
	localClusterClient kubernetes.Interface
	mcm                wrangler.MultiClusterManager
}

func Register(ctx context.Context,
	clusters v32.ClusterController,
	localClusterClient kubernetes.Interface,
	mcm wrangler.MultiClusterManager,
) {
	h := &handler{
		ctx:                ctx,
		clusters:           clusters,
		localClusterClient: localClusterClient,
		mcm:                mcm,
	}
	clusters.OnChange(ctx, "kubernetes-provider", h.OnChange)
}

func (h *handler) OnChange(key string, cluster *v3.Cluster) (*v3.Cluster, error) {
	if cluster == nil || cluster.Status.Provider != "" {
		return cluster, nil
	}

	var client kubernetes.Interface
	if cluster.Spec.Internal {
		client = h.localClusterClient
	} else if k8s, err := h.mcm.K8sClient(cluster.Name); err != nil {
		return nil, err
	} else if k8s != nil {
		client = k8s
	}

	if client == nil {
		return cluster, nil
	}

	provider, err := detector.DetectProvider(h.ctx, client)
	var u detector.ErrUnknownProvider
	if errors.Is(err, &u) {
		return cluster, nil
	} else if err != nil {
		return cluster, err
	}
	cluster = cluster.DeepCopy()
	cluster.Status.Provider = provider
	return h.clusters.Update(cluster)
}
<commit_msg>Don't fail getting the kubernetes provider before the cluster is ready<commit_after>package kubernetesprovider

import (
	"context"
	"errors"

	detector "github.com/rancher/kubernetes-provider-detector"
	v3 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	v32 "github.com/rancher/rancher/pkg/generated/controllers/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/wrangler"
	"k8s.io/client-go/kubernetes"
)

type handler struct {
	ctx                context.Context
	clusters           v32.ClusterClient
	localClusterClient kubernetes.Interface
	mcm                wrangler.MultiClusterManager
}

func Register(ctx context.Context,
	clusters v32.ClusterController,
	localClusterClient kubernetes.Interface,
	mcm wrangler.MultiClusterManager,
) {
	h := &handler{
		ctx:                ctx,
		clusters:           clusters,
		localClusterClient: localClusterClient,
		mcm:                mcm,
	}
	clusters.OnChange(ctx, "kubernetes-provider", h.OnChange)
}

func (h *handler) OnChange(key string, cluster *v3.Cluster) (*v3.Cluster, error) {
	if cluster == nil || cluster.Status.Provider != "" {
		return cluster, nil
	}

	if !v3.ClusterConditionReady.IsTrue(cluster) {
		return cluster, nil
	}

	var client kubernetes.Interface
	if cluster.Spec.Internal {
		client = h.localClusterClient
	} else if k8s, err := h.mcm.K8sClient(cluster.Name); err != nil {
		// ignore error. If we can't get a client just ignore it. The cluster probably isn't happy
		// yet and we will get an update later when it is.
		return nil, nil
	} else if k8s != nil {
		client = k8s
	}

	if client == nil {
		return cluster, nil
	}

	provider, err := detector.DetectProvider(h.ctx, client)
	var u detector.ErrUnknownProvider
	if errors.Is(err, &u) {
		return cluster, nil
	} else if err != nil {
		return cluster, err
	}
	cluster = cluster.DeepCopy()
	cluster.Status.Provider = provider
	return h.clusters.Update(cluster)
}
<|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 sqlparser

import (
	"fmt"
	"strconv"
	"strings"
	"unicode"
)

const (
	// DirectiveMultiShardAutocommit is the query comment directive to allow
	// single round trip autocommit with a multi-shard statement.
	DirectiveMultiShardAutocommit = "MULTI_SHARD_AUTOCOMMIT"
	// DirectiveSkipQueryPlanCache skips query plan cache when set.
	DirectiveSkipQueryPlanCache = "SKIP_QUERY_PLAN_CACHE"
	// DirectiveQueryTimeout sets a query timeout in vtgate. Only supported for SELECTS.
	DirectiveQueryTimeout = "QUERY_TIMEOUT_MS"
	// DirectiveScatterErrorsAsWarnings enables partial success scatter select queries
	DirectiveScatterErrorsAsWarnings = "SCATTER_ERRORS_AS_WARNINGS"
	// DirectiveIgnoreMaxPayloadSize skips payload size validation when set.
	DirectiveIgnoreMaxPayloadSize = "IGNORE_MAX_PAYLOAD_SIZE"
	// DirectiveIgnoreMaxMemoryRows skips memory row validation when set.
	DirectiveIgnoreMaxMemoryRows = "IGNORE_MAX_MEMORY_ROWS"
)

func isNonSpace(r rune) bool {
	return !unicode.IsSpace(r)
}

// leadingCommentEnd returns the first index after all leading comments, or
// 0 if there are no leading comments.
func leadingCommentEnd(text string) (end int) {
	hasComment := false
	pos := 0
	for pos < len(text) {
		// Eat up any whitespace. Trailing whitespace will be considered part of
		// the leading comments.
		nextVisibleOffset := strings.IndexFunc(text[pos:], isNonSpace)
		if nextVisibleOffset < 0 {
			break
		}
		pos += nextVisibleOffset
		remainingText := text[pos:]

		// Found visible characters. Look for '/*' at the beginning
		// and '*/' somewhere after that.
		if len(remainingText) < 4 || remainingText[:2] != "/*" || remainingText[2] == '!' {
			break
		}
		commentLength := 4 + strings.Index(remainingText[2:], "*/")
		if commentLength < 4 {
			// Missing end comment :/
			break
		}

		hasComment = true
		pos += commentLength
	}

	if hasComment {
		return pos
	}
	return 0
}

// trailingCommentStart returns the first index of trailing comments.
// If there are no trailing comments, returns the length of the input string.
func trailingCommentStart(text string) (start int) {
	hasComment := false
	reducedLen := len(text)
	for reducedLen > 0 {
		// Eat up any whitespace. Leading whitespace will be considered part of
		// the trailing comments.
		nextReducedLen := strings.LastIndexFunc(text[:reducedLen], isNonSpace) + 1
		if nextReducedLen == 0 {
			break
		}
		reducedLen = nextReducedLen
		if reducedLen < 4 || text[reducedLen-2:reducedLen] != "*/" {
			break
		}

		// Find the beginning of the comment
		startCommentPos := strings.LastIndex(text[:reducedLen-2], "/*")
		if startCommentPos < 0 || text[startCommentPos+2] == '!' {
			// Badly formatted sql, or a special /*! comment
			break
		}

		hasComment = true
		reducedLen = startCommentPos
	}

	if hasComment {
		return reducedLen
	}
	return len(text)
}

// MarginComments holds the leading and trailing comments that surround a query.
type MarginComments struct {
	Leading  string
	Trailing string
}

// SplitMarginComments pulls out any leading or trailing comments from a raw sql query.
// This function also trims leading (if there's a comment) and trailing whitespace.
func SplitMarginComments(sql string) (query string, comments MarginComments) {
	trailingStart := trailingCommentStart(sql)
	leadingEnd := leadingCommentEnd(sql[:trailingStart])
	comments = MarginComments{
		Leading:  strings.TrimLeftFunc(sql[:leadingEnd], unicode.IsSpace),
		Trailing: strings.TrimRightFunc(sql[trailingStart:], unicode.IsSpace),
	}
	return strings.TrimFunc(sql[leadingEnd:trailingStart], func(c rune) bool {
		return unicode.IsSpace(c) || c == ';'
	}), comments
}

// StripLeadingComments trims the SQL string and removes any leading comments
func StripLeadingComments(sql string) string {
	sql = strings.TrimFunc(sql, unicode.IsSpace)

	for hasCommentPrefix(sql) {
		switch sql[0] {
		case '/':
			// Multi line comment
			index := strings.Index(sql, "*/")
			if index <= 1 {
				return sql
			}
			// don't strip /*! ... */ or /*!50700 ... */
			if len(sql) > 2 && sql[2] == '!' {
				return sql
			}
			sql = sql[index+2:]
		case '-':
			// Single line comment
			index := strings.Index(sql, "\n")
			if index == -1 {
				return ""
			}
			sql = sql[index+1:]
		}

		sql = strings.TrimFunc(sql, unicode.IsSpace)
	}

	return sql
}

func hasCommentPrefix(sql string) bool {
	return len(sql) > 1 && ((sql[0] == '/' && sql[1] == '*') || (sql[0] == '-' && sql[1] == '-'))
}

// ExtractMysqlComment extracts the version and SQL from a comment-only query
// such as /*!50708 sql here */
func ExtractMysqlComment(sql string) (string, string) {
	sql = sql[3 : len(sql)-2]

	digitCount := 0
	endOfVersionIndex := strings.IndexFunc(sql, func(c rune) bool {
		digitCount++
		return !unicode.IsDigit(c) || digitCount == 6
	})
	if endOfVersionIndex < 0 {
		return "", ""
	}
	if endOfVersionIndex < 5 {
		endOfVersionIndex = 0
	}
	version := sql[0:endOfVersionIndex]
	innerSQL := strings.TrimFunc(sql[endOfVersionIndex:], unicode.IsSpace)

	return version, innerSQL
}

const commentDirectivePreamble = "/*vt+"

// CommentDirectives is the parsed representation for execution directives
// conveyed in query comments
type CommentDirectives map[string]interface{}

// ExtractCommentDirectives parses the comment list for any execution directives
// of the form:
//
//     /*vt+ OPTION_ONE=1 OPTION_TWO OPTION_THREE=abcd */
//
// It returns the map of the directive values or nil if there aren't any.
func ExtractCommentDirectives(comments Comments) CommentDirectives {
	if comments == nil {
		return nil
	}

	var vals map[string]interface{}

	for _, comment := range comments {
		commentStr := string(comment)
		if commentStr[0:5] != commentDirectivePreamble {
			continue
		}

		if vals == nil {
			vals = make(map[string]interface{})
		}

		// Split on whitespace and ignore the first and last directive
		// since they contain the comment start/end
		directives := strings.Fields(commentStr)
		for i := 1; i < len(directives)-1; i++ {
			directive := directives[i]
			sep := strings.IndexByte(directive, '=')

			// No value is equivalent to a true boolean
			if sep == -1 {
				vals[directive] = true
				continue
			}

			strVal := directive[sep+1:]
			directive = directive[:sep]

			intVal, err := strconv.Atoi(strVal)
			if err == nil {
				vals[directive] = intVal
				continue
			}

			boolVal, err := strconv.ParseBool(strVal)
			if err == nil {
				vals[directive] = boolVal
				continue
			}

			vals[directive] = strVal
		}
	}
	return vals
}

// IsSet checks the directive map for the named directive and returns
// true if the directive is set and has a true/false or 0/1 value
func (d CommentDirectives) IsSet(key string) bool {
	if d == nil {
		return false
	}

	val, ok := d[key]
	if !ok {
		return false
	}

	boolVal, ok := val.(bool)
	if ok {
		return boolVal
	}

	intVal, ok := val.(int)
	if ok {
		return intVal == 1
	}
	return false
}

// GetString gets a directive value as string, with default value if not found
func (d CommentDirectives) GetString(key string, defaultVal string) string {
	val, ok := d[key]
	if !ok {
		return defaultVal
	}
	stringVal := fmt.Sprintf("%v", val)
	if unquoted, err := strconv.Unquote(stringVal); err == nil {
		stringVal = unquoted
	}
	return stringVal
}

// SkipQueryPlanCacheDirective returns true if skip query plan cache directive is set to true in query.
func SkipQueryPlanCacheDirective(stmt Statement) bool {
	switch stmt := stmt.(type) {
	case *Select:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveSkipQueryPlanCache) {
			return true
		}
	case *Insert:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveSkipQueryPlanCache) {
			return true
		}
	case *Update:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveSkipQueryPlanCache) {
			return true
		}
	case *Delete:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveSkipQueryPlanCache) {
			return true
		}
	default:
		return false
	}
	return false
}

// IgnoreMaxPayloadSizeDirective returns true if the max payload size override
// directive is set to true.
func IgnoreMaxPayloadSizeDirective(stmt Statement) bool {
	switch stmt := stmt.(type) {
	case *Select:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxPayloadSize)
	case *Insert:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxPayloadSize)
	case *Update:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxPayloadSize)
	case *Delete:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxPayloadSize)
	default:
		return false
	}
}

// IgnoreMaxMaxMemoryRowsDirective returns true if the max memory rows override
// directive is set to true.
func IgnoreMaxMaxMemoryRowsDirective(stmt Statement) bool {
	switch stmt := stmt.(type) {
	case *Select:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
	case *Insert:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
	case *Update:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
	case *Delete:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
	default:
		return false
	}
}
<commit_msg>add MultiShardAutocommitDirective utility<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 sqlparser

import (
	"fmt"
	"strconv"
	"strings"
	"unicode"
)

const (
	// DirectiveMultiShardAutocommit is the query comment directive to allow
	// single round trip autocommit with a multi-shard statement.
	DirectiveMultiShardAutocommit = "MULTI_SHARD_AUTOCOMMIT"
	// DirectiveSkipQueryPlanCache skips query plan cache when set.
	DirectiveSkipQueryPlanCache = "SKIP_QUERY_PLAN_CACHE"
	// DirectiveQueryTimeout sets a query timeout in vtgate. Only supported for SELECTS.
	DirectiveQueryTimeout = "QUERY_TIMEOUT_MS"
	// DirectiveScatterErrorsAsWarnings enables partial success scatter select queries
	DirectiveScatterErrorsAsWarnings = "SCATTER_ERRORS_AS_WARNINGS"
	// DirectiveIgnoreMaxPayloadSize skips payload size validation when set.
	DirectiveIgnoreMaxPayloadSize = "IGNORE_MAX_PAYLOAD_SIZE"
	// DirectiveIgnoreMaxMemoryRows skips memory row validation when set.
	DirectiveIgnoreMaxMemoryRows = "IGNORE_MAX_MEMORY_ROWS"
)

func isNonSpace(r rune) bool {
	return !unicode.IsSpace(r)
}

// leadingCommentEnd returns the first index after all leading comments, or
// 0 if there are no leading comments.
func leadingCommentEnd(text string) (end int) {
	hasComment := false
	pos := 0
	for pos < len(text) {
		// Eat up any whitespace. Trailing whitespace will be considered part of
		// the leading comments.
		nextVisibleOffset := strings.IndexFunc(text[pos:], isNonSpace)
		if nextVisibleOffset < 0 {
			break
		}
		pos += nextVisibleOffset
		remainingText := text[pos:]

		// Found visible characters. Look for '/*' at the beginning
		// and '*/' somewhere after that.
		if len(remainingText) < 4 || remainingText[:2] != "/*" || remainingText[2] == '!' {
			break
		}
		commentLength := 4 + strings.Index(remainingText[2:], "*/")
		if commentLength < 4 {
			// Missing end comment :/
			break
		}

		hasComment = true
		pos += commentLength
	}

	if hasComment {
		return pos
	}
	return 0
}

// trailingCommentStart returns the first index of trailing comments.
// If there are no trailing comments, returns the length of the input string.
func trailingCommentStart(text string) (start int) {
	hasComment := false
	reducedLen := len(text)
	for reducedLen > 0 {
		// Eat up any whitespace. Leading whitespace will be considered part of
		// the trailing comments.
		nextReducedLen := strings.LastIndexFunc(text[:reducedLen], isNonSpace) + 1
		if nextReducedLen == 0 {
			break
		}
		reducedLen = nextReducedLen
		if reducedLen < 4 || text[reducedLen-2:reducedLen] != "*/" {
			break
		}

		// Find the beginning of the comment
		startCommentPos := strings.LastIndex(text[:reducedLen-2], "/*")
		if startCommentPos < 0 || text[startCommentPos+2] == '!' {
			// Badly formatted sql, or a special /*! comment
			break
		}

		hasComment = true
		reducedLen = startCommentPos
	}

	if hasComment {
		return reducedLen
	}
	return len(text)
}

// MarginComments holds the leading and trailing comments that surround a query.
type MarginComments struct {
	Leading  string
	Trailing string
}

// SplitMarginComments pulls out any leading or trailing comments from a raw sql query.
// This function also trims leading (if there's a comment) and trailing whitespace.
func SplitMarginComments(sql string) (query string, comments MarginComments) {
	trailingStart := trailingCommentStart(sql)
	leadingEnd := leadingCommentEnd(sql[:trailingStart])
	comments = MarginComments{
		Leading:  strings.TrimLeftFunc(sql[:leadingEnd], unicode.IsSpace),
		Trailing: strings.TrimRightFunc(sql[trailingStart:], unicode.IsSpace),
	}
	return strings.TrimFunc(sql[leadingEnd:trailingStart], func(c rune) bool {
		return unicode.IsSpace(c) || c == ';'
	}), comments
}

// StripLeadingComments trims the SQL string and removes any leading comments
func StripLeadingComments(sql string) string {
	sql = strings.TrimFunc(sql, unicode.IsSpace)

	for hasCommentPrefix(sql) {
		switch sql[0] {
		case '/':
			// Multi line comment
			index := strings.Index(sql, "*/")
			if index <= 1 {
				return sql
			}
			// don't strip /*! ... */ or /*!50700 ... */
			if len(sql) > 2 && sql[2] == '!' {
				return sql
			}
			sql = sql[index+2:]
		case '-':
			// Single line comment
			index := strings.Index(sql, "\n")
			if index == -1 {
				return ""
			}
			sql = sql[index+1:]
		}

		sql = strings.TrimFunc(sql, unicode.IsSpace)
	}

	return sql
}

func hasCommentPrefix(sql string) bool {
	return len(sql) > 1 && ((sql[0] == '/' && sql[1] == '*') || (sql[0] == '-' && sql[1] == '-'))
}

// ExtractMysqlComment extracts the version and SQL from a comment-only query
// such as /*!50708 sql here */
func ExtractMysqlComment(sql string) (string, string) {
	sql = sql[3 : len(sql)-2]

	digitCount := 0
	endOfVersionIndex := strings.IndexFunc(sql, func(c rune) bool {
		digitCount++
		return !unicode.IsDigit(c) || digitCount == 6
	})
	if endOfVersionIndex < 0 {
		return "", ""
	}
	if endOfVersionIndex < 5 {
		endOfVersionIndex = 0
	}
	version := sql[0:endOfVersionIndex]
	innerSQL := strings.TrimFunc(sql[endOfVersionIndex:], unicode.IsSpace)

	return version, innerSQL
}

const commentDirectivePreamble = "/*vt+"

// CommentDirectives is the parsed representation for execution directives
// conveyed in query comments
type CommentDirectives map[string]interface{}

// ExtractCommentDirectives parses the comment list for any execution directives
// of the form:
//
//     /*vt+ OPTION_ONE=1 OPTION_TWO OPTION_THREE=abcd */
//
// It returns the map of the directive values or nil if there aren't any.
func ExtractCommentDirectives(comments Comments) CommentDirectives {
	if comments == nil {
		return nil
	}

	var vals map[string]interface{}

	for _, comment := range comments {
		commentStr := string(comment)
		if commentStr[0:5] != commentDirectivePreamble {
			continue
		}

		if vals == nil {
			vals = make(map[string]interface{})
		}

		// Split on whitespace and ignore the first and last directive
		// since they contain the comment start/end
		directives := strings.Fields(commentStr)
		for i := 1; i < len(directives)-1; i++ {
			directive := directives[i]
			sep := strings.IndexByte(directive, '=')

			// No value is equivalent to a true boolean
			if sep == -1 {
				vals[directive] = true
				continue
			}

			strVal := directive[sep+1:]
			directive = directive[:sep]

			intVal, err := strconv.Atoi(strVal)
			if err == nil {
				vals[directive] = intVal
				continue
			}

			boolVal, err := strconv.ParseBool(strVal)
			if err == nil {
				vals[directive] = boolVal
				continue
			}

			vals[directive] = strVal
		}
	}
	return vals
}

// IsSet checks the directive map for the named directive and returns
// true if the directive is set and has a true/false or 0/1 value
func (d CommentDirectives) IsSet(key string) bool {
	if d == nil {
		return false
	}

	val, ok := d[key]
	if !ok {
		return false
	}

	boolVal, ok := val.(bool)
	if ok {
		return boolVal
	}

	intVal, ok := val.(int)
	if ok {
		return intVal == 1
	}
	return false
}

// GetString gets a directive value as string, with default value if not found
func (d CommentDirectives) GetString(key string, defaultVal string) string {
	val, ok := d[key]
	if !ok {
		return defaultVal
	}
	stringVal := fmt.Sprintf("%v", val)
	if unquoted, err := strconv.Unquote(stringVal); err == nil {
		stringVal = unquoted
	}
	return stringVal
}

// MultiShardAutocommitDirective returns true if multishard autocommit directive is set to true in query.
func MultiShardAutocommitDirective(stmt Statement) bool {
	switch stmt := stmt.(type) {
	case *Insert:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveMultiShardAutocommit) {
			return true
		}
	case *Update:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveMultiShardAutocommit) {
			return true
		}
	case *Delete:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveMultiShardAutocommit) {
			return true
		}
	default:
		return false
	}
	return false
}

// SkipQueryPlanCacheDirective returns true if skip query plan cache directive is set to true in query.
func SkipQueryPlanCacheDirective(stmt Statement) bool {
	switch stmt := stmt.(type) {
	case *Select:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveSkipQueryPlanCache) {
			return true
		}
	case *Insert:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveSkipQueryPlanCache) {
			return true
		}
	case *Update:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveSkipQueryPlanCache) {
			return true
		}
	case *Delete:
		directives := ExtractCommentDirectives(stmt.Comments)
		if directives.IsSet(DirectiveSkipQueryPlanCache) {
			return true
		}
	default:
		return false
	}
	return false
}

// IgnoreMaxPayloadSizeDirective returns true if the max payload size override
// directive is set to true.
func IgnoreMaxPayloadSizeDirective(stmt Statement) bool {
	switch stmt := stmt.(type) {
	case *Select:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxPayloadSize)
	case *Insert:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxPayloadSize)
	case *Update:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxPayloadSize)
	case *Delete:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxPayloadSize)
	default:
		return false
	}
}

// IgnoreMaxMaxMemoryRowsDirective returns true if the max memory rows override
// directive is set to true.
func IgnoreMaxMaxMemoryRowsDirective(stmt Statement) bool {
	switch stmt := stmt.(type) {
	case *Select:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
	case *Insert:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
	case *Update:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
	case *Delete:
		directives := ExtractCommentDirectives(stmt.Comments)
		return directives.IsSet(DirectiveIgnoreMaxMemoryRows)
	default:
		return false
	}
}
<|endoftext|>


<commit_before>package typed

// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

import (
	"bytes"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSimple(t *testing.T) {
	buf := make([]byte, 200)

	var r ReadBuffer
	var w WriteBuffer

	{
		w.Wrap(buf)
		w.WriteByte(0xFC)
		r.Wrap(buf)
		assert.Equal(t, byte(0xFC), r.ReadByte())
	}

	{
		w.Wrap(buf)
		w.WriteUint16(0xDEAD)
		r.Wrap(buf)
		assert.Equal(t, uint16(0xDEAD), r.ReadUint16())
	}

	{
		w.Wrap(buf)
		w.WriteUint32(0xBEEFDEAD)
		r.Wrap(buf)
		assert.Equal(t, uint32(0xBEEFDEAD), r.ReadUint32())
	}

}

func TestReadWrite(t *testing.T) {
	s := "the small brown fix"
	bslice := []byte("jumped over the lazy dog")

	w := NewWriteBufferWithSize(1024)
	w.WriteUint64(0x0123456789ABCDEF)
	w.WriteUint32(0xABCDEF01)
	w.WriteUint16(0x2345)
	w.WriteByte(0xFF)
	w.WriteString(s)
	w.WriteBytes(bslice)
	w.WriteLen8String("hello")
	w.WriteLen16String("This is a much larger string")
	require.NoError(t, w.Err())

	var b bytes.Buffer
	w.FlushTo(&b)

	r := NewReadBufferWithSize(1024)
	r.FillFrom(bytes.NewReader(b.Bytes()), len(b.Bytes()))

	assert.Equal(t, uint64(0x0123456789ABCDEF), r.ReadUint64())
	assert.Equal(t, uint32(0xABCDEF01), r.ReadUint32())
	assert.Equal(t, uint16(0x2345), r.ReadUint16())
	assert.Equal(t, byte(0xFF), r.ReadByte())
	assert.Equal(t, s, r.ReadString(len(s)))
	assert.Equal(t, bslice, r.ReadBytes(len(bslice)))
	assert.Equal(t, "hello", r.ReadLen8String())
	assert.Equal(t, "This is a much larger string", r.ReadLen16String())

	require.NoError(t, r.Err())
}

func TestDeferredWrites(t *testing.T) {
	w := NewWriteBufferWithSize(1024)
	u16ref := w.DeferUint16()
	require.NotNil(t, u16ref)

	u32ref := w.DeferUint32()
	require.NotNil(t, u32ref)

	u64ref := w.DeferUint64()
	require.NotNil(t, u64ref)

	bref := w.DeferBytes(5)
	require.NotNil(t, bref)

	sref := w.DeferBytes(5)
	require.NotNil(t, sref)

	byteref := w.DeferByte()
	require.NotNil(t, byteref)

	assert.Equal(t, 2+4+8+5+5+1, w.BytesWritten())

	u16ref.Update(2040)
	u32ref.Update(495404)
	u64ref.Update(0x40950459)
	bref.Update([]byte{0x30, 0x12, 0x45, 0x55, 0x65})
	sref.UpdateString("where")
	byteref.Update(0x44)

	var buf bytes.Buffer
	w.FlushTo(&buf)

	r := NewReadBuffer(buf.Bytes())

	u16 := r.ReadUint16()
	assert.Equal(t, uint16(2040), u16)

	u32 := r.ReadUint32()
	assert.Equal(t, uint32(495404), u32)

	u64 := r.ReadUint64()
	assert.Equal(t, uint32(0x40950459), u64)

	b := r.ReadBytes(5)
	assert.Equal(t, []byte{0x30, 0x12, 0x45, 0x55, 0x65}, b)

	s := r.ReadString(5)
	assert.Equal(t, "where", s)

	u8 := r.ReadByte()
	assert.Equal(t, byte(0x44), u8)
	assert.NoError(t, r.Err())
}
<commit_msg>Fix buffer_test failure caused by uint32 vs uint64 comparison<commit_after>package typed

// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

import (
	"bytes"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestSimple(t *testing.T) {
	buf := make([]byte, 200)

	var r ReadBuffer
	var w WriteBuffer

	{
		w.Wrap(buf)
		w.WriteByte(0xFC)
		r.Wrap(buf)
		assert.Equal(t, byte(0xFC), r.ReadByte())
	}

	{
		w.Wrap(buf)
		w.WriteUint16(0xDEAD)
		r.Wrap(buf)
		assert.Equal(t, uint16(0xDEAD), r.ReadUint16())
	}

	{
		w.Wrap(buf)
		w.WriteUint32(0xBEEFDEAD)
		r.Wrap(buf)
		assert.Equal(t, uint32(0xBEEFDEAD), r.ReadUint32())
	}

}

func TestReadWrite(t *testing.T) {
	s := "the small brown fix"
	bslice := []byte("jumped over the lazy dog")

	w := NewWriteBufferWithSize(1024)
	w.WriteUint64(0x0123456789ABCDEF)
	w.WriteUint32(0xABCDEF01)
	w.WriteUint16(0x2345)
	w.WriteByte(0xFF)
	w.WriteString(s)
	w.WriteBytes(bslice)
	w.WriteLen8String("hello")
	w.WriteLen16String("This is a much larger string")
	require.NoError(t, w.Err())

	var b bytes.Buffer
	w.FlushTo(&b)

	r := NewReadBufferWithSize(1024)
	r.FillFrom(bytes.NewReader(b.Bytes()), len(b.Bytes()))

	assert.Equal(t, uint64(0x0123456789ABCDEF), r.ReadUint64())
	assert.Equal(t, uint32(0xABCDEF01), r.ReadUint32())
	assert.Equal(t, uint16(0x2345), r.ReadUint16())
	assert.Equal(t, byte(0xFF), r.ReadByte())
	assert.Equal(t, s, r.ReadString(len(s)))
	assert.Equal(t, bslice, r.ReadBytes(len(bslice)))
	assert.Equal(t, "hello", r.ReadLen8String())
	assert.Equal(t, "This is a much larger string", r.ReadLen16String())

	require.NoError(t, r.Err())
}

func TestDeferredWrites(t *testing.T) {
	w := NewWriteBufferWithSize(1024)
	u16ref := w.DeferUint16()
	require.NotNil(t, u16ref)

	u32ref := w.DeferUint32()
	require.NotNil(t, u32ref)

	u64ref := w.DeferUint64()
	require.NotNil(t, u64ref)

	bref := w.DeferBytes(5)
	require.NotNil(t, bref)

	sref := w.DeferBytes(5)
	require.NotNil(t, sref)

	byteref := w.DeferByte()
	require.NotNil(t, byteref)

	assert.Equal(t, 2+4+8+5+5+1, w.BytesWritten())

	u16ref.Update(2040)
	u32ref.Update(495404)
	u64ref.Update(0x40950459)
	bref.Update([]byte{0x30, 0x12, 0x45, 0x55, 0x65})
	sref.UpdateString("where")
	byteref.Update(0x44)

	var buf bytes.Buffer
	w.FlushTo(&buf)

	r := NewReadBuffer(buf.Bytes())

	u16 := r.ReadUint16()
	assert.Equal(t, uint16(2040), u16)

	u32 := r.ReadUint32()
	assert.Equal(t, uint32(495404), u32)

	u64 := r.ReadUint64()
	assert.Equal(t, uint64(0x40950459), u64)

	b := r.ReadBytes(5)
	assert.Equal(t, []byte{0x30, 0x12, 0x45, 0x55, 0x65}, b)

	s := r.ReadString(5)
	assert.Equal(t, "where", s)

	u8 := r.ReadByte()
	assert.Equal(t, byte(0x44), u8)
	assert.NoError(t, r.Err())
}
<|endoftext|>


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

// This is a games support
// Ask a question and get an answer support
// Just type "bye" to quit

package main

import (
    "bufio"
    "fmt"
    "math/rand"
    "os"
    "time"
)

func Welcome() {
    fmt.Println("Welcome to the games support.")
    fmt.Println("")
    fmt.Println("Write your question.")
    fmt.Println("We will try to solve your problem.")
    fmt.Println("Type 'bye' for quit the support.")
}

func Goodbye() {
    fmt.Println("We were happy to help. Bye Bye...")
}

type Responses []string

type Support struct {
    Name string
    Resp Responses
}

func NewSupport(name string) *Support {
    return &Support{name, make(Responses, 0)}
}

func (s *Support) AddResponse(str string) *Support {
    s.Resp = append(s.Resp, str)
    return s
}

func (s Support) GetResponses() Responses {
    return s.Resp
}

func (s Support) GenerateResp() string {
    rand.Seed(time.Now().Unix())
    responses := s.Resp
    return responses[rand.Intn(len(responses))]
}

func (s *Support) Init() {
    s.AddResponse("It seems strange. Could you describe the problem more precisely?")
    s.AddResponse("No other customer has ever complained about this.")
    s.AddResponse("What is your system configuration?")
    s.AddResponse("That sounds interesting. Tell me more ...")
    s.AddResponse("This is explained in the manual")
    s.AddResponse("Your description is not clear. Will he an expert you could describe this more precisely?")
    s.AddResponse("This is not a bug but a feature!")
    s.AddResponse("Could you clarify?")
}

func play() {
    support := NewSupport("Games Support")
    support.Init()

    Welcome()
    var input string
    scanner := bufio.NewScanner(os.Stdin)

    for scanner.Scan() {
        input = scanner.Text()
        if input == "bye" {
            Goodbye()
            break
        } else {
            fmt.Println(support.GenerateResp())
        }
    }
}

func main() {
    play()
}
<commit_msg>Improved code<commit_after>// Copyright 2014 The Nanoninja Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This is a games support
// Ask a question and get an answer support
// Just type "bye" to quit

package main

import (
    "bufio"
    "fmt"
    "math/rand"
    "os"
    "time"
)

func Welcome() {
    fmt.Println("Welcome to the games support.")
    fmt.Println("")
    fmt.Println("Write your question.")
    fmt.Println("We will try to solve your problem.")
    fmt.Println("Type 'bye' for quit the support.")
}

func Goodbye() {
    fmt.Println("We were happy to help. Bye Bye...")
}

type Responses []string

type Support struct {
    Name string
    Resp Responses
}

func NewSupport(name string) *Support {
    return &Support{name, make(Responses, 0)}
}

func (s *Support) AddResponse(str string) *Support {
    s.Resp = append(s.Resp, str)
    return s
}

func (s Support) GetResponses() Responses {
    return s.Resp
}

func (s Support) GenerateResp() string {
    rand.Seed(time.Now().Unix())
    responses := s.Resp
    return responses[rand.Intn(len(responses))]
}

func (s *Support) Init() {
    s.AddResponse("It seems strange. Could you describe the problem more precisely?")
    s.AddResponse("No other customer has ever complained about this.")
    s.AddResponse("What is your system configuration?")
    s.AddResponse("That sounds interesting. Tell me more ...")
    s.AddResponse("This is explained in the manual")
    s.AddResponse("Your description is not clear. Will he an expert you could describe this more precisely?")
    s.AddResponse("This is not a bug but a feature!")
    s.AddResponse("Could you clarify?")
}

func (s *Support) play() {
    s.Init()

    Welcome()
    var input string
    scanner := bufio.NewScanner(os.Stdin)

    for scanner.Scan() {
        input = scanner.Text()
        if input == "bye" {
            Goodbye()
            break
        } else {
            fmt.Println(s.GenerateResp())
        }
    }
}

func main() {
    support := NewSupport("Games Support")
    support.play()
}
<|endoftext|>