<commit_before>// Copyright 2015 The rkt 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 image

import (
	"container/list"
	"errors"
	"fmt"
	"net/url"
	"runtime"

	"github.com/coreos/rkt/common"
	"github.com/coreos/rkt/common/apps"
	"github.com/coreos/rkt/stage0"
	"github.com/coreos/rkt/store"
	"github.com/hashicorp/errwrap"

	"github.com/appc/spec/discovery"
	"github.com/appc/spec/schema/types"
)

// Fetcher will try to fetch images into the store.
type Fetcher action

// FetchImage will take an image as either a path, a URL or a name
// string and import it into the store if found. If ascPath is not "",
// it must exist as a local file and will be used as the signature
// file for verification, unless verification is disabled. If
// f.WithDeps is true also image dependencies are fetched.
func (f *Fetcher) FetchImage(img string, ascPath string, imgType apps.AppImageType) (string, error) {
	ensureLogger(f.Debug)
	a := f.getAsc(ascPath)
	hash, err := f.fetchSingleImage(img, a, imgType)
	if err != nil {
		return "", err
	}
	if f.WithDeps {
		err = f.fetchImageDeps(hash)
		if err != nil {
			return "", err
		}
	}
	if common.SupportsOverlay() {
		if _, _, err := f.S.RenderTreeStore(hash, false); err != nil {
			return "", errwrap.Wrap(errors.New("error rendering tree store"), err)
		}
	}
	return hash, nil
}

func (f *Fetcher) getAsc(ascPath string) *asc {
	if ascPath != "" {
		return &asc{
			Location: ascPath,
			Fetcher:  &localAscFetcher{},
		}
	}
	return &asc{}
}

// fetchImageDeps will recursively fetch all the image dependencies
func (f *Fetcher) fetchImageDeps(hash string) error {
	imgsl := list.New()
	seen := map[string]struct{}{}
	f.addImageDeps(hash, imgsl, seen)
	for el := imgsl.Front(); el != nil; el = el.Next() {
		a := &asc{}
		img := el.Value.(string)
		hash, err := f.fetchSingleImage(img, a, apps.AppImageName)
		if err != nil {
			return err
		}
		f.addImageDeps(hash, imgsl, seen)
	}
	return nil
}

func (f *Fetcher) addImageDeps(hash string, imgsl *list.List, seen map[string]struct{}) error {
	dependencies, err := f.getImageDeps(hash)
	if err != nil {
		return errwrap.Wrap(fmt.Errorf("failed to get dependencies for image ID %q", hash), err)
	}
	for _, d := range dependencies {
		imgName := d.ImageName.String()
		app, err := discovery.NewApp(imgName, d.Labels.ToMap())
		if err != nil {
			return errwrap.Wrap(fmt.Errorf("one of image ID's %q dependencies (image %q) is invalid", hash, imgName), err)
		}
		appStr := app.String()
		if _, ok := seen[appStr]; ok {
			continue
		}
		imgsl.PushBack(app.String())
		seen[appStr] = struct{}{}
	}
	return nil
}

func (f *Fetcher) getImageDeps(hash string) (types.Dependencies, error) {
	key, err := f.S.ResolveKey(hash)
	if err != nil {
		return nil, err
	}
	im, err := f.S.GetImageManifest(key)
	if err != nil {
		return nil, err
	}
	return im.Dependencies, nil
}

func (f *Fetcher) fetchSingleImage(img string, a *asc, imgType apps.AppImageType) (string, error) {
	if imgType == apps.AppImageGuess {
		imgType = guessImageType(img)
	}
	if imgType == apps.AppImageHash {
		return "", fmt.Errorf("cannot fetch a hash %q, expected either a URL, a path or an image name", img)
	}

	switch imgType {
	case apps.AppImageURL:
		return f.fetchSingleImageByURL(img, a)
	case apps.AppImagePath:
		return f.fetchSingleImageByPath(img, a)
	case apps.AppImageName:
		return f.fetchSingleImageByName(img, a)
	default:
		return "", fmt.Errorf("unknown image type %d", imgType)
	}
}

type remoteCheck int

const (
	remoteCheckLax remoteCheck = iota
	remoteCheckStrict
)

func (f *Fetcher) fetchSingleImageByURL(urlStr string, a *asc) (string, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return "", errwrap.Wrap(fmt.Errorf("invalid image URL %q", urlStr), err)
	}

	switch u.Scheme {
	case "http", "https":
		return f.fetchSingleImageByHTTPURL(u, a)
	case "docker":
		return f.fetchSingleImageByDockerURL(u)
	case "file":
		return f.fetchSingleImageByPath(u.Path, a)
	case "":
		return "", fmt.Errorf("expected image URL %q to contain a scheme", urlStr)
	default:
		return "", fmt.Errorf("an unsupported URL scheme %q - the only URL schemes supported by rkt are docker, http, https and file", u.Scheme)
	}
}

func (f *Fetcher) fetchSingleImageByHTTPURL(u *url.URL, a *asc) (string, error) {
	rem, err := f.getRemoteForURL(u)
	if err != nil {
		return "", err
	}
	if h := f.maybeCheckRemoteFromStore(rem, remoteCheckStrict); h != "" {
		return h, nil
	}
	if h, err := f.maybeFetchHTTPURLFromRemote(rem, u, a); h != "" || err != nil {
		return h, err
	}
	return "", fmt.Errorf("unable to fetch image from URL %q: either image was not found in the store or store was disabled and fetching from remote yielded nothing or it was disabled", u.String())
}

func (f *Fetcher) fetchSingleImageByDockerURL(u *url.URL) (string, error) {
	rem, err := f.getRemoteForURL(u)
	if err != nil {
		return "", err
	}
	// TODO(krnowak): use strict checking when we implement
	// setting CacheMaxAge in store.Remote for docker images
	if h := f.maybeCheckRemoteFromStore(rem, remoteCheckLax); h != "" {
		return h, nil
	}
	if h, err := f.maybeFetchDockerURLFromRemote(u); h != "" || err != nil {
		return h, err
	}
	return "", fmt.Errorf("unable to fetch docker image from URL %q: either image was not found in the store or store was disabled and fetching from remote yielded nothing or it was disabled", u.String())
}

func (f *Fetcher) getRemoteForURL(u *url.URL) (*store.Remote, error) {
	if f.NoCache {
		return nil, nil
	}
	urlStr := u.String()
	if rem, ok, err := f.S.GetRemote(urlStr); err != nil {
		return nil, errwrap.Wrap(fmt.Errorf("failed to fetch URL %q", urlStr), err)
	} else if ok {
		return rem, nil
	}
	return nil, nil
}

func (f *Fetcher) maybeCheckRemoteFromStore(rem *store.Remote, check remoteCheck) string {
	if f.NoStore || rem == nil {
		return ""
	}
	useBlobKey := false
	switch check {
	case remoteCheckLax:
		useBlobKey = true
	case remoteCheckStrict:
		useBlobKey = useCached(rem.DownloadTime, rem.CacheMaxAge)
	}
	if useBlobKey {
		log.Printf("using image from local store for url %s", rem.ACIURL)
		return rem.BlobKey
	}
	return ""
}

func (f *Fetcher) maybeFetchHTTPURLFromRemote(rem *store.Remote, u *url.URL, a *asc) (string, error) {
	if !f.StoreOnly {
		log.Printf("remote fetching from URL %q", u.String())
		hf := &httpFetcher{
			InsecureFlags: f.InsecureFlags,
			S:             f.S,
			Ks:            f.Ks,
			Rem:           rem,
			Debug:         f.Debug,
			Headers:       f.Headers,
		}
		return hf.GetHash(u, a)
	}
	return "", nil
}

func (f *Fetcher) maybeFetchDockerURLFromRemote(u *url.URL) (string, error) {
	if !f.StoreOnly {
		log.Printf("remote fetching from URL %q", u.String())
		df := &dockerFetcher{
			InsecureFlags: f.InsecureFlags,
			DockerAuth:    f.DockerAuth,
			S:             f.S,
			Debug:         f.Debug,
		}
		return df.GetHash(u)
	}
	return "", nil
}

func (f *Fetcher) fetchSingleImageByPath(path string, a *asc) (string, error) {
	log.Printf("using image from file %s", path)
	ff := &fileFetcher{
		InsecureFlags: f.InsecureFlags,
		S:             f.S,
		Ks:            f.Ks,
		Debug:         f.Debug,
	}
	return ff.GetHash(path, a)
}

type appBundle struct {
	App *discovery.App
	Str string
}

func newAppBundle(name string) (*appBundle, error) {
	app, err := discovery.NewAppFromString(name)
	if err != nil {
		return nil, errwrap.Wrap(fmt.Errorf("invalid image name %q", name), err)
	}
	if _, ok := app.Labels["arch"]; !ok {
		app.Labels["arch"] = runtime.GOARCH
	}
	if _, ok := app.Labels["os"]; !ok {
		app.Labels["os"] = runtime.GOOS
	}
	if err := types.IsValidOSArch(app.Labels, stage0.ValidOSArch); err != nil {
		return nil, errwrap.Wrap(fmt.Errorf("invalid image name %q", name), err)
	}
	bundle := &appBundle{
		App: app,
		Str: name,
	}
	return bundle, nil
}

func (f *Fetcher) fetchSingleImageByName(name string, a *asc) (string, error) {
	app, err := newAppBundle(name)
	if err != nil {
		return "", err
	}
	if h, err := f.maybeCheckStoreForApp(app); h != "" || err != nil {
		return h, err
	}
	if h, err := f.maybeFetchImageFromRemote(app, a); h != "" || err != nil {
		return h, err
	}
	return "", fmt.Errorf("unable to fetch image from image name %q: either image was not found in the store or store was disabled and fetching from remote yielded nothing or it was disabled", name)
}

func (f *Fetcher) maybeCheckStoreForApp(app *appBundle) (string, error) {
	if !f.NoStore {
		key, err := f.getStoreKeyFromApp(app)
		if err == nil {
			log.Printf("using image from local store for image name %s", app.Str)
			return key, nil
		}
		switch err.(type) {
		case store.ACINotFoundError:
			// ignore the "not found" error
		default:
			return "", err
		}
	}
	return "", nil
}

func (f *Fetcher) getStoreKeyFromApp(app *appBundle) (string, error) {
	labels, err := types.LabelsFromMap(app.App.Labels)
	if err != nil {
		return "", errwrap.Wrap(fmt.Errorf("invalid labels in the name %q", app.Str), err)
	}
	key, err := f.S.GetACI(app.App.Name, labels)
	if err != nil {
		switch err.(type) {
		case store.ACINotFoundError:
			return "", err
		default:
			return "", errwrap.Wrap(fmt.Errorf("cannot find image %q", app.Str), err)
		}
	}
	return key, nil
}

func (f *Fetcher) maybeFetchImageFromRemote(app *appBundle, a *asc) (string, error) {
	if !f.StoreOnly {
		nf := &nameFetcher{
			InsecureFlags:      f.InsecureFlags,
			S:                  f.S,
			Ks:                 f.Ks,
			Debug:              f.Debug,
			Headers:            f.Headers,
			TrustKeysFromHTTPS: f.TrustKeysFromHTTPS,
		}
		return nf.GetHash(app.App, a)
	}
	return "", nil
}
<commit_msg>rkt/image: render images only if we're root<commit_after>// Copyright 2015 The rkt 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 image

import (
	"container/list"
	"errors"
	"fmt"
	"net/url"
	"os"
	"runtime"

	"github.com/coreos/rkt/common"
	"github.com/coreos/rkt/common/apps"
	"github.com/coreos/rkt/stage0"
	"github.com/coreos/rkt/store"
	"github.com/hashicorp/errwrap"

	"github.com/appc/spec/discovery"
	"github.com/appc/spec/schema/types"
)

// Fetcher will try to fetch images into the store.
type Fetcher action

// FetchImage will take an image as either a path, a URL or a name
// string and import it into the store if found. If ascPath is not "",
// it must exist as a local file and will be used as the signature
// file for verification, unless verification is disabled. If
// f.WithDeps is true also image dependencies are fetched.
func (f *Fetcher) FetchImage(img string, ascPath string, imgType apps.AppImageType) (string, error) {
	ensureLogger(f.Debug)
	a := f.getAsc(ascPath)
	hash, err := f.fetchSingleImage(img, a, imgType)
	if err != nil {
		return "", err
	}
	if f.WithDeps {
		err = f.fetchImageDeps(hash)
		if err != nil {
			return "", err
		}
	}
	// we need to be able to do a chroot and access to the tree store
	// directories, check if we're root
	if common.SupportsOverlay() && os.Geteuid() == 0 {
		if _, _, err := f.S.RenderTreeStore(hash, false); err != nil {
			return "", errwrap.Wrap(errors.New("error rendering tree store"), err)
		}
	}
	return hash, nil
}

func (f *Fetcher) getAsc(ascPath string) *asc {
	if ascPath != "" {
		return &asc{
			Location: ascPath,
			Fetcher:  &localAscFetcher{},
		}
	}
	return &asc{}
}

// fetchImageDeps will recursively fetch all the image dependencies
func (f *Fetcher) fetchImageDeps(hash string) error {
	imgsl := list.New()
	seen := map[string]struct{}{}
	f.addImageDeps(hash, imgsl, seen)
	for el := imgsl.Front(); el != nil; el = el.Next() {
		a := &asc{}
		img := el.Value.(string)
		hash, err := f.fetchSingleImage(img, a, apps.AppImageName)
		if err != nil {
			return err
		}
		f.addImageDeps(hash, imgsl, seen)
	}
	return nil
}

func (f *Fetcher) addImageDeps(hash string, imgsl *list.List, seen map[string]struct{}) error {
	dependencies, err := f.getImageDeps(hash)
	if err != nil {
		return errwrap.Wrap(fmt.Errorf("failed to get dependencies for image ID %q", hash), err)
	}
	for _, d := range dependencies {
		imgName := d.ImageName.String()
		app, err := discovery.NewApp(imgName, d.Labels.ToMap())
		if err != nil {
			return errwrap.Wrap(fmt.Errorf("one of image ID's %q dependencies (image %q) is invalid", hash, imgName), err)
		}
		appStr := app.String()
		if _, ok := seen[appStr]; ok {
			continue
		}
		imgsl.PushBack(app.String())
		seen[appStr] = struct{}{}
	}
	return nil
}

func (f *Fetcher) getImageDeps(hash string) (types.Dependencies, error) {
	key, err := f.S.ResolveKey(hash)
	if err != nil {
		return nil, err
	}
	im, err := f.S.GetImageManifest(key)
	if err != nil {
		return nil, err
	}
	return im.Dependencies, nil
}

func (f *Fetcher) fetchSingleImage(img string, a *asc, imgType apps.AppImageType) (string, error) {
	if imgType == apps.AppImageGuess {
		imgType = guessImageType(img)
	}
	if imgType == apps.AppImageHash {
		return "", fmt.Errorf("cannot fetch a hash %q, expected either a URL, a path or an image name", img)
	}

	switch imgType {
	case apps.AppImageURL:
		return f.fetchSingleImageByURL(img, a)
	case apps.AppImagePath:
		return f.fetchSingleImageByPath(img, a)
	case apps.AppImageName:
		return f.fetchSingleImageByName(img, a)
	default:
		return "", fmt.Errorf("unknown image type %d", imgType)
	}
}

type remoteCheck int

const (
	remoteCheckLax remoteCheck = iota
	remoteCheckStrict
)

func (f *Fetcher) fetchSingleImageByURL(urlStr string, a *asc) (string, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return "", errwrap.Wrap(fmt.Errorf("invalid image URL %q", urlStr), err)
	}

	switch u.Scheme {
	case "http", "https":
		return f.fetchSingleImageByHTTPURL(u, a)
	case "docker":
		return f.fetchSingleImageByDockerURL(u)
	case "file":
		return f.fetchSingleImageByPath(u.Path, a)
	case "":
		return "", fmt.Errorf("expected image URL %q to contain a scheme", urlStr)
	default:
		return "", fmt.Errorf("an unsupported URL scheme %q - the only URL schemes supported by rkt are docker, http, https and file", u.Scheme)
	}
}

func (f *Fetcher) fetchSingleImageByHTTPURL(u *url.URL, a *asc) (string, error) {
	rem, err := f.getRemoteForURL(u)
	if err != nil {
		return "", err
	}
	if h := f.maybeCheckRemoteFromStore(rem, remoteCheckStrict); h != "" {
		return h, nil
	}
	if h, err := f.maybeFetchHTTPURLFromRemote(rem, u, a); h != "" || err != nil {
		return h, err
	}
	return "", fmt.Errorf("unable to fetch image from URL %q: either image was not found in the store or store was disabled and fetching from remote yielded nothing or it was disabled", u.String())
}

func (f *Fetcher) fetchSingleImageByDockerURL(u *url.URL) (string, error) {
	rem, err := f.getRemoteForURL(u)
	if err != nil {
		return "", err
	}
	// TODO(krnowak): use strict checking when we implement
	// setting CacheMaxAge in store.Remote for docker images
	if h := f.maybeCheckRemoteFromStore(rem, remoteCheckLax); h != "" {
		return h, nil
	}
	if h, err := f.maybeFetchDockerURLFromRemote(u); h != "" || err != nil {
		return h, err
	}
	return "", fmt.Errorf("unable to fetch docker image from URL %q: either image was not found in the store or store was disabled and fetching from remote yielded nothing or it was disabled", u.String())
}

func (f *Fetcher) getRemoteForURL(u *url.URL) (*store.Remote, error) {
	if f.NoCache {
		return nil, nil
	}
	urlStr := u.String()
	if rem, ok, err := f.S.GetRemote(urlStr); err != nil {
		return nil, errwrap.Wrap(fmt.Errorf("failed to fetch URL %q", urlStr), err)
	} else if ok {
		return rem, nil
	}
	return nil, nil
}

func (f *Fetcher) maybeCheckRemoteFromStore(rem *store.Remote, check remoteCheck) string {
	if f.NoStore || rem == nil {
		return ""
	}
	useBlobKey := false
	switch check {
	case remoteCheckLax:
		useBlobKey = true
	case remoteCheckStrict:
		useBlobKey = useCached(rem.DownloadTime, rem.CacheMaxAge)
	}
	if useBlobKey {
		log.Printf("using image from local store for url %s", rem.ACIURL)
		return rem.BlobKey
	}
	return ""
}

func (f *Fetcher) maybeFetchHTTPURLFromRemote(rem *store.Remote, u *url.URL, a *asc) (string, error) {
	if !f.StoreOnly {
		log.Printf("remote fetching from URL %q", u.String())
		hf := &httpFetcher{
			InsecureFlags: f.InsecureFlags,
			S:             f.S,
			Ks:            f.Ks,
			Rem:           rem,
			Debug:         f.Debug,
			Headers:       f.Headers,
		}
		return hf.GetHash(u, a)
	}
	return "", nil
}

func (f *Fetcher) maybeFetchDockerURLFromRemote(u *url.URL) (string, error) {
	if !f.StoreOnly {
		log.Printf("remote fetching from URL %q", u.String())
		df := &dockerFetcher{
			InsecureFlags: f.InsecureFlags,
			DockerAuth:    f.DockerAuth,
			S:             f.S,
			Debug:         f.Debug,
		}
		return df.GetHash(u)
	}
	return "", nil
}

func (f *Fetcher) fetchSingleImageByPath(path string, a *asc) (string, error) {
	log.Printf("using image from file %s", path)
	ff := &fileFetcher{
		InsecureFlags: f.InsecureFlags,
		S:             f.S,
		Ks:            f.Ks,
		Debug:         f.Debug,
	}
	return ff.GetHash(path, a)
}

type appBundle struct {
	App *discovery.App
	Str string
}

func newAppBundle(name string) (*appBundle, error) {
	app, err := discovery.NewAppFromString(name)
	if err != nil {
		return nil, errwrap.Wrap(fmt.Errorf("invalid image name %q", name), err)
	}
	if _, ok := app.Labels["arch"]; !ok {
		app.Labels["arch"] = runtime.GOARCH
	}
	if _, ok := app.Labels["os"]; !ok {
		app.Labels["os"] = runtime.GOOS
	}
	if err := types.IsValidOSArch(app.Labels, stage0.ValidOSArch); err != nil {
		return nil, errwrap.Wrap(fmt.Errorf("invalid image name %q", name), err)
	}
	bundle := &appBundle{
		App: app,
		Str: name,
	}
	return bundle, nil
}

func (f *Fetcher) fetchSingleImageByName(name string, a *asc) (string, error) {
	app, err := newAppBundle(name)
	if err != nil {
		return "", err
	}
	if h, err := f.maybeCheckStoreForApp(app); h != "" || err != nil {
		return h, err
	}
	if h, err := f.maybeFetchImageFromRemote(app, a); h != "" || err != nil {
		return h, err
	}
	return "", fmt.Errorf("unable to fetch image from image name %q: either image was not found in the store or store was disabled and fetching from remote yielded nothing or it was disabled", name)
}

func (f *Fetcher) maybeCheckStoreForApp(app *appBundle) (string, error) {
	if !f.NoStore {
		key, err := f.getStoreKeyFromApp(app)
		if err == nil {
			log.Printf("using image from local store for image name %s", app.Str)
			return key, nil
		}
		switch err.(type) {
		case store.ACINotFoundError:
			// ignore the "not found" error
		default:
			return "", err
		}
	}
	return "", nil
}

func (f *Fetcher) getStoreKeyFromApp(app *appBundle) (string, error) {
	labels, err := types.LabelsFromMap(app.App.Labels)
	if err != nil {
		return "", errwrap.Wrap(fmt.Errorf("invalid labels in the name %q", app.Str), err)
	}
	key, err := f.S.GetACI(app.App.Name, labels)
	if err != nil {
		switch err.(type) {
		case store.ACINotFoundError:
			return "", err
		default:
			return "", errwrap.Wrap(fmt.Errorf("cannot find image %q", app.Str), err)
		}
	}
	return key, nil
}

func (f *Fetcher) maybeFetchImageFromRemote(app *appBundle, a *asc) (string, error) {
	if !f.StoreOnly {
		nf := &nameFetcher{
			InsecureFlags:      f.InsecureFlags,
			S:                  f.S,
			Ks:                 f.Ks,
			Debug:              f.Debug,
			Headers:            f.Headers,
			TrustKeysFromHTTPS: f.TrustKeysFromHTTPS,
		}
		return nf.GetHash(app.App, a)
	}
	return "", nil
}
<|endoftext|>


<commit_before>// Package game keeps track of basic poker game information. Any Player may use
// the game.Play function in order to play a game.
package game

import (
	"fmt"
	"log"
	"net"

	"poker/game/diff"
)

const (
	PreFlop = iota
	Flop
	Turn
	River
	Showdown
)

func splitCards(s string) []string {
	a := make([]string, len(s)/2)
	na := 0
	for i, c := range s {
		if c == 'c' || c == 'd' || c == 'h' || c == 's' {
			a[na] = s[i-1 : i+1]
			na++
		}
	}
	return a[:na]
}

type Player interface {
	Play(g *Game) (action string)
	Observe(g *Game)
}

type Game struct {
	Round     int       // 0-4: pre-flop, flop, turn, river, showdown.
	Bets      []float64 // The chips put in for each player this round.
	Holes     []string  // All of the viewable hole cards.
	Board     []string  // All of the board cards.
	Raises    int       // The number of raises this round.
	Actions   []byte    // The last action taken by each player.
	Actor     int       // The player whose turn it is to act.
	*Rules              // The set of rules to use to play the game.
	pot       float64   // Chips in the pot from previous rounds.
	Event     interface{} // The most recent event
	*diff.Players
}

func NewGame(rules string) (*Game, error) {
	r, err := ChooseRules(rules)
	if err != nil {
		return nil, err
	}
	return &Game{
		Rules:  r,
		Actions: make([]byte, r.numPlayers),
		Bets:   make([]float64, r.numPlayers)}, nil
}

func (this *Game) String() string {
	s := fmt.Sprintln(this.Holes, this.Board)
	if this.Actor != -1 {
		s += fmt.Sprintln(this.Pot(), this.Bets, this.CallAmt(), this.RaiseAmt())
	}
	return s
}

// NumActive returns how many players are still in the hand.
func (this *Game) NumActive() int {
	var count int
	for _, action := range this.Actions {
		if action != 'f' {
			count++
		}
	}
	return count
}

// LegalActions returns a string containing the currently legal actions.
func (this *Game) LegalActions() string {
	actions := "c"
	if this.CallAmt() > 0 {
		actions += "f"
	}
	if this.Raises < this.maxRaises[this.Round] {
		actions += "r"
	}
	return actions
}

func (this *Game) CallAmt() float64 {
	var max float64
	for _, chips := range this.Bets {
		if chips > max {
			max = chips
		}
	}
	return max - this.Bets[this.Actor]
}

func (this *Game) RaiseAmt() float64 {
	return this.CallAmt() + this.raiseSize[this.Round]
}

func (this *Game) Pot() float64 {
	var sum float64
	for _, chips := range this.Bets {
		sum += chips
	}
	return this.pot + sum
}

// Is there anyone in the hand who has not acted yet? If not, then does everyone
// who has not folded have the same amount in the pot?
// Pre-condition: two players have not folded.
func (this *Game) evenBets() bool {
	var j int
	var b0 float64
	for i, a := range this.Actions {
		if a == 0 {
			return false
		}
		if a != 'f' {
			b0 = this.Bets[i]
			j = i
			break
		}
	}
	for i, a := range this.Actions[j+1:] {
		if a == 0 || (a != 'f' && this.Bets[i+j+1] != b0) {
			return false
		}
	}
	return true
}

func (this *Game) Update(event interface{}) {
	this.Event = event
	switch e := event.(type) {
	case *diff.Players:
		this.Players = e
		this.Round = -1
		this.pot = 0
		copy(this.Bets, this.blind)
		this.Board = nil
		for i := range this.Actions {
			this.Actions[i] = 0
		}
	case diff.Cards:
		for i, a := range this.Actions {
			if a != 'f' {
				this.Actions[i] = 0
			}
		}
		cards := splitCards(string(e))
		this.Round++
		this.Raises = 0
		switch this.Round {
		case PreFlop:
			this.Actor = this.firstPlayer[this.Round] - 1
			this.Holes = cards
		case Flop, Turn, River:
			this.Actor = this.firstPlayer[this.Round] - 1
			this.pot = this.Pot()
			for i := range this.Bets {
				this.Bets[i] = 0
			}
			this.Board = append(this.Board, cards...)
		case Showdown:
			this.Actor = -1
			this.Holes = cards
		}
	case diff.Action:
		action := string(e)
		this.Actions[this.Actor] = action[0]
		switch action {
		case "c":
			this.Bets[this.Actor] += this.CallAmt()
		case "r":
			this.Raises++
			this.Bets[this.Actor] += this.RaiseAmt()
		}
		if this.NumActive() < 2  || this.evenBets() {
			this.Actor = -1
			return
		}
		i := this.Actor
		for {
			i = (i + 1) % len(this.Actions)
			if this.Actions[i] != 'f' {
				this.Actor = i
				return
			}
		}
	default:
		panic("game: Invalid event passed to Update")
	}
}

// Start playing a game.
//	rules -- a String naming the game to play
//	p     -- an object that implements the Player interface.
//	host  -- the InetAddress of the dealer passed as a String
//	port  -- the port the dealer is listening on for the client passed as a String
func Play(rules string, p Player, host, port string) {
	addr := net.JoinHostPort(host, port)
	fmt.Printf("Connecting to dealer at %s to play %s...\n", addr, rules)
	in, out, err := diff.NewACPC(addr)
	if err != nil {
		log.Fatalln(err)
	}
	game, err := NewGame(rules)
	if err != nil {
		log.Fatalln(err)
	}
	for event := range in {
		game.Update(event)
		if game.Actor == game.Viewer {
			out <- p.Play(game)
		} else {
			p.Observe(game)
		}
	}
	fmt.Println("GAME OVER")
}
<commit_msg>Add code to game for generating trees<commit_after>// Package game keeps track of basic poker game information. Any Player may use
// the game.Play function in order to play a game.
package game

import (
	"fmt"
	"log"
	"net"

	"poker/game/diff"
)

const (
	PreFlop = iota
	Flop
	Turn
	River
	Showdown
)

func splitCards(s string) []string {
	a := make([]string, len(s)/2)
	na := 0
	for i, c := range s {
		if c == 'c' || c == 'd' || c == 'h' || c == 's' {
			a[na] = s[i-1 : i+1]
			na++
		}
	}
	return a[:na]
}

type Player interface {
	Play(g *Game) (action string)
	Observe(g *Game)
}

type Game struct {
	Round     int          // 0-4: pre-flop, flop, turn, river, showdown.
	Bets      [4][]float64 // The chips put in for each player for each round.
	Holes     []string     // All of the viewable hole cards.
	Board     []string     // All of the board cards.
	Raises    int          // The number of raises this round.
	Actions   []byte       // The last action taken by each player.
	Actor     int          // The player whose turn it is to act.
	*Rules                 // The set of rules to use to play the game.
	Event     interface{}  // The most recent event
	*diff.Players
}

// Makes a partial copy of the Game that is specifically suited for creating
// game trees.
func (this *Game) Copy() *Game {
	g := new(Game)
	g.Round = this.Round
	for i := range g.Bets {
		g.Bets[i] = make([]float64, len(this.Bets[0]))
		copy(g.Bets[i], this.Bets[i])
	}
	g.Holes = make([]string, len(this.Holes))
	copy(g.Holes, this.Holes)
	g.Board = make([]string, len(this.Board))
	copy(g.Board, this.Board)
	g.Raises = this.Raises
	g.Actions = make([]byte, len(this.Actions))
	copy(g.Actions, this.Actions)
	g.Actor = this.Actor
	g.Rules = this.Rules
	//Event     interface{}
	g.Players = this.Players
	return g
}

func NewGame(rules string) (*Game, error) {
	r, err := ChooseRules(rules)
	if err != nil {
		return nil, err
	}
	g := new(Game)
	g.Rules = r
	g.Actions = make([]byte, r.numPlayers)
	for i := range g.Bets {
		g.Bets[i] = make([]float64, r.numPlayers)
	}
	return g, nil
}

func (this *Game) String() string {
	s := fmt.Sprintln(this.Holes, this.Board)
	if this.Actor != -1 {
		s += fmt.Sprintln(this.Pot(), this.Bets, this.CallAmt(), this.RaiseAmt())
	}
	return s
}

// NumActive returns how many players are still in the hand.
func (this *Game) NumActive() int {
	var count int
	for _, action := range this.Actions {
		if action != 'f' {
			count++
		}
	}
	return count
}

// LegalActions returns a string containing the currently legal actions.
func (this *Game) LegalActions() string {
	actions := "c"
	if this.CallAmt() > 0 {
		actions += "f"
	}
	if this.Raises < this.maxRaises[this.Round] {
		actions += "r"
	}
	return actions
}

func (this *Game) CallAmt() float64 {
	var max float64
	for _, chips := range this.Bets[this.Round] {
		if chips > max {
			max = chips
		}
	}
	return max - this.Bets[this.Round][this.Actor]
}

func (this *Game) RaiseAmt() float64 {
	return this.CallAmt() + this.raiseSize[this.Round]
}

func (this *Game) Pot() float64 {
	var sum float64
	for _, bets := range this.Bets {
		for _, chips := range bets {
			sum += chips
		}
	}
	return sum
}

// Is there anyone in the hand who has not acted yet? If not, then does everyone
// who has not folded have the same amount in the pot?
// Pre-condition: two players have not folded.
func (this *Game) evenBets() bool {
	var j int
	var b0 float64
	for i, a := range this.Actions {
		if a == 0 {
			return false
		}
		if a != 'f' {
			b0 = this.Bets[this.Round][i]
			j = i
			break
		}
	}
	for i, a := range this.Actions[j+1:] {
		if a == 0 || (a != 'f' && this.Bets[this.Round][i+j+1] != b0) {
			return false
		}
	}
	return true
}

func (this *Game) Update(event interface{}) {
	this.Event = event
	switch e := event.(type) {
	case *diff.Players:
		this.Actor = -1
		this.Players = e
		this.Round = -1
		for _, bets := range this.Bets {
			for i := range bets {
				bets[i] = 0
			}
		}
		copy(this.Bets[0], this.blind)
		this.Board = nil
		for i := range this.Actions {
			this.Actions[i] = 0
		}
	case diff.Cards:
		for i, a := range this.Actions {
			if a != 'f' {
				this.Actions[i] = 0
			}
		}
		cards := splitCards(string(e))
		this.Round++
		this.Raises = 0
		switch this.Round {
		case PreFlop:
			this.Actor = this.firstPlayer[this.Round] - 1
			this.Holes = cards
		case Flop, Turn, River:
			this.Actor = this.firstPlayer[this.Round] - 1
			this.Board = append(this.Board, cards...)
		case Showdown:
			this.Actor = -1
			this.Holes = cards
		}
	case diff.Action:
		action := string(e)
		this.Actions[this.Actor] = action[0]
		switch action {
		case "c":
			this.Bets[this.Round][this.Actor] += this.CallAmt()
		case "r":
			this.Raises++
			this.Bets[this.Round][this.Actor] += this.RaiseAmt()
		}
		if this.NumActive() < 2  || this.evenBets() {
			this.Actor = -1
			return
		}
		i := this.Actor
		for {
			i = (i + 1) % len(this.Actions)
			if this.Actions[i] != 'f' {
				this.Actor = i
				return
			}
		}
	default:
		panic("game: Invalid event passed to Update")
	}
}

// Start playing a game.
//	rules -- a String naming the game to play
//	p     -- an object that implements the Player interface.
//	host  -- the InetAddress of the dealer passed as a String
//	port  -- the port the dealer is listening on for the client passed as a String
func Play(rules string, p Player, host, port string) {
	addr := net.JoinHostPort(host, port)
	fmt.Printf("Connecting to dealer at %s to play %s...\n", addr, rules)
	in, out, err := diff.NewACPC(addr)
	if err != nil {
		log.Fatalln(err)
	}
	game, err := NewGame(rules)
	if err != nil {
		log.Fatalln(err)
	}
	for event := range in {
		game.Update(event)
		if game.Actor == game.Viewer {
			out <- p.Play(game)
		} else {
			p.Observe(game)
		}
	}
	fmt.Println("GAME OVER")
}
<|endoftext|>


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

// Package qemu provides a Go API for starting QEMU VMs.
//
// qemu is mainly suitable for running QEMU-based integration tests.
//
// The environment variable `UROOT_QEMU` overrides the path to QEMU and the
// first few arguments (defaults to "qemu"). For example, I use:
//
//     UROOT_QEMU='qemu-system-x86_64 -L . -m 4096 -enable-kvm'
//
// For CI, this environment variable is set in `.circleci/images/integration/Dockerfile`.
package qemu

import (
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"
	"time"

	expect "github.com/google/goexpect"
)

// DefaultTimeout for `Expect` and `ExpectRE` functions.
var DefaultTimeout = 7 * time.Second

// TimeoutMultiplier increases all timeouts proportionally. Useful when running
// QEMU on a slow machine.
var TimeoutMultiplier = 2.0

// Options are VM start-up parameters.
type Options struct {
	// QEMUPath is the path to the QEMU binary to invoke.
	//
	// If left unspecified, the UROOT_QEMU env var will be used.
	// If the env var is unspecified, "qemu" is the default.
	QEMUPath string

	// Path to the kernel to boot.
	Kernel string

	// Path to the initramfs.
	Initramfs string

	// Extra kernel command-line arguments.
	KernelArgs string

	// Where to send serial output.
	SerialOutput io.WriteCloser

	// Timeout is the expect timeout.
	Timeout time.Duration

	// Devices are devices to expose to the QEMU VM.
	Devices []Device
}

// Start starts a QEMU VM.
func (o *Options) Start() (*VM, error) {
	cmdline := cmdline(o)

	gExpect, _, err := expect.SpawnWithArgs(cmdline, -1,
		expect.Tee(o.SerialOutput),
		expect.CheckDuration(2*time.Millisecond))
	if err != nil {
		return nil, err
	}
	return &VM{
		Options: o,
		cmdline: cmdline,
		gExpect: gExpect,
	}, nil
}

// cmdline returns the command line arguments used to start QEMU. These
// arguments are derived from the given QEMU struct.
func cmdline(o *Options) []string {
	var args []string
	if len(o.QEMUPath) > 0 {
		args = append(args, o.QEMUPath)
	} else {
		// Read first few arguments for env.
		env := os.Getenv("UROOT_QEMU")
		if env == "" {
			env = "qemu" // default
		}
		args = append(args, strings.Fields(env)...)
	}

	// Disable graphics because we are using serial.
	args = append(args, "-nographic")

	// Arguments passed to the kernel:
	//
	// - earlyprintk=ttyS0: print very early debug messages to the serial
	// - console=ttyS0: /dev/console points to /dev/ttyS0 (the serial port)
	// - o.KernelArgs: extra, optional kernel arguments
	if len(o.Kernel) != 0 {
		args = append(args, "-kernel", o.Kernel)
		cmdline := "console=ttyS0 earlyprintk=ttyS0"
		if len(o.KernelArgs) != 0 {
			cmdline += " " + o.KernelArgs
		}
		args = append(args, "-append", cmdline)
	}
	if len(o.Initramfs) != 0 {
		args = append(args, "-initrd", o.Initramfs)
	}

	for _, dev := range o.Devices {
		if dev != nil {
			if c := dev.Cmdline(); c != nil {
				args = append(args, c...)
			}
		}
	}
	return args
}

// VM is a running QEMU virtual machine.
type VM struct {
	Options *Options
	cmdline []string
	gExpect *expect.GExpect
}

func (v *VM) Cmdline() []string {
	// Maybe return a copy?
	return v.cmdline
}

// CmdlineQuoted quotes any of QEMU's command line arguments containing a space
// so it is easy to copy-n-paste into a shell for debugging.
func (v *VM) CmdlineQuoted() string {
	args := make([]string, len(v.cmdline))
	for i, arg := range v.cmdline {
		if strings.ContainsAny(arg, " \t\n") {
			args[i] = fmt.Sprintf("'%s'", arg)
		} else {
			args[i] = arg
		}
	}
	return strings.Join(args, " ")
}

// Close stops QEMU.
func (v *VM) Close() {
	v.gExpect.Close()
	v.gExpect = nil
}

// Send sends a string to QEMU's serial.
func (v *VM) Send(in string) {
	v.gExpect.Send(in)
}

func (v *VM) TimeoutOr() time.Duration {
	if v.Options.Timeout == 0 {
		return DefaultTimeout
	}
	return v.Options.Timeout
}

// Expect returns an error if the given string is not found in vEMU's serial
// output within `DefaultTimeout`.
func (v *VM) Expect(search string) error {
	return v.ExpectTimeout(search, v.TimeoutOr())
}

// ExpectTimeout returns an error if the given string is not found in vEMU's serial
// output within the given timeout.
func (v *VM) ExpectTimeout(search string, timeout time.Duration) error {
	_, err := v.ExpectRETimeout(regexp.MustCompile(regexp.QuoteMeta(search)), timeout)
	return err
}

// ExpectRE returns an error if the given regular expression is not found in
// vEMU's serial output within `DefaultTimeout`. The matched string is
// returned.
func (v *VM) ExpectRE(pattern *regexp.Regexp) (string, error) {
	return v.ExpectRETimeout(pattern, v.TimeoutOr())
}

// ExpectRETimeout returns an error if the given regular expression is not
// found in vEMU's serial output within the given timeout. The matched string
// is returned.
func (v *VM) ExpectRETimeout(pattern *regexp.Regexp, timeout time.Duration) (string, error) {
	scaled := time.Duration(float64(timeout) * TimeoutMultiplier)
	str, _, err := v.gExpect.Expect(pattern, scaled)
	return str, err
}
<commit_msg>qemu: add a VM.Wait() method.<commit_after>// Copyright 2018 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package qemu provides a Go API for starting QEMU VMs.
//
// qemu is mainly suitable for running QEMU-based integration tests.
//
// The environment variable `UROOT_QEMU` overrides the path to QEMU and the
// first few arguments (defaults to "qemu"). For example, I use:
//
//     UROOT_QEMU='qemu-system-x86_64 -L . -m 4096 -enable-kvm'
//
// For CI, this environment variable is set in `.circleci/images/integration/Dockerfile`.
package qemu

import (
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"
	"time"

	expect "github.com/google/goexpect"
)

// DefaultTimeout for `Expect` and `ExpectRE` functions.
var DefaultTimeout = 7 * time.Second

// TimeoutMultiplier increases all timeouts proportionally. Useful when running
// QEMU on a slow machine.
var TimeoutMultiplier = 2.0

// Options are VM start-up parameters.
type Options struct {
	// QEMUPath is the path to the QEMU binary to invoke.
	//
	// If left unspecified, the UROOT_QEMU env var will be used.
	// If the env var is unspecified, "qemu" is the default.
	QEMUPath string

	// Path to the kernel to boot.
	Kernel string

	// Path to the initramfs.
	Initramfs string

	// Extra kernel command-line arguments.
	KernelArgs string

	// Where to send serial output.
	SerialOutput io.WriteCloser

	// Timeout is the expect timeout.
	Timeout time.Duration

	// Devices are devices to expose to the QEMU VM.
	Devices []Device
}

// Start starts a QEMU VM.
func (o *Options) Start() (*VM, error) {
	cmdline := o.Cmdline()

	gExpect, ch, err := expect.SpawnWithArgs(cmdline, -1,
		expect.Tee(o.SerialOutput),
		expect.CheckDuration(2*time.Millisecond))
	if err != nil {
		return nil, err
	}
	return &VM{
		Options: o,
		errCh:   ch,
		cmdline: cmdline,
		gExpect: gExpect,
	}, nil
}

// cmdline returns the command line arguments used to start QEMU. These
// arguments are derived from the given QEMU struct.
func (o *Options) Cmdline() []string {
	var args []string
	if len(o.QEMUPath) > 0 {
		args = append(args, o.QEMUPath)
	} else {
		// Read first few arguments for env.
		env := os.Getenv("UROOT_QEMU")
		if env == "" {
			env = "qemu" // default
		}
		args = append(args, strings.Fields(env)...)
	}

	// Disable graphics because we are using serial.
	args = append(args, "-nographic")

	// Arguments passed to the kernel:
	//
	// - earlyprintk=ttyS0: print very early debug messages to the serial
	// - console=ttyS0: /dev/console points to /dev/ttyS0 (the serial port)
	// - o.KernelArgs: extra, optional kernel arguments
	if len(o.Kernel) != 0 {
		args = append(args, "-kernel", o.Kernel)
		cmdline := "console=ttyS0 earlyprintk=ttyS0"
		if len(o.KernelArgs) != 0 {
			cmdline += " " + o.KernelArgs
		}
		args = append(args, "-append", cmdline)
	}
	if len(o.Initramfs) != 0 {
		args = append(args, "-initrd", o.Initramfs)
	}

	for _, dev := range o.Devices {
		if dev != nil {
			if c := dev.Cmdline(); c != nil {
				args = append(args, c...)
			}
		}
	}
	return args
}

// VM is a running QEMU virtual machine.
type VM struct {
	Options *Options
	cmdline []string
	errCh   <-chan error
	gExpect *expect.GExpect
}

// Wait waits for the VM to exit.
func (v *VM) Wait() error {
	return <-v.errCh
}

// Cmdline is the command-line the VM was started with.
func (v *VM) Cmdline() []string {
	// Maybe return a copy?
	return v.cmdline
}

// CmdlineQuoted quotes any of QEMU's command line arguments containing a space
// so it is easy to copy-n-paste into a shell for debugging.
func (v *VM) CmdlineQuoted() string {
	args := make([]string, len(v.cmdline))
	for i, arg := range v.cmdline {
		if strings.ContainsAny(arg, " \t\n") {
			args[i] = fmt.Sprintf("'%s'", arg)
		} else {
			args[i] = arg
		}
	}
	return strings.Join(args, " ")
}

// Close stops QEMU.
func (v *VM) Close() {
	v.gExpect.Close()
	v.gExpect = nil
}

// Send sends a string to QEMU's serial.
func (v *VM) Send(in string) {
	v.gExpect.Send(in)
}

func (v *VM) TimeoutOr() time.Duration {
	if v.Options.Timeout == 0 {
		return DefaultTimeout
	}
	return v.Options.Timeout
}

// Expect returns an error if the given string is not found in vEMU's serial
// output within `DefaultTimeout`.
func (v *VM) Expect(search string) error {
	return v.ExpectTimeout(search, v.TimeoutOr())
}

// ExpectTimeout returns an error if the given string is not found in vEMU's serial
// output within the given timeout.
func (v *VM) ExpectTimeout(search string, timeout time.Duration) error {
	_, err := v.ExpectRETimeout(regexp.MustCompile(regexp.QuoteMeta(search)), timeout)
	return err
}

// ExpectRE returns an error if the given regular expression is not found in
// vEMU's serial output within `DefaultTimeout`. The matched string is
// returned.
func (v *VM) ExpectRE(pattern *regexp.Regexp) (string, error) {
	return v.ExpectRETimeout(pattern, v.TimeoutOr())
}

// ExpectRETimeout returns an error if the given regular expression is not
// found in vEMU's serial output within the given timeout. The matched string
// is returned.
func (v *VM) ExpectRETimeout(pattern *regexp.Regexp, timeout time.Duration) (string, error) {
	scaled := time.Duration(float64(timeout) * TimeoutMultiplier)
	str, _, err := v.gExpect.Expect(pattern, scaled)
	return str, err
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 site

import (
	"context"
	"fmt"
	"html/template"
	"time"

	"github.com/google/triage-party/pkg/hubbub"
	"github.com/google/triage-party/pkg/triage"
	"k8s.io/klog/v2"
)

const (
	// OpenStatsName is the name of the rule containing open items stats
	OpenStatsName = "__open__"
	// VelocityStatsName is the name of the rulee containing velocity stats
	VelocityStatsName = "__velocity__"
)

func (h *Handlers) collectionPage(ctx context.Context, id string, refresh bool) (*Page, error) {
	start := time.Now()
	dataAge := time.Time{}

	defer func() {
		klog.Infof("Served %q request within %s from data %s old", id, time.Since(start), time.Since(dataAge))
	}()

	s, err := h.party.LookupCollection(id)
	if err != nil {
		return nil, fmt.Errorf("lookup collection: %w", err)
	}

	sts, err := h.party.ListCollections()
	if err != nil {
		return nil, fmt.Errorf("list collections: %w", err)
	}

	var result *triage.CollectionResult
	if refresh {
		result = h.updater.ForceRefresh(ctx, id)
		klog.Infof("refresh %q result: %d items", id, len(result.RuleResults))
	} else {
		result = h.updater.Lookup(ctx, id, false)
		if result == nil {
			klog.Errorf("lookup %q returned no data", id)
			result = &triage.CollectionResult{}
		} else if result.RuleResults == nil {
			klog.Errorf("lookup %q returned no results: %+v", id, result)
		}

		klog.V(2).Infof("lookup %q result: %d items", id, len(result.RuleResults))
	}

	total := 0
	for _, o := range result.RuleResults {
		total += len(o.Items)
	}

	age := result.LatestInput
	if result.NewerThan.After(age) {
		age = result.NewerThan
	}

	unique := uniqueItems(result.RuleResults)

	p := &Page{
		ID:               s.ID,
		Version:          VERSION,
		SiteName:         h.siteName,
		Title:            s.Name,
		Collection:       s,
		Collections:      sts,
		Description:      s.Description,
		CollectionResult: result,
		Total:            len(unique),
		Types:            "Issues",
		UniqueItems:      unique,
		ResultAge:        time.Since(age),
	}

	dataAge = result.LatestInput

	if result.RuleResults == nil {
		p.Warning = template.HTML(fmt.Sprintf(`Service started %s ago, and is still downloading required data. Page will refresh after 5 seconds ...`, humanDuration(time.Since(h.startTime))))
	} else if time.Since(result.LatestInput) > h.warnAge {
		p.Notification = template.HTML(fmt.Sprintf(`Service started %s ago, and is still downloading fresh data. Data may be up to %s old. You may use Shift-Reload to force a data refresh at any time.`, humanDuration(time.Since(h.startTime)), humanDuration(time.Since(result.LatestInput))))
	}

	for _, s := range sts {
		if s.UsedForStats {
			if s.ID == VelocityStatsName {
				p.VelocityStats = h.updater.Lookup(ctx, s.ID, false)
				continue
			}
			// Older configs may not use OpenStatsName
			if s.ID == OpenStatsName || p.OpenStats == nil {
				p.OpenStats = h.updater.Lookup(ctx, s.ID, false)
				continue
			}
		}
	}

	return p, nil
}

func uniqueItems(results []*triage.RuleResult) []*hubbub.Conversation {
	items := []*hubbub.Conversation{}
	seen := map[string]bool{}

	for _, r := range results {
		for _, i := range r.Items {
			if !seen[i.URL] {
				seen[i.URL] = true
				items = append(items, i)
			}
		}
	}
	return items
}
<commit_msg>Fix data age calculations<commit_after>// Copyright 2020 Google LLC
//
// 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 site

import (
	"context"
	"fmt"
	"html/template"
	"time"

	"github.com/google/triage-party/pkg/hubbub"
	"github.com/google/triage-party/pkg/triage"
	"k8s.io/klog/v2"
)

const (
	// OpenStatsName is the name of the rule containing open items stats
	OpenStatsName = "__open__"
	// VelocityStatsName is the name of the rulee containing velocity stats
	VelocityStatsName = "__velocity__"
)

func (h *Handlers) collectionPage(ctx context.Context, id string, refresh bool) (*Page, error) {
	start := time.Now()
	dataAge := time.Time{}

	defer func() {
		klog.Infof("Served %q request within %s from data %s old", id, time.Since(start), time.Since(dataAge))
	}()

	s, err := h.party.LookupCollection(id)
	if err != nil {
		return nil, fmt.Errorf("lookup collection: %w", err)
	}

	sts, err := h.party.ListCollections()
	if err != nil {
		return nil, fmt.Errorf("list collections: %w", err)
	}

	var result *triage.CollectionResult
	if refresh {
		result = h.updater.ForceRefresh(ctx, id)
		klog.Infof("refresh %q result: %d items", id, len(result.RuleResults))
	} else {
		result = h.updater.Lookup(ctx, id, false)
		if result == nil {
			klog.Errorf("lookup %q returned no data", id)
			result = &triage.CollectionResult{}
		} else if result.RuleResults == nil {
			klog.Errorf("lookup %q returned no results: %+v", id, result)
		}

		klog.V(2).Infof("lookup %q result: %d items", id, len(result.RuleResults))
	}

	total := 0
	for _, o := range result.RuleResults {
		total += len(o.Items)
	}

	dataAge = result.LatestInput
	if result.NewerThan.After(dataAge) {
		dataAge = result.NewerThan
	}

	unique := uniqueItems(result.RuleResults)

	p := &Page{
		ID:               s.ID,
		Version:          VERSION,
		SiteName:         h.siteName,
		Title:            s.Name,
		Collection:       s,
		Collections:      sts,
		Description:      s.Description,
		CollectionResult: result,
		Total:            len(unique),
		Types:            "Issues",
		UniqueItems:      unique,
		ResultAge:        time.Since(dataAge),
	}

	if result.RuleResults == nil {
		p.Warning = template.HTML(fmt.Sprintf(`Service started %s ago, and is still downloading required data. Page will refresh after 5 seconds ...`, humanDuration(time.Since(h.startTime))))
	} else if time.Since(dataAge) > h.warnAge {
		p.Notification = template.HTML(fmt.Sprintf(`Service started %s ago, and is still downloading fresh data. Data may be up to %s old. You may use Shift-Reload to force a data refresh at any time.`, humanDuration(time.Since(h.startTime)), humanDuration(time.Since(dataAge))))
	}

	for _, s := range sts {
		if s.UsedForStats {
			if s.ID == VelocityStatsName {
				p.VelocityStats = h.updater.Lookup(ctx, s.ID, false)
				continue
			}
			// Older configs may not use OpenStatsName
			if s.ID == OpenStatsName || p.OpenStats == nil {
				p.OpenStats = h.updater.Lookup(ctx, s.ID, false)
				continue
			}
		}
	}

	return p, nil
}

func uniqueItems(results []*triage.RuleResult) []*hubbub.Conversation {
	items := []*hubbub.Conversation{}
	seen := map[string]bool{}

	for _, r := range results {
		for _, i := range r.Items {
			if !seen[i.URL] {
				seen[i.URL] = true
				items = append(items, i)
			}
		}
	}
	return items
}
<|endoftext|>


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

package gcs

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"net"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/api/googleapi"
)

// A bucket that wraps another, calling its methods in a retry loop with
// randomized exponential backoff.
type retryBucket struct {
	maxSleep time.Duration
	wrapped  Bucket
}

func newRetryBucket(
	maxSleep time.Duration,
	wrapped Bucket) (b Bucket) {
	b = &retryBucket{
		maxSleep: maxSleep,
		wrapped:  wrapped,
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func shouldRetry(err error) (b bool) {
	// HTTP 50x errors.
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code >= 500 && typed.Code < 600 {
			b = true
			return
		}
	}

	// HTTP 429 errors (GCS uses these for rate limiting).
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code == 429 {
			b = true
			return
		}
	}

	// Network errors, which tend to show up transiently when doing lots of
	// operations in parallel. For example:
	//
	//     dial tcp 74.125.203.95:443: too many open files
	//
	if _, ok := err.(*net.OpError); ok {
		b = true
		return
	}

	// The HTTP package returns ErrUnexpectedEOF in several places. This seems to
	// come up when the server terminates the connection in the middle of an
	// object read.
	if err == io.ErrUnexpectedEOF {
		b = true
		return
	}

	// The HTTP library also appears to leak EOF errors from... somewhere in its
	// guts as URL errors sometimes.
	if urlErr, ok := err.(*url.Error); ok {
		if urlErr.Err == io.EOF {
			b = true
			return
		}
	}

	// Sometimes the HTTP package helpfully encapsulates the real error in a URL
	// error.
	if urlErr, ok := err.(*url.Error); ok {
		b = shouldRetry(urlErr.Err)
		return
	}

	return
}

// Choose an appropriate delay for exponential backoff, given that we have
// already slept the given number of times for this logical request.
func chooseDelay(prevSleepCount uint) (d time.Duration) {
	const baseDelay = time.Millisecond

	// Choose a a delay in [0, 2^prevSleepCount * baseDelay).
	d = (1 << prevSleepCount) * baseDelay
	d = time.Duration(rand.Int63n(int64(d)))

	return
}

// Exponential backoff for a function that might fail.
//
// This is essentially what is described in the "Best practices" section of the
// "Upload Objects" docs:
//
//     https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload
//
// with the following exceptions:
//
//  *  We perform backoff for all operations.
//
//  *  The random component scales with the delay, so that the first sleep
//     cannot be as long as one second. The algorithm used matches the
//     description at http://en.wikipedia.org/wiki/Exponential_backoff.
//
//  *  We retry more types of errors; see shouldRetry above.
//
// State for total sleep time and number of previous sleeps is housed outside
// of this function to allow it to be "resumed" by multiple invocations of
// retryObjectReader.Read.
func expBackoff(
	ctx context.Context,
	desc string,
	maxSleep time.Duration,
	f func() error,
	prevSleepCount *uint,
	prevSleepDuration *time.Duration) (err error) {
	for {
		// Make an attempt. Stop if successful.
		err = f()
		if err == nil {
			return
		}

		// Do we want to retry?
		if !shouldRetry(err) {
			// Special case: don't spam up the logs for EOF, which io.Reader returns
			// in the normal course of things.
			if err != io.EOF {
				log.Printf(
					"Not retrying %s after error of type %T (%q): %#v",
					desc,
					err,
					err.Error(),
					err)
			}

			return
		}

		// Choose a a delay.
		d := chooseDelay(*prevSleepCount)
		*prevSleepCount++

		// Are we out of credit?
		if *prevSleepDuration+d > maxSleep {
			// Return the most recent error.
			return
		}

		// Sleep, returning early if cancelled.
		log.Printf(
			"Retrying %s after error of type %T (%q) in %v",
			desc,
			err,
			err,
			d)

		select {
		case <-ctx.Done():
			// On cancellation, return the last error we saw.
			return

		case <-time.After(d):
			*prevSleepDuration += d
			continue
		}
	}
}

// Like expBackoff, but assumes that we've never slept before (and won't need
// to sleep again).
func oneShotExpBackoff(
	ctx context.Context,
	desc string,
	maxSleep time.Duration,
	f func() error) (err error) {
	var prevSleepCount uint
	var prevSleepDuration time.Duration

	err = expBackoff(
		ctx,
		desc,
		maxSleep,
		f,
		&prevSleepCount,
		&prevSleepDuration)

	return
}

////////////////////////////////////////////////////////////////////////
// Read support
////////////////////////////////////////////////////////////////////////

type retryObjectReader struct {
	bucket *retryBucket

	// The context we should watch when sleeping for retries.
	ctx context.Context

	// What we are trying to read.
	name       string
	generation int64
	byteRange  ByteRange

	// nil when we start or have seen a permanent error.
	wrapped io.ReadCloser

	// If we've seen an error that we shouldn't retry for, this will be non-nil
	// and should be returned permanently.
	permanentErr error

	// The number of times we've slept so far, and the total amount of time we've
	// spent sleeping.
	sleepCount    uint
	sleepDuration time.Duration
}

// Set up the wrapped reader.
func (rc *retryObjectReader) setUpWrapped() (err error) {
	// Call through to create the reader.
	req := &ReadObjectRequest{
		Name:       rc.name,
		Generation: rc.generation,
		Range:      &rc.byteRange,
	}

	wrapped, err := rc.bucket.wrapped.NewReader(rc.ctx, req)
	if err != nil {
		return
	}

	rc.wrapped = wrapped
	return
}

// Set up the wrapped reader if necessary, and make one attempt to read through
// it.
//
// Clears the wrapped reader on error.
func (rc *retryObjectReader) readOnce(p []byte) (n int, err error) {
	// Set up the wrapped reader if it's not already around.
	if rc.wrapped == nil {
		err = rc.setUpWrapped()
		if err != nil {
			return
		}
	}

	// Attempt to read from it.
	n, err = rc.wrapped.Read(p)
	if err != nil {
		rc.wrapped.Close()
		rc.wrapped = nil
		return
	}

	return
}

// Invariant: we never return an error from this function unless we've given up
// on retrying. In particular, we won't return a short read because the wrapped
// reader returned a short read and an error.
func (rc *retryObjectReader) Read(p []byte) (n int, err error) {
	// Whatever we do, accumulate the bytes that we're returning to the user.
	defer func() {
		if n < 0 {
			panic(fmt.Sprintf("Negative byte count: %d", n))
		}

		rc.byteRange.Start += uint64(n)
	}()

	// If we've already decided on a permanent error, return that.
	if rc.permanentErr != nil {
		err = rc.permanentErr
		return
	}

	// If we let an error escape below, it must be a permanent one.
	defer func() {
		if err != nil {
			rc.permanentErr = err
		}
	}()

	// We will repeatedly make single attempts until we get a successful request.
	// Don't forget to accumulate the result each time.
	tryOnce := func() (err error) {
		var bytesRead int
		bytesRead, err = rc.readOnce(p)
		n += bytesRead
		p = p[bytesRead:]

		return
	}

	err = expBackoff(
		rc.ctx,
		fmt.Sprintf("Read(%q, %d)", rc.name, rc.generation),
		rc.bucket.maxSleep,
		tryOnce,
		&rc.sleepCount,
		&rc.sleepDuration)

	return
}

func (rc *retryObjectReader) Close() (err error) {
	// If we don't have a wrapped reader, there is nothing useful that we can or
	// need to do here.
	if rc.wrapped == nil {
		return
	}

	// Call through.
	err = rc.wrapped.Close()

	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (rb *retryBucket) Name() (name string) {
	name = rb.wrapped.Name()
	return
}

func (rb *retryBucket) NewReader(
	ctx context.Context,
	req *ReadObjectRequest) (rc io.ReadCloser, err error) {
	// If the user specified the latest generation, we need to figure out what
	// that is so that we can create a reader that knows how to keep a stable
	// generation despite retrying repeatedly.
	var generation int64 = req.Generation
	var sleepCount uint
	var sleepDuration time.Duration

	if generation == 0 {
		findGeneration := func() (err error) {
			o, err := rb.wrapped.StatObject(
				ctx,
				&StatObjectRequest{
					Name: req.Name,
				})

			if err != nil {
				return
			}

			generation = o.Generation
			return
		}

		err = expBackoff(
			ctx,
			fmt.Sprintf("FindLatestGeneration(%q)", req.Name),
			rb.maxSleep,
			findGeneration,
			&sleepCount,
			&sleepDuration)

		if err != nil {
			return
		}
	}

	// Choose an appropriate byte range.
	byteRange := ByteRange{0, math.MaxUint64}
	if req.Range != nil {
		byteRange = *req.Range
	}

	// Now that we know what generation we're looking for, return an appropriate
	// reader that knows how to retry when the connection fails. Make sure to
	// inherit the time spent sleeping above.
	rc = &retryObjectReader{
		bucket: rb,
		ctx:    ctx,

		name:       req.Name,
		generation: generation,
		byteRange:  byteRange,

		sleepCount:    sleepCount,
		sleepDuration: sleepDuration,
	}

	return
}

func (rb *retryBucket) CreateObject(
	ctx context.Context,
	req *CreateObjectRequest) (o *Object, err error) {
	var seeker io.ReadSeeker
	if readSeeker, ok := req.Contents.(io.ReadSeeker); ok {
		seeker = readSeeker
	} else {
		// We can't simply replay the request multiple times, because the first
		// attempt might exhaust some of the req.Contents reader, leaving
		// missing contents for the second attempt.
		//
		// So, copy out all contents and create a copy of the request that we
		// will modify to serve from memory for each call.
		data, err := ioutil.ReadAll(req.Contents)
		if err != nil {
			err = fmt.Errorf("ioutil.ReadAll: %v", err)
			return nil, err
		}
		seeker = bytes.NewReader(data)
	}

	reqCopy := *req

	// Call through with that request.
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CreateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			if _, err = seeker.Seek(0, io.SeekStart); err != nil {
				return
			}
			reqCopy.Contents = seeker
			o, err = rb.wrapped.CreateObject(ctx, &reqCopy)
			return
		})

	return
}

func (rb *retryBucket) CopyObject(
	ctx context.Context,
	req *CopyObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CopyObject(%q, %q)", req.SrcName, req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.CopyObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ComposeObjects(
	ctx context.Context,
	req *ComposeObjectsRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ComposeObjects(%q)", req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.ComposeObjects(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) StatObject(
	ctx context.Context,
	req *StatObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("StatObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.StatObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ListObjects(
	ctx context.Context,
	req *ListObjectsRequest) (listing *Listing, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ListObjects(%q)", req.Prefix),
		rb.maxSleep,
		func() (err error) {
			listing, err = rb.wrapped.ListObjects(ctx, req)
			return
		})
	return
}

func (rb *retryBucket) UpdateObject(
	ctx context.Context,
	req *UpdateObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("UpdateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.UpdateObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) DeleteObject(
	ctx context.Context,
	req *DeleteObjectRequest) (err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("DeleteObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			err = rb.wrapped.DeleteObject(ctx, req)
			return
		})

	return
}
<commit_msg>Reduce noise in the retry logs (#19)<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcs

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/api/googleapi"
)

// A bucket that wraps another, calling its methods in a retry loop with
// randomized exponential backoff.
type retryBucket struct {
	maxSleep time.Duration
	wrapped  Bucket
}

func newRetryBucket(
	maxSleep time.Duration,
	wrapped Bucket) (b Bucket) {
	b = &retryBucket{
		maxSleep: maxSleep,
		wrapped:  wrapped,
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func shouldRetry(err error) (b bool) {
	// HTTP 50x errors.
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code >= 500 && typed.Code < 600 {
			b = true
			return
		}
	}

	// HTTP 429 errors (GCS uses these for rate limiting).
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code == 429 {
			b = true
			return
		}
	}

	// Network errors, which tend to show up transiently when doing lots of
	// operations in parallel. For example:
	//
	//     dial tcp 74.125.203.95:443: too many open files
	//
	if _, ok := err.(*net.OpError); ok {
		b = true
		return
	}

	// The HTTP package returns ErrUnexpectedEOF in several places. This seems to
	// come up when the server terminates the connection in the middle of an
	// object read.
	if err == io.ErrUnexpectedEOF {
		b = true
		return
	}

	// The HTTP library also appears to leak EOF errors from... somewhere in its
	// guts as URL errors sometimes.
	if urlErr, ok := err.(*url.Error); ok {
		if urlErr.Err == io.EOF {
			b = true
			return
		}
	}

	// Sometimes the HTTP package helpfully encapsulates the real error in a URL
	// error.
	if urlErr, ok := err.(*url.Error); ok {
		b = shouldRetry(urlErr.Err)
		return
	}

	return
}

// Choose an appropriate delay for exponential backoff, given that we have
// already slept the given number of times for this logical request.
func chooseDelay(prevSleepCount uint) (d time.Duration) {
	const baseDelay = time.Millisecond

	// Choose a a delay in [0, 2^prevSleepCount * baseDelay).
	d = (1 << prevSleepCount) * baseDelay
	d = time.Duration(rand.Int63n(int64(d)))

	return
}

// Exponential backoff for a function that might fail.
//
// This is essentially what is described in the "Best practices" section of the
// "Upload Objects" docs:
//
//     https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload
//
// with the following exceptions:
//
//  *  We perform backoff for all operations.
//
//  *  The random component scales with the delay, so that the first sleep
//     cannot be as long as one second. The algorithm used matches the
//     description at http://en.wikipedia.org/wiki/Exponential_backoff.
//
//  *  We retry more types of errors; see shouldRetry above.
//
// State for total sleep time and number of previous sleeps is housed outside
// of this function to allow it to be "resumed" by multiple invocations of
// retryObjectReader.Read.
func expBackoff(
	ctx context.Context,
	desc string,
	maxSleep time.Duration,
	f func() error,
	prevSleepCount *uint,
	prevSleepDuration *time.Duration) (err error) {
	for {
		// Make an attempt. Stop if successful.
		err = f()
		if err == nil {
			return
		}

		// Do we want to retry?
		if !shouldRetry(err) {
			// Special case: don't spam up the logs for EOF, which io.Reader returns
			// in the normal course of things.
			if err != io.EOF {
				err = fmt.Errorf("not retrying %s: %w", desc, err)
			}

			return
		}

		// Choose a a delay.
		d := chooseDelay(*prevSleepCount)
		*prevSleepCount++

		// Are we out of credit?
		if *prevSleepDuration+d > maxSleep {
			// Return the most recent error.
			return
		}

		// Sleep, returning early if cancelled.
		select {
		case <-ctx.Done():
			// On cancellation, return the last error we saw.
			return

		case <-time.After(d):
			*prevSleepDuration += d
			continue
		}
	}
}

// Like expBackoff, but assumes that we've never slept before (and won't need
// to sleep again).
func oneShotExpBackoff(
	ctx context.Context,
	desc string,
	maxSleep time.Duration,
	f func() error) (err error) {
	var prevSleepCount uint
	var prevSleepDuration time.Duration

	err = expBackoff(
		ctx,
		desc,
		maxSleep,
		f,
		&prevSleepCount,
		&prevSleepDuration)

	return
}

////////////////////////////////////////////////////////////////////////
// Read support
////////////////////////////////////////////////////////////////////////

type retryObjectReader struct {
	bucket *retryBucket

	// The context we should watch when sleeping for retries.
	ctx context.Context

	// What we are trying to read.
	name       string
	generation int64
	byteRange  ByteRange

	// nil when we start or have seen a permanent error.
	wrapped io.ReadCloser

	// If we've seen an error that we shouldn't retry for, this will be non-nil
	// and should be returned permanently.
	permanentErr error

	// The number of times we've slept so far, and the total amount of time we've
	// spent sleeping.
	sleepCount    uint
	sleepDuration time.Duration
}

// Set up the wrapped reader.
func (rc *retryObjectReader) setUpWrapped() (err error) {
	// Call through to create the reader.
	req := &ReadObjectRequest{
		Name:       rc.name,
		Generation: rc.generation,
		Range:      &rc.byteRange,
	}

	wrapped, err := rc.bucket.wrapped.NewReader(rc.ctx, req)
	if err != nil {
		return
	}

	rc.wrapped = wrapped
	return
}

// Set up the wrapped reader if necessary, and make one attempt to read through
// it.
//
// Clears the wrapped reader on error.
func (rc *retryObjectReader) readOnce(p []byte) (n int, err error) {
	// Set up the wrapped reader if it's not already around.
	if rc.wrapped == nil {
		err = rc.setUpWrapped()
		if err != nil {
			return
		}
	}

	// Attempt to read from it.
	n, err = rc.wrapped.Read(p)
	if err != nil {
		rc.wrapped.Close()
		rc.wrapped = nil
		return
	}

	return
}

// Invariant: we never return an error from this function unless we've given up
// on retrying. In particular, we won't return a short read because the wrapped
// reader returned a short read and an error.
func (rc *retryObjectReader) Read(p []byte) (n int, err error) {
	// Whatever we do, accumulate the bytes that we're returning to the user.
	defer func() {
		if n < 0 {
			panic(fmt.Sprintf("Negative byte count: %d", n))
		}

		rc.byteRange.Start += uint64(n)
	}()

	// If we've already decided on a permanent error, return that.
	if rc.permanentErr != nil {
		err = rc.permanentErr
		return
	}

	// If we let an error escape below, it must be a permanent one.
	defer func() {
		if err != nil {
			rc.permanentErr = err
		}
	}()

	// We will repeatedly make single attempts until we get a successful request.
	// Don't forget to accumulate the result each time.
	tryOnce := func() (err error) {
		var bytesRead int
		bytesRead, err = rc.readOnce(p)
		n += bytesRead
		p = p[bytesRead:]

		return
	}

	err = expBackoff(
		rc.ctx,
		fmt.Sprintf("Read(%q, %d)", rc.name, rc.generation),
		rc.bucket.maxSleep,
		tryOnce,
		&rc.sleepCount,
		&rc.sleepDuration)

	return
}

func (rc *retryObjectReader) Close() (err error) {
	// If we don't have a wrapped reader, there is nothing useful that we can or
	// need to do here.
	if rc.wrapped == nil {
		return
	}

	// Call through.
	err = rc.wrapped.Close()

	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (rb *retryBucket) Name() (name string) {
	name = rb.wrapped.Name()
	return
}

func (rb *retryBucket) NewReader(
	ctx context.Context,
	req *ReadObjectRequest) (rc io.ReadCloser, err error) {
	// If the user specified the latest generation, we need to figure out what
	// that is so that we can create a reader that knows how to keep a stable
	// generation despite retrying repeatedly.
	var generation int64 = req.Generation
	var sleepCount uint
	var sleepDuration time.Duration

	if generation == 0 {
		findGeneration := func() (err error) {
			o, err := rb.wrapped.StatObject(
				ctx,
				&StatObjectRequest{
					Name: req.Name,
				})

			if err != nil {
				return
			}

			generation = o.Generation
			return
		}

		err = expBackoff(
			ctx,
			fmt.Sprintf("FindLatestGeneration(%q)", req.Name),
			rb.maxSleep,
			findGeneration,
			&sleepCount,
			&sleepDuration)

		if err != nil {
			return
		}
	}

	// Choose an appropriate byte range.
	byteRange := ByteRange{0, math.MaxUint64}
	if req.Range != nil {
		byteRange = *req.Range
	}

	// Now that we know what generation we're looking for, return an appropriate
	// reader that knows how to retry when the connection fails. Make sure to
	// inherit the time spent sleeping above.
	rc = &retryObjectReader{
		bucket: rb,
		ctx:    ctx,

		name:       req.Name,
		generation: generation,
		byteRange:  byteRange,

		sleepCount:    sleepCount,
		sleepDuration: sleepDuration,
	}

	return
}

func (rb *retryBucket) CreateObject(
	ctx context.Context,
	req *CreateObjectRequest) (o *Object, err error) {
	var seeker io.ReadSeeker
	if readSeeker, ok := req.Contents.(io.ReadSeeker); ok {
		seeker = readSeeker
	} else {
		// We can't simply replay the request multiple times, because the first
		// attempt might exhaust some of the req.Contents reader, leaving
		// missing contents for the second attempt.
		//
		// So, copy out all contents and create a copy of the request that we
		// will modify to serve from memory for each call.
		data, err := ioutil.ReadAll(req.Contents)
		if err != nil {
			err = fmt.Errorf("ioutil.ReadAll: %v", err)
			return nil, err
		}
		seeker = bytes.NewReader(data)
	}

	reqCopy := *req

	// Call through with that request.
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CreateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			if _, err = seeker.Seek(0, io.SeekStart); err != nil {
				return
			}
			reqCopy.Contents = seeker
			o, err = rb.wrapped.CreateObject(ctx, &reqCopy)
			return
		})

	return
}

func (rb *retryBucket) CopyObject(
	ctx context.Context,
	req *CopyObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CopyObject(%q, %q)", req.SrcName, req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.CopyObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ComposeObjects(
	ctx context.Context,
	req *ComposeObjectsRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ComposeObjects(%q)", req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.ComposeObjects(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) StatObject(
	ctx context.Context,
	req *StatObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("StatObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.StatObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ListObjects(
	ctx context.Context,
	req *ListObjectsRequest) (listing *Listing, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ListObjects(%q)", req.Prefix),
		rb.maxSleep,
		func() (err error) {
			listing, err = rb.wrapped.ListObjects(ctx, req)
			return
		})
	return
}

func (rb *retryBucket) UpdateObject(
	ctx context.Context,
	req *UpdateObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("UpdateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.UpdateObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) DeleteObject(
	ctx context.Context,
	req *DeleteObjectRequest) (err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("DeleteObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			err = rb.wrapped.DeleteObject(ctx, req)
			return
		})

	return
}
<|endoftext|>


<commit_before>// Package gddo is a simple client library for accessing the godoc.org API.
//
// It provides a single utility to fetch the importers of a Go package.
package gddo

import (
	"encoding/json"
	"fmt"
	"net/http"
)

// Client manages communication with the godoc.org API.
type Client struct {
	// UserAgent is used for outbound requests to godoc.org API, if set to non-empty value.
	UserAgent string
}

// GetImporters fetches the importers of Go package with specified importPath via godoc.org API.
func (c *Client) GetImporters(importPath string) (Importers, error) {
	req, err := http.NewRequest("GET", "http://api.godoc.org/importers/"+importPath, nil)
	if err != nil {
		return Importers{}, err
	}
	if c.UserAgent != "" {
		req.Header.Set("User-Agent", c.UserAgent)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return Importers{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return Importers{}, fmt.Errorf("non-200 status code: %v", resp.StatusCode)
	}
	var importers Importers
	err = json.NewDecoder(resp.Body).Decode(&importers)
	if err != nil {
		return Importers{}, err
	}
	return importers, nil
}

// Importers contains the list of Go packages that import a given Go package.
type Importers struct {
	Results []Package
}

// Package represents a Go package.
type Package struct {
	Path     string // Import path of the package.
	Synopsis string // Synopsis of the package.
}
<commit_msg>gddo: Use https for making API requests.<commit_after>// Package gddo is a simple client library for accessing the godoc.org API.
//
// It provides a single utility to fetch the importers of a Go package.
package gddo

import (
	"encoding/json"
	"fmt"
	"net/http"
)

// Client manages communication with the godoc.org API.
type Client struct {
	// UserAgent is used for outbound requests to godoc.org API, if set to non-empty value.
	UserAgent string
}

// GetImporters fetches the importers of Go package with specified importPath via godoc.org API.
func (c *Client) GetImporters(importPath string) (Importers, error) {
	req, err := http.NewRequest("GET", "https://api.godoc.org/importers/"+importPath, nil)
	if err != nil {
		return Importers{}, err
	}
	if c.UserAgent != "" {
		req.Header.Set("User-Agent", c.UserAgent)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return Importers{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return Importers{}, fmt.Errorf("non-200 status code: %v", resp.StatusCode)
	}
	var importers Importers
	err = json.NewDecoder(resp.Body).Decode(&importers)
	if err != nil {
		return Importers{}, err
	}
	return importers, nil
}

// Importers contains the list of Go packages that import a given Go package.
type Importers struct {
	Results []Package
}

// Package represents a Go package.
type Package struct {
	Path     string // Import path of the package.
	Synopsis string // Synopsis of the package.
}
<|endoftext|>


<commit_before>package gen

import (
	"github.com/leanovate/gopter"
)

// Sized derives a generator from based on size
// This honors the `MinSize` and `MaxSize` of the `GenParameters` of the test suite.
func Sized(f func(int) gopter.Gen) gopter.Gen {
	return func(params *gopter.GenParameters) *gopter.GenResult {
		size := params.Rng.Intn(params.MaxSize-params.MinSize) + params.MinSize
		return f(size)(params)
	}
}
<commit_msg>Fix MaxSize=MinSize case<commit_after>package gen

import (
	"github.com/leanovate/gopter"
)

// Sized derives a generator from based on size
// This honors the `MinSize` and `MaxSize` of the `GenParameters` of the test suite.
// Keep an eye on memory consumption, by default MaxSize is 100.
func Sized(f func(int) gopter.Gen) gopter.Gen {
	return func(params *gopter.GenParameters) *gopter.GenResult {
		var size int
		if params.MaxSize == params.MinSize {
			size = params.MaxSize
		} else {
			size = params.Rng.Intn(params.MaxSize-params.MinSize) + params.MinSize
		}
		return f(size)(params)
	}
}
<|endoftext|>


<commit_before>package main // import "sourcegraph.com/sourcegraph/gen-mocks"

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/tools/imports"
)

var (
	ifacePkgDir = flag.String("p", ".", "directory of package containing interface types")
	ifacePat    = flag.String("i", ".+Service", "regexp pattern for selecting interface types by name")
	writeFiles  = flag.Bool("w", false, "write over existing files in output directory (default: writes to stdout)")
	outDir      = flag.String("o", ".", "output directory")
	outPkg      = flag.String("outpkg", "", "output pkg name (default: same as input pkg)")
	namePrefix  = flag.String("name_prefix", "Mock", "output: name prefix of mock impl types (e.g., T -> MockT)")
	noPassArgs  = flag.String("no_pass_args", "", "don't pass args with this name from the interface method to the mock impl func")

	fset = token.NewFileSet()
)

func main() {
	flag.Parse()
	log.SetFlags(0)

	bpkg, err := build.Import(*ifacePkgDir, ".", build.FindOnly)
	if err != nil {
		log.Fatal(err)
	}

	pat, err := regexp.Compile(*ifacePat)
	if err != nil {
		log.Fatal(err)
	}

	pkgs, err := parser.ParseDir(fset, *ifacePkgDir, nil, parser.AllErrors)
	if err != nil {
		log.Fatal(err)
	}

	for _, pkg := range pkgs {
		if pkg.Name == "main" || strings.HasSuffix(pkg.Name, "_test") {
			continue
		}

		ifaces, err := readIfaces(pkg, pat)
		if err != nil {
			log.Fatal(err)
		}
		if len(ifaces) == 0 {
			log.Printf("warning: package %q has no interface types matching %q", pkg.Name, *ifacePat)
			continue
		}

		var pkgName string
		if *outPkg == "" {
			pkgName = pkg.Name
		} else {
			pkgName = *outPkg
		}

		if err := writeMockImplFiles(*outDir, pkgName, pkg.Name, bpkg.ImportPath, ifaces); err != nil {
			log.Fatal(err)
		}
	}
}

// readIfaces returns a list of interface types in pkg that should be
// mocked.
func readIfaces(pkg *ast.Package, pat *regexp.Regexp) ([]*ast.TypeSpec, error) {
	var ifaces []*ast.TypeSpec
	ast.Walk(visitFn(func(node ast.Node) bool {
		switch node := node.(type) {
		case *ast.GenDecl:
			if node.Tok == token.TYPE {
				for _, spec := range node.Specs {
					tspec := spec.(*ast.TypeSpec)
					if _, ok := tspec.Type.(*ast.InterfaceType); !ok {
						continue
					}
					if name := tspec.Name.Name; pat.MatchString(name) {
						ifaces = append(ifaces, tspec)
					}
				}
			}
			return false
		default:
			return true
		}
	}), pkg)
	return ifaces, nil
}

type visitFn func(node ast.Node) (descend bool)

func (v visitFn) Visit(node ast.Node) ast.Visitor {
	descend := v(node)
	if descend {
		return v
	} else {
		return nil
	}
}

func writeMockImplFiles(outDir, outPkg, ifacePkgName, ifacePkgPath string, svcIfaces []*ast.TypeSpec) error {
	if err := os.MkdirAll(outDir, 0700); err != nil {
		return err
	}
	decls := map[string][]ast.Decl{} // file -> decls
	for _, iface := range svcIfaces {
		filename := fset.Position(iface.Pos()).Filename
		filename = filepath.Join(outDir, strings.TrimSuffix(filepath.Base(filename), ".go")+"_mock.go")

		// mock method fields on struct
		var methFields []*ast.Field
		for _, methField := range iface.Type.(*ast.InterfaceType).Methods.List {
			if meth, ok := methField.Type.(*ast.FuncType); ok {
				methFields = append(methFields, &ast.Field{
					Names: []*ast.Ident{ast.NewIdent(methField.Names[0].Name + "_")},
					Type:  meth,
				})
			}
		}

		// struct implementation type
		mockTypeName := *namePrefix + iface.Name.Name
		implType := &ast.GenDecl{Tok: token.TYPE, Specs: []ast.Spec{&ast.TypeSpec{
			Name: ast.NewIdent(mockTypeName),
			Type: &ast.StructType{Fields: &ast.FieldList{List: methFields}},
		}}}
		decls[filename] = append(decls[filename], implType)

		// struct methods
		for _, methField := range iface.Type.(*ast.InterfaceType).Methods.List {
			if meth, ok := methField.Type.(*ast.FuncType); ok {
				synthesizeFieldNamesIfMissing(meth.Params)
				if ifacePkgName != outPkg {
					// TODO(sqs): check for import paths or dirs unequal, not pkg name
					qualifyPkgRefs(meth, ifacePkgName)
				}
				decls[filename] = append(decls[filename], &ast.FuncDecl{
					Recv: &ast.FieldList{List: []*ast.Field{
						{
							Names: []*ast.Ident{ast.NewIdent("s")},
							Type:  &ast.StarExpr{X: ast.NewIdent(mockTypeName)},
						},
					}},
					Name: ast.NewIdent(methField.Names[0].Name),
					Type: meth,
					Body: &ast.BlockStmt{List: []ast.Stmt{
						&ast.ReturnStmt{Results: []ast.Expr{
							&ast.CallExpr{
								Fun: &ast.SelectorExpr{
									X:   ast.NewIdent("s"),
									Sel: ast.NewIdent(methField.Names[0].Name + "_"),
								},
								Args:     fieldListToIdentList(meth.Params),
								Ellipsis: ellipsisIfNeeded(meth.Params),
							},
						}},
					}},
				})
			}
		}

		// compile-time implements checks
		var ifaceType ast.Expr
		if ifacePkgName == outPkg {
			ifaceType = ast.NewIdent(iface.Name.Name)
		} else {
			ifaceType = &ast.SelectorExpr{X: ast.NewIdent(ifacePkgName), Sel: ast.NewIdent(iface.Name.Name)}
		}
		decls[filename] = append(decls[filename], &ast.GenDecl{
			Tok: token.VAR,
			Specs: []ast.Spec{
				&ast.ValueSpec{
					Names: []*ast.Ident{ast.NewIdent("_")},
					Type:  ifaceType,
					Values: []ast.Expr{
						&ast.CallExpr{
							Fun:  &ast.ParenExpr{X: &ast.StarExpr{X: ast.NewIdent(mockTypeName)}},
							Args: []ast.Expr{ast.NewIdent("nil")},
						},
					},
				},
			},
		})
	}

	for filename, decls := range decls {
		file := &ast.File{
			Name:  ast.NewIdent(outPkg),
			Decls: decls,
		}
		log.Println("#", filename)
		var w io.Writer
		if *writeFiles {
			f, err := os.Create(filename)
			if err != nil {
				return err
			}
			defer f.Close()
			w = f
		} else {
			w = os.Stdout
		}

		var buf bytes.Buffer
		if err := printer.Fprint(&buf, fset, file); err != nil {
			return err
		}

		// Always put blank lines between funcs.
		src := bytes.Replace(buf.Bytes(), []byte("}\nfunc"), []byte("}\n\nfunc"), -1)

		var err error
		src, err = imports.Process(filename, src, nil)
		if err != nil {
			return err
		}

		fmt.Fprintln(w, "// generated by gen-mocks; DO NOT EDIT")
		fmt.Fprintln(w)
		w.Write(src)
	}
	return nil
}

// qualifyPkgRefs qualifies all refs to non-package-qualified non-builtin types in f so that they refer to definitions in pkg. E.g., 'func(x MyType) -> func (x pkg.MyType)'.
func qualifyPkgRefs(f *ast.FuncType, pkg string) {
	var qualify func(x ast.Expr) ast.Expr
	qualify = func(x ast.Expr) ast.Expr {
		switch y := x.(type) {
		case *ast.Ident:
			if ast.IsExported(y.Name) {
				return &ast.SelectorExpr{X: ast.NewIdent(pkg), Sel: y}
			}
		case *ast.StarExpr:
			y.X = qualify(y.X)
		case *ast.ArrayType:
			y.Elt = qualify(y.Elt)
		case *ast.MapType:
			y.Key = qualify(y.Key)
			y.Value = qualify(y.Value)
		}
		return x
	}
	for _, p := range f.Params.List {
		p.Type = qualify(p.Type)
	}
	for _, r := range f.Results.List {
		r.Type = qualify(r.Type)
	}
}

// synthesizeFieldNamesIfMissing adds synthesized variable names to fl
// if it contains fields with no name. E.g., the field list in
// `func(string, int)` would be converted to `func(v0 string, v1
// int)`.
func synthesizeFieldNamesIfMissing(fl *ast.FieldList) {
	for i, f := range fl.List {
		if len(f.Names) == 0 {
			f.Names = []*ast.Ident{ast.NewIdent(fmt.Sprintf("v%d", i))}
		}
	}
}

func fieldListToIdentList(fl *ast.FieldList) []ast.Expr {
	var fs []ast.Expr
	for _, f := range fl.List {
		for _, name := range f.Names {
			if *noPassArgs == name.Name {
				continue
			}
			x := ast.Expr(ast.NewIdent(name.Name))
			fs = append(fs, x)
		}
	}
	return fs
}

func hasEllipsis(fl *ast.FieldList) bool {
	if fl.List == nil {
		return false
	}
	lastField := fl.List[len(fl.List)-1]
	if len(lastField.Names) > 0 && lastField.Names[0].Name == *noPassArgs {
		return false
	}
	_, ok := lastField.Type.(*ast.Ellipsis)
	return ok
}

func ellipsisIfNeeded(fl *ast.FieldList) token.Pos {
	if hasEllipsis(fl) {
		return 1
	}
	return 0
}

func astString(x ast.Expr) string {
	var buf bytes.Buffer
	if err := printer.Fprint(&buf, fset, x); err != nil {
		panic(err)
	}
	return buf.String()
}
<commit_msg>omit "no pass args" from struct func types too<commit_after>package main // import "sourcegraph.com/sourcegraph/gen-mocks"

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/tools/imports"
)

var (
	ifacePkgDir = flag.String("p", ".", "directory of package containing interface types")
	ifacePat    = flag.String("i", ".+Service", "regexp pattern for selecting interface types by name")
	writeFiles  = flag.Bool("w", false, "write over existing files in output directory (default: writes to stdout)")
	outDir      = flag.String("o", ".", "output directory")
	outPkg      = flag.String("outpkg", "", "output pkg name (default: same as input pkg)")
	namePrefix  = flag.String("name_prefix", "Mock", "output: name prefix of mock impl types (e.g., T -> MockT)")
	noPassArgs  = flag.String("no_pass_args", "", "don't pass args with this name from the interface method to the mock impl func")

	fset = token.NewFileSet()
)

func main() {
	flag.Parse()
	log.SetFlags(0)

	bpkg, err := build.Import(*ifacePkgDir, ".", build.FindOnly)
	if err != nil {
		log.Fatal(err)
	}

	pat, err := regexp.Compile(*ifacePat)
	if err != nil {
		log.Fatal(err)
	}

	pkgs, err := parser.ParseDir(fset, *ifacePkgDir, nil, parser.AllErrors)
	if err != nil {
		log.Fatal(err)
	}

	for _, pkg := range pkgs {
		if pkg.Name == "main" || strings.HasSuffix(pkg.Name, "_test") {
			continue
		}

		ifaces, err := readIfaces(pkg, pat)
		if err != nil {
			log.Fatal(err)
		}
		if len(ifaces) == 0 {
			log.Printf("warning: package %q has no interface types matching %q", pkg.Name, *ifacePat)
			continue
		}

		var pkgName string
		if *outPkg == "" {
			pkgName = pkg.Name
		} else {
			pkgName = *outPkg
		}

		if err := writeMockImplFiles(*outDir, pkgName, pkg.Name, bpkg.ImportPath, ifaces); err != nil {
			log.Fatal(err)
		}
	}
}

// readIfaces returns a list of interface types in pkg that should be
// mocked.
func readIfaces(pkg *ast.Package, pat *regexp.Regexp) ([]*ast.TypeSpec, error) {
	var ifaces []*ast.TypeSpec
	ast.Walk(visitFn(func(node ast.Node) bool {
		switch node := node.(type) {
		case *ast.GenDecl:
			if node.Tok == token.TYPE {
				for _, spec := range node.Specs {
					tspec := spec.(*ast.TypeSpec)
					if _, ok := tspec.Type.(*ast.InterfaceType); !ok {
						continue
					}
					if name := tspec.Name.Name; pat.MatchString(name) {
						ifaces = append(ifaces, tspec)
					}
				}
			}
			return false
		default:
			return true
		}
	}), pkg)
	return ifaces, nil
}

type visitFn func(node ast.Node) (descend bool)

func (v visitFn) Visit(node ast.Node) ast.Visitor {
	descend := v(node)
	if descend {
		return v
	} else {
		return nil
	}
}

func writeMockImplFiles(outDir, outPkg, ifacePkgName, ifacePkgPath string, svcIfaces []*ast.TypeSpec) error {
	if err := os.MkdirAll(outDir, 0700); err != nil {
		return err
	}
	decls := map[string][]ast.Decl{} // file -> decls
	for _, iface := range svcIfaces {
		filename := fset.Position(iface.Pos()).Filename
		filename = filepath.Join(outDir, strings.TrimSuffix(filepath.Base(filename), ".go")+"_mock.go")

		// mock method fields on struct
		var methFields []*ast.Field
		for _, methField := range iface.Type.(*ast.InterfaceType).Methods.List {
			if meth, ok := methField.Type.(*ast.FuncType); ok {
				methFields = append(methFields, &ast.Field{
					Names: []*ast.Ident{ast.NewIdent(methField.Names[0].Name + "_")},
					Type:  omitNoPassArgs(meth),
				})
			}
		}

		// struct implementation type
		mockTypeName := *namePrefix + iface.Name.Name
		implType := &ast.GenDecl{Tok: token.TYPE, Specs: []ast.Spec{&ast.TypeSpec{
			Name: ast.NewIdent(mockTypeName),
			Type: &ast.StructType{Fields: &ast.FieldList{List: methFields}},
		}}}
		decls[filename] = append(decls[filename], implType)

		// struct methods
		for _, methField := range iface.Type.(*ast.InterfaceType).Methods.List {
			if meth, ok := methField.Type.(*ast.FuncType); ok {
				synthesizeFieldNamesIfMissing(meth.Params)
				if ifacePkgName != outPkg {
					// TODO(sqs): check for import paths or dirs unequal, not pkg name
					qualifyPkgRefs(meth, ifacePkgName)
				}
				decls[filename] = append(decls[filename], &ast.FuncDecl{
					Recv: &ast.FieldList{List: []*ast.Field{
						{
							Names: []*ast.Ident{ast.NewIdent("s")},
							Type:  &ast.StarExpr{X: ast.NewIdent(mockTypeName)},
						},
					}},
					Name: ast.NewIdent(methField.Names[0].Name),
					Type: meth,
					Body: &ast.BlockStmt{List: []ast.Stmt{
						&ast.ReturnStmt{Results: []ast.Expr{
							&ast.CallExpr{
								Fun: &ast.SelectorExpr{
									X:   ast.NewIdent("s"),
									Sel: ast.NewIdent(methField.Names[0].Name + "_"),
								},
								Args:     fieldListToIdentList(meth.Params),
								Ellipsis: ellipsisIfNeeded(meth.Params),
							},
						}},
					}},
				})
			}
		}

		// compile-time implements checks
		var ifaceType ast.Expr
		if ifacePkgName == outPkg {
			ifaceType = ast.NewIdent(iface.Name.Name)
		} else {
			ifaceType = &ast.SelectorExpr{X: ast.NewIdent(ifacePkgName), Sel: ast.NewIdent(iface.Name.Name)}
		}
		decls[filename] = append(decls[filename], &ast.GenDecl{
			Tok: token.VAR,
			Specs: []ast.Spec{
				&ast.ValueSpec{
					Names: []*ast.Ident{ast.NewIdent("_")},
					Type:  ifaceType,
					Values: []ast.Expr{
						&ast.CallExpr{
							Fun:  &ast.ParenExpr{X: &ast.StarExpr{X: ast.NewIdent(mockTypeName)}},
							Args: []ast.Expr{ast.NewIdent("nil")},
						},
					},
				},
			},
		})
	}

	for filename, decls := range decls {
		file := &ast.File{
			Name:  ast.NewIdent(outPkg),
			Decls: decls,
		}
		log.Println("#", filename)
		var w io.Writer
		if *writeFiles {
			f, err := os.Create(filename)
			if err != nil {
				return err
			}
			defer f.Close()
			w = f
		} else {
			w = os.Stdout
		}

		var buf bytes.Buffer
		if err := printer.Fprint(&buf, fset, file); err != nil {
			return err
		}

		// Always put blank lines between funcs.
		src := bytes.Replace(buf.Bytes(), []byte("}\nfunc"), []byte("}\n\nfunc"), -1)

		var err error
		src, err = imports.Process(filename, src, nil)
		if err != nil {
			return err
		}

		fmt.Fprintln(w, "// generated by gen-mocks; DO NOT EDIT")
		fmt.Fprintln(w)
		w.Write(src)
	}
	return nil
}

// qualifyPkgRefs qualifies all refs to non-package-qualified non-builtin types in f so that they refer to definitions in pkg. E.g., 'func(x MyType) -> func (x pkg.MyType)'.
func qualifyPkgRefs(f *ast.FuncType, pkg string) {
	var qualify func(x ast.Expr) ast.Expr
	qualify = func(x ast.Expr) ast.Expr {
		switch y := x.(type) {
		case *ast.Ident:
			if ast.IsExported(y.Name) {
				return &ast.SelectorExpr{X: ast.NewIdent(pkg), Sel: y}
			}
		case *ast.StarExpr:
			y.X = qualify(y.X)
		case *ast.ArrayType:
			y.Elt = qualify(y.Elt)
		case *ast.MapType:
			y.Key = qualify(y.Key)
			y.Value = qualify(y.Value)
		}
		return x
	}
	for _, p := range f.Params.List {
		p.Type = qualify(p.Type)
	}
	for _, r := range f.Results.List {
		r.Type = qualify(r.Type)
	}
}

// synthesizeFieldNamesIfMissing adds synthesized variable names to fl
// if it contains fields with no name. E.g., the field list in
// `func(string, int)` would be converted to `func(v0 string, v1
// int)`.
func synthesizeFieldNamesIfMissing(fl *ast.FieldList) {
	for i, f := range fl.List {
		if len(f.Names) == 0 {
			f.Names = []*ast.Ident{ast.NewIdent(fmt.Sprintf("v%d", i))}
		}
	}
}

func fieldListToIdentList(fl *ast.FieldList) []ast.Expr {
	var fs []ast.Expr
	for _, f := range fl.List {
		for _, name := range f.Names {
			if *noPassArgs == name.Name {
				continue
			}
			x := ast.Expr(ast.NewIdent(name.Name))
			fs = append(fs, x)
		}
	}
	return fs
}

func hasEllipsis(fl *ast.FieldList) bool {
	if fl.List == nil {
		return false
	}
	lastField := fl.List[len(fl.List)-1]
	if len(lastField.Names) > 0 && lastField.Names[0].Name == *noPassArgs {
		return false
	}
	_, ok := lastField.Type.(*ast.Ellipsis)
	return ok
}

func ellipsisIfNeeded(fl *ast.FieldList) token.Pos {
	if hasEllipsis(fl) {
		return 1
	}
	return 0
}

func omitNoPassArgs(ft *ast.FuncType) *ast.FuncType {
	tmp := *ft // copy
	ft = &tmp

	tmp2 := *ft.Params
	ft.Params = &tmp2
	var keepParams []*ast.Field
	for _, p := range ft.Params.List {
		if len(p.Names) == 1 && p.Names[0].Name == *noPassArgs {
			continue
		}
		keepParams = append(keepParams, p)
	}
	ft.Params.List = keepParams
	return ft
}

func astString(x ast.Expr) string {
	var buf bytes.Buffer
	if err := printer.Fprint(&buf, fset, x); err != nil {
		panic(err)
	}
	return buf.String()
}
<|endoftext|>