<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package lxdclient

import (
	"fmt"
	"net"
	"net/url"
	"strings"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"
)

// TODO(ericsnow) This file should be removed as soon as possible.
// It is based on https://github.com/lxc/lxd/blob/master/lxc/remote.go
// and the code there should instead be cleaned up and exported so we
// could use it here.

// TODO(ericsnow) Address licensing.

// addServer adds the given remote info to the provided config.
// The implementation is derived from:
//  https://github.com/lxc/lxd/blob/master/lxc/remote.go
// Note that we've removed some validation code (we don't need it).
func addServer(config *lxd.Config, server string, addr string) error {
	addr, err := fixAddr(addr, shared.PathExists)
	if err != nil {
		return err
	}

	if config.Remotes == nil {
		config.Remotes = make(map[string]lxd.RemoteConfig)
	}

	/* Actually add the remote */
	// TODO(ericsnow) Fail on collision?
	config.Remotes[server] = lxd.RemoteConfig{Addr: addr}

	return nil
}

// fixAddr is the portion of the original addServer function that fixes
// up the provided addr to be a fully qualified URL.
func fixAddr(addr string, pathExists func(string) bool) (string, error) {
	if addr == "" {
		return addr, nil
	}

	var r_scheme string
	var r_host string
	var r_port string

	/* Complex remote URL parsing */
	remote_url, err := url.Parse(addr)
	if err != nil {
		return "", err
	}

	if remote_url.Scheme != "" {
		if remote_url.Scheme != "unix" && remote_url.Scheme != "https" {
			r_scheme = "https"
		} else {
			r_scheme = remote_url.Scheme
		}
	} else if addr[0] == '/' {
		r_scheme = "unix"
	} else {
		if !pathExists(addr) {
			r_scheme = "https"
		} else {
			r_scheme = "unix"
		}
	}

	if remote_url.Host != "" {
		r_host = remote_url.Host
	} else {
		r_host = addr
	}

	host, port, err := net.SplitHostPort(r_host)
	if err == nil {
		r_host = host
		r_port = port
	} else {
		r_port = shared.DefaultPort
	}

	if r_scheme == "unix" {
		if addr[0:5] == "unix:" {
			if len(addr) > 6 && addr[0:7] == "unix://" {
				if len(addr) > 8 {
					r_host = addr[8:]
				} else {
					r_host = ""
				}
			} else {
				r_host = addr[6:]
			}
		}
		r_port = ""
	}

	if strings.Contains(r_host, ":") && !strings.HasPrefix(r_host, "[") {
		r_host = fmt.Sprintf("[%s]", r_host)
	}

	if r_port != "" {
		addr = r_scheme + "://" + r_host + ":" + r_port
	} else {
		addr = r_scheme + "://" + r_host
	}

	return addr, nil
}
<commit_msg>Clean up fixAddrs().<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package lxdclient

import (
	"fmt"
	"net"
	"net/url"
	"strings"

	"github.com/lxc/lxd"
	"github.com/lxc/lxd/shared"
)

// TODO(ericsnow) This file should be removed as soon as possible.
// It is based on https://github.com/lxc/lxd/blob/master/lxc/remote.go
// and the code there should instead be cleaned up and exported so we
// could use it here.

// TODO(ericsnow) Address licensing.

// addServer adds the given remote info to the provided config.
// The implementation is derived from:
//  https://github.com/lxc/lxd/blob/master/lxc/remote.go
// Note that we've removed some validation code (we don't need it).
func addServer(config *lxd.Config, server string, addr string) error {
	addr, err := fixAddr(addr, shared.PathExists)
	if err != nil {
		return err
	}

	if config.Remotes == nil {
		config.Remotes = make(map[string]lxd.RemoteConfig)
	}

	/* Actually add the remote */
	// TODO(ericsnow) Fail on collision?
	config.Remotes[server] = lxd.RemoteConfig{Addr: addr}

	return nil
}

// fixAddr is the portion of the original addServer function that fixes
// up the provided addr to be a fully qualified URL.
func fixAddr(addr string, pathExists func(string) bool) (string, error) {
	if addr == "" {
		// TODO(ericsnow) Return lxd.LocalRemote.Addr?
		return addr, nil
	}

	parsedURL, err := url.Parse(addr)
	if err != nil {
		return "", err
	}
	remoteURL := url.URL{
		Scheme: parsedURL.Scheme,
		Host:   parsedURL.Host,
		//Path:   parsedURL.Path,
	}

	// Fix the scheme.
	switch remoteURL.Scheme {
	case "https":
	case "unix":
	case "":
		// TODO(ericsnow) Decided based on whether or not Host is set?
		remoteURL.Scheme = "https"
		// TODO(ericsnow) Use remoteURL.Path/.Opaque here instead of addr.
		if addr[0] == '/' || pathExists(addr) {
			remoteURL.Scheme = "unix"
		}
	default:
		// TODO(ericsnow) Fail by default?
		remoteURL.Scheme = "https"
	}

	// Fix the host.
	if remoteURL.Host == "" {
		// TODO(ericsnow) Instead, make use of Path and Opaque below.
		remoteURL.Host = addr
	}
	host, port, err := net.SplitHostPort(remoteURL.Host)
	if err != nil {
		port = shared.DefaultPort
	} else {
		remoteURL.Host = host
	}

	// Special-case "unix".
	if remoteURL.Scheme == "unix" {
		// TODO(ericsnow) All this could be done earlier.
		// TODO(ericsnow) Use Path and Opaque instead of addr.
		if addr[0:5] == "unix:" {
			if len(addr) > 6 && addr[0:7] == "unix://" {
				if len(addr) > 8 {
					remoteURL.Host = addr[8:]
				} else {
					remoteURL.Host = ""
				}
			} else {
				remoteURL.Host = addr[6:]
			}
		}
		port = ""
	}

	// Handle IPv6 hosts.
	if strings.Contains(remoteURL.Host, ":") && !strings.HasPrefix(remoteURL.Host, "[") {
		remoteURL.Host = fmt.Sprintf("[%s]", remoteURL.Host)
	}

	// Add the port, if applicable.
	if port != "" {
		remoteURL.Host += ":" + port
	}

	// TODO(ericsnow) Use remoteUrl.String()
	return fmt.Sprintf("%s://%s", remoteURL.Scheme, remoteURL.Host), nil
}
<|endoftext|>


<commit_before>package torrent

import (
	"math"
	"sync"
	"time"

	"github.com/anacrolix/missinggo/v2/bitmap"
	"github.com/anacrolix/missinggo/v2/prioritybitmap"

	pp "github.com/anacrolix/torrent/peer_protocol"
)

type requestStrategyPiece interface {
	numChunks() pp.Integer
	dirtyChunks() bitmap.Bitmap
	chunkIndexRequest(i pp.Integer) request
}

type requestStrategyTorrent interface {
	numConns() int
	numReaders() int
	numPieces() int
	readerPiecePriorities() (now, readahead bitmap.Bitmap)
	ignorePieces() bitmap.Bitmap
	pendingPieces() *prioritybitmap.PriorityBitmap
}

type requestStrategyConnection interface {
	torrent() requestStrategyTorrent
	peerPieces() bitmap.Bitmap
	pieceRequestOrder() *prioritybitmap.PriorityBitmap
	fastest() bool
	stats() *ConnStats
	totalExpectingTime() time.Duration
	peerMaxRequests() int
	chunksReceivedWhileExpecting() int64
}

type requestStrategy interface {
	iterPendingPieces(requestStrategyConnection, func(pieceIndex) bool) bool
	iterUndirtiedChunks(requestStrategyPiece, func(chunkSpec) bool) bool
	nominalMaxRequests(requestStrategyConnection) int
	shouldRequestWithoutBias(requestStrategyConnection) bool
	piecePriority(requestStrategyConnection, pieceIndex, piecePriority, int) int
	hooks() requestStrategyHooks
}

type requestStrategyHooks struct {
	sentRequest    func(request)
	deletedRequest func(request)
}

type requestStrategyCallbacks interface {
	requestTimedOut(request)
}

// Favour higher priority pieces with some fuzzing to reduce overlaps and wastage across
// connections.
type requestStrategyFuzzing struct {
	requestStrategyDefaults
}

// The fastest connection downloads strictly in order of priority, while all others adhere to their
// piece inclinations.
type requestStrategyFastest struct {
	requestStrategyDefaults
}

func newRequestStrategyMaker(rs requestStrategy) requestStrategyMaker {
	return func(requestStrategyCallbacks, sync.Locker) requestStrategy {
		return rs
	}
}

func RequestStrategyFastest() requestStrategyMaker {
	return newRequestStrategyMaker(requestStrategyFastest{})
}

func RequestStrategyFuzzing() requestStrategyMaker {
	return newRequestStrategyMaker(requestStrategyFuzzing{})
}

func (requestStrategyFastest) shouldRequestWithoutBias(cn requestStrategyConnection) bool {
	if cn.torrent().numReaders() == 0 {
		return false
	}
	if cn.torrent().numConns() == 1 {
		return true
	}
	if cn.fastest() {
		return true
	}
	return false
}

// Requests are strictly by piece priority, and not duplicated until duplicateRequestTimeout is
// reached.
type requestStrategyDuplicateRequestTimeout struct {
	requestStrategyDefaults
	// How long to avoid duplicating a pending request.
	duplicateRequestTimeout time.Duration

	callbacks requestStrategyCallbacks

	// The last time we requested a chunk. Deleting the request from any connection will clear this
	// value.
	lastRequested map[request]*time.Timer
	// The lock to take when running a request timeout handler.
	timeoutLocker sync.Locker
}

// Generates a request strategy instance for a given torrent. callbacks are probably specific to the torrent.
type requestStrategyMaker func(callbacks requestStrategyCallbacks, clientLocker sync.Locker) requestStrategy

func RequestStrategyDuplicateRequestTimeout(duplicateRequestTimeout time.Duration) requestStrategyMaker {
	return func(callbacks requestStrategyCallbacks, clientLocker sync.Locker) requestStrategy {
		return requestStrategyDuplicateRequestTimeout{
			duplicateRequestTimeout: duplicateRequestTimeout,
			callbacks:               callbacks,
			lastRequested:           make(map[request]*time.Timer),
			timeoutLocker:           clientLocker,
		}
	}
}

func (rs requestStrategyDuplicateRequestTimeout) hooks() requestStrategyHooks {
	return requestStrategyHooks{
		deletedRequest: func(r request) {
			if t, ok := rs.lastRequested[r]; ok {
				t.Stop()
				delete(rs.lastRequested, r)
			}
		},
		sentRequest: rs.onSentRequest,
	}
}

func (rs requestStrategyDuplicateRequestTimeout) iterUndirtiedChunks(p requestStrategyPiece, f func(chunkSpec) bool) bool {
	for i := pp.Integer(0); i < pp.Integer(p.numChunks()); i++ {
		if p.dirtyChunks().Get(bitmap.BitIndex(i)) {
			continue
		}
		r := p.chunkIndexRequest(i)
		if rs.wouldDuplicateRecent(r) {
			continue
		}
		if !f(r.chunkSpec) {
			return false
		}
	}
	return true
}

func (requestStrategyFuzzing) piecePriority(cn requestStrategyConnection, piece pieceIndex, tpp piecePriority, prio int) int {
	switch tpp {
	case PiecePriorityNormal:
	case PiecePriorityReadahead:
		prio -= int(cn.torrent().numPieces())
	case PiecePriorityNext, PiecePriorityNow:
		prio -= 2 * int(cn.torrent().numPieces())
	default:
		panic(tpp)
	}
	prio += int(piece / 3)
	return prio
}

func (requestStrategyDuplicateRequestTimeout) iterPendingPieces(cn requestStrategyConnection, f func(pieceIndex) bool) bool {
	return iterUnbiasedPieceRequestOrder(cn, f)
}
func defaultIterPendingPieces(rs requestStrategy, cn requestStrategyConnection, f func(pieceIndex) bool) bool {
	if rs.shouldRequestWithoutBias(cn) {
		return iterUnbiasedPieceRequestOrder(cn, f)
	} else {
		return cn.pieceRequestOrder().IterTyped(func(i int) bool {
			return f(pieceIndex(i))
		})
	}
}
func (rs requestStrategyFuzzing) iterPendingPieces(cn requestStrategyConnection, cb func(pieceIndex) bool) bool {
	return defaultIterPendingPieces(rs, cn, cb)
}
func (rs requestStrategyFastest) iterPendingPieces(cn requestStrategyConnection, cb func(pieceIndex) bool) bool {
	return defaultIterPendingPieces(rs, cn, cb)
}

func (rs requestStrategyDuplicateRequestTimeout) onSentRequest(r request) {
	rs.lastRequested[r] = time.AfterFunc(rs.duplicateRequestTimeout, func() {
		rs.timeoutLocker.Lock()
		delete(rs.lastRequested, r)
		rs.timeoutLocker.Unlock()
		rs.callbacks.requestTimedOut(r)
	})
}

// The actual value to use as the maximum outbound requests.
func (rs requestStrategyDuplicateRequestTimeout) nominalMaxRequests(cn requestStrategyConnection) (ret int) {
	expectingTime := int64(cn.totalExpectingTime())
	if expectingTime == 0 {
		expectingTime = math.MaxInt64
	} else {
		expectingTime *= 2
	}
	return int(clamp(
		1,
		int64(cn.peerMaxRequests()),
		max(
			// It makes sense to always pipeline at least one connection, since latency must be
			// non-zero.
			2,
			// Request only as many as we expect to receive in the duplicateRequestTimeout
			// window. We are trying to avoid having to duplicate requests.
			cn.chunksReceivedWhileExpecting()*int64(rs.duplicateRequestTimeout)/expectingTime,
		),
	))
}
func (rs requestStrategyDuplicateRequestTimeout) wouldDuplicateRecent(r request) bool {
	// This piece has been requested on another connection, and the duplicate request timer is still
	// running.
	_, ok := rs.lastRequested[r]
	return ok
}
<commit_msg>Move request strategy doc comments onto their public functions<commit_after>package torrent

import (
	"math"
	"sync"
	"time"

	"github.com/anacrolix/missinggo/v2/bitmap"
	"github.com/anacrolix/missinggo/v2/prioritybitmap"

	pp "github.com/anacrolix/torrent/peer_protocol"
)

type requestStrategyPiece interface {
	numChunks() pp.Integer
	dirtyChunks() bitmap.Bitmap
	chunkIndexRequest(i pp.Integer) request
}

type requestStrategyTorrent interface {
	numConns() int
	numReaders() int
	numPieces() int
	readerPiecePriorities() (now, readahead bitmap.Bitmap)
	ignorePieces() bitmap.Bitmap
	pendingPieces() *prioritybitmap.PriorityBitmap
}

type requestStrategyConnection interface {
	torrent() requestStrategyTorrent
	peerPieces() bitmap.Bitmap
	pieceRequestOrder() *prioritybitmap.PriorityBitmap
	fastest() bool
	stats() *ConnStats
	totalExpectingTime() time.Duration
	peerMaxRequests() int
	chunksReceivedWhileExpecting() int64
}

type requestStrategy interface {
	iterPendingPieces(requestStrategyConnection, func(pieceIndex) bool) bool
	iterUndirtiedChunks(requestStrategyPiece, func(chunkSpec) bool) bool
	nominalMaxRequests(requestStrategyConnection) int
	shouldRequestWithoutBias(requestStrategyConnection) bool
	piecePriority(requestStrategyConnection, pieceIndex, piecePriority, int) int
	hooks() requestStrategyHooks
}

type requestStrategyHooks struct {
	sentRequest    func(request)
	deletedRequest func(request)
}

type requestStrategyCallbacks interface {
	requestTimedOut(request)
}

type requestStrategyFuzzing struct {
	requestStrategyDefaults
}

type requestStrategyFastest struct {
	requestStrategyDefaults
}

func newRequestStrategyMaker(rs requestStrategy) requestStrategyMaker {
	return func(requestStrategyCallbacks, sync.Locker) requestStrategy {
		return rs
	}
}

// The fastest connection downloads strictly in order of priority, while all others adhere to their
// piece inclinations.
func RequestStrategyFastest() requestStrategyMaker {
	return newRequestStrategyMaker(requestStrategyFastest{})
}

// Favour higher priority pieces with some fuzzing to reduce overlaps and wastage across
// connections.
func RequestStrategyFuzzing() requestStrategyMaker {
	return newRequestStrategyMaker(requestStrategyFuzzing{})
}

func (requestStrategyFastest) shouldRequestWithoutBias(cn requestStrategyConnection) bool {
	if cn.torrent().numReaders() == 0 {
		return false
	}
	if cn.torrent().numConns() == 1 {
		return true
	}
	if cn.fastest() {
		return true
	}
	return false
}

type requestStrategyDuplicateRequestTimeout struct {
	requestStrategyDefaults
	// How long to avoid duplicating a pending request.
	duplicateRequestTimeout time.Duration

	callbacks requestStrategyCallbacks

	// The last time we requested a chunk. Deleting the request from any connection will clear this
	// value.
	lastRequested map[request]*time.Timer
	// The lock to take when running a request timeout handler.
	timeoutLocker sync.Locker
}

// Generates a request strategy instance for a given torrent. callbacks are probably specific to the torrent.
type requestStrategyMaker func(callbacks requestStrategyCallbacks, clientLocker sync.Locker) requestStrategy

// Requests are strictly by piece priority, and not duplicated until duplicateRequestTimeout is
// reached.
func RequestStrategyDuplicateRequestTimeout(duplicateRequestTimeout time.Duration) requestStrategyMaker {
	return func(callbacks requestStrategyCallbacks, clientLocker sync.Locker) requestStrategy {
		return requestStrategyDuplicateRequestTimeout{
			duplicateRequestTimeout: duplicateRequestTimeout,
			callbacks:               callbacks,
			lastRequested:           make(map[request]*time.Timer),
			timeoutLocker:           clientLocker,
		}
	}
}

func (rs requestStrategyDuplicateRequestTimeout) hooks() requestStrategyHooks {
	return requestStrategyHooks{
		deletedRequest: func(r request) {
			if t, ok := rs.lastRequested[r]; ok {
				t.Stop()
				delete(rs.lastRequested, r)
			}
		},
		sentRequest: rs.onSentRequest,
	}
}

func (rs requestStrategyDuplicateRequestTimeout) iterUndirtiedChunks(p requestStrategyPiece, f func(chunkSpec) bool) bool {
	for i := pp.Integer(0); i < pp.Integer(p.numChunks()); i++ {
		if p.dirtyChunks().Get(bitmap.BitIndex(i)) {
			continue
		}
		r := p.chunkIndexRequest(i)
		if rs.wouldDuplicateRecent(r) {
			continue
		}
		if !f(r.chunkSpec) {
			return false
		}
	}
	return true
}

func (requestStrategyFuzzing) piecePriority(cn requestStrategyConnection, piece pieceIndex, tpp piecePriority, prio int) int {
	switch tpp {
	case PiecePriorityNormal:
	case PiecePriorityReadahead:
		prio -= int(cn.torrent().numPieces())
	case PiecePriorityNext, PiecePriorityNow:
		prio -= 2 * int(cn.torrent().numPieces())
	default:
		panic(tpp)
	}
	prio += int(piece / 3)
	return prio
}

func (requestStrategyDuplicateRequestTimeout) iterPendingPieces(cn requestStrategyConnection, f func(pieceIndex) bool) bool {
	return iterUnbiasedPieceRequestOrder(cn, f)
}
func defaultIterPendingPieces(rs requestStrategy, cn requestStrategyConnection, f func(pieceIndex) bool) bool {
	if rs.shouldRequestWithoutBias(cn) {
		return iterUnbiasedPieceRequestOrder(cn, f)
	} else {
		return cn.pieceRequestOrder().IterTyped(func(i int) bool {
			return f(pieceIndex(i))
		})
	}
}
func (rs requestStrategyFuzzing) iterPendingPieces(cn requestStrategyConnection, cb func(pieceIndex) bool) bool {
	return defaultIterPendingPieces(rs, cn, cb)
}
func (rs requestStrategyFastest) iterPendingPieces(cn requestStrategyConnection, cb func(pieceIndex) bool) bool {
	return defaultIterPendingPieces(rs, cn, cb)
}

func (rs requestStrategyDuplicateRequestTimeout) onSentRequest(r request) {
	rs.lastRequested[r] = time.AfterFunc(rs.duplicateRequestTimeout, func() {
		rs.timeoutLocker.Lock()
		delete(rs.lastRequested, r)
		rs.timeoutLocker.Unlock()
		rs.callbacks.requestTimedOut(r)
	})
}

// The actual value to use as the maximum outbound requests.
func (rs requestStrategyDuplicateRequestTimeout) nominalMaxRequests(cn requestStrategyConnection) (ret int) {
	expectingTime := int64(cn.totalExpectingTime())
	if expectingTime == 0 {
		expectingTime = math.MaxInt64
	} else {
		expectingTime *= 2
	}
	return int(clamp(
		1,
		int64(cn.peerMaxRequests()),
		max(
			// It makes sense to always pipeline at least one connection, since latency must be
			// non-zero.
			2,
			// Request only as many as we expect to receive in the duplicateRequestTimeout
			// window. We are trying to avoid having to duplicate requests.
			cn.chunksReceivedWhileExpecting()*int64(rs.duplicateRequestTimeout)/expectingTime,
		),
	))
}
func (rs requestStrategyDuplicateRequestTimeout) wouldDuplicateRecent(r request) bool {
	// This piece has been requested on another connection, and the duplicate request timer is still
	// running.
	_, ok := rs.lastRequested[r]
	return ok
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"sort"
	"strings"
	"time"
	"unicode/utf8"

	ioutil "io/ioutil"

	git "github.com/libgit2/git2go"
	"github.com/mgutz/ansi"
)

var (
	Red    string = ansi.ColorCode("red")
	Yellow        = ansi.ColorCode("yellow")
	Green         = ansi.ColorCode("green")
)

const (
	CachePath = ".git/go_gb_cache.json"
)

func exit(msg string, args ...interface{}) {
	msg = fmt.Sprintf(msg, args...)
	fmt.Println(msg)
	os.Exit(1)
}

func NewRepo() *git.Repository {
	repo, err := git.OpenRepository(".")
	if err != nil {
		wd, _ := os.Getwd()
		exit("Could not open repository at '%s'", wd)
	}
	return repo
}

func NewBranchIterator(repo *git.Repository) *git.BranchIterator {
	i, err := repo.NewBranchIterator(git.BranchLocal)
	if err != nil {
		wd, _ := os.Getwd()
		exit("Failed to list branches for '%s'", wd)
	}
	return i
}

func LookupBaseOid(repo *git.Repository, baseBranchName string) *git.Oid {
	base_branch, err := repo.LookupBranch(baseBranchName, git.BranchLocal)
	if err != nil {
		exit("Error looking up base branch '%s'", baseBranchName)
	}

	return base_branch.Target()
}

type Comparison struct {
	Repo    *git.Repository
	BaseOid *git.Oid
	Branch  *git.Branch
	Oid     *git.Oid

	IsMerged bool
	Ahead    int
	Behind   int
}

func NewComparison(repo *git.Repository, base_oid *git.Oid, branch *git.Branch, store CacheStore) *Comparison {
	c := new(Comparison)

	c.Repo = repo
	c.BaseOid = base_oid

	c.Branch = branch
	c.Oid = branch.Target()

	cache := store[c.CacheKey()]

	if cache != nil {
		c.Ahead = cache.Ahead
		c.Behind = cache.Behind
		c.IsMerged = cache.IsMerged
	} else {
		c.IsMerged = false
		c.Ahead = -1
		c.Behind = -1
	}

	return c
}

func (c *Comparison) Name() string {
	name, err := c.Branch.Name()
	if err != nil {
		exit("Can't get branch name for '%s'", c.Oid)
	}
	return name
}

func (c *Comparison) IsHead() bool {
	head, err := c.Branch.IsHead()
	if err != nil {
		exit("Can't get IsHead for '%s'", c.Name())
	}
	return head
}

func (c *Comparison) Commit() *git.Commit {
	commit, err := c.Repo.LookupCommit(c.Oid)
	if err != nil {
		exit("Could not lookup commit '%s'.", c.Oid.String())
	}
	return commit
}

func (c *Comparison) ColorCode() string {
	hours, _ := time.ParseDuration("336h") // two weeks
	two_weeks := time.Now().Add(-hours)

	if c.IsHead() {
		return Green
	} else if c.When().Before(two_weeks) {
		return Red
	} else {
		return Yellow
	}
}

func (c *Comparison) When() time.Time {
	sig := c.Commit().Committer()
	return sig.When
}

func (c *Comparison) FormattedWhen() string {
	return c.When().Format("2006-01-02 15:04PM")
}

func (c *Comparison) CacheKey() string {
	strs := []string{c.BaseOid.String(), c.Oid.String()}
	return strings.Join(strs, "..")
}

func (c *Comparison) SetIsMerged() {
	if c.Oid.String() == c.BaseOid.String() {
		c.IsMerged = true
	} else {
		merged, err := c.Repo.DescendantOf(c.BaseOid, c.Oid)
		if err != nil {
			exit("Could not get descendant of '%s' and '%s'.", c.BaseOid.String(), c.Oid.String())
		}
		c.IsMerged = merged
	}
}

func (c *Comparison) SetAheadBehind() {
	var err error
	c.Ahead, c.Behind, err = c.Repo.AheadBehind(c.Oid, c.BaseOid)
	if err != nil {
		exit("Error getting ahead/behind", c.BaseOid.String())
	}
}

func (c *Comparison) Execute() {
	if c.Ahead > -1 && c.Behind > -1 {
		return
	}

	c.SetIsMerged()
	c.SetAheadBehind()
}

type Comparisons []*Comparison

func (cs Comparisons) MaxBranchLength() int {
	max := 30

	for _, comp := range cs {
		length := utf8.RuneCountInString(comp.Name())
		if length > max {
			max = length
		}
	}
	return max
}

type ComparisonsByWhen Comparisons

func (a ComparisonsByWhen) Len() int {
	return len(a)
}

func (a ComparisonsByWhen) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

func (a ComparisonsByWhen) Less(i, j int) bool {
	return a[i].When().Unix() < a[j].When().Unix()
}

type Options struct {
	Ahead      int
	Behind     int
	Merged     bool
	NoMerged   bool
	ClearCache bool
	BaseBranch string
}

func NewOptions() *Options {
	o := new(Options)

	flag.IntVar(&o.Ahead, "ahead", -1, "only show branches that are <ahead> commits ahead.")
	flag.IntVar(&o.Behind, "behind", -1, "only show branches that are <behind> commits behind.")
	flag.BoolVar(&o.Merged, "merged", false, "only show branches that are merged.")
	flag.BoolVar(&o.NoMerged, "no-merged", false, "only show branches that are not merged.")
	flag.BoolVar(&o.ClearCache, "clear-cache", false, "clear cache of comparisons.")
	flag.StringVar(&o.BaseBranch, "base", "master", "base branch to make comparisons against.")

	flag.Parse()

	return o
}

type CacheStore map[string]*Comparison

func NewCacheStore() CacheStore {
	bits, err := ioutil.ReadFile(CachePath)
	if err != nil {
		// no-op: `cache.json` will be written on exit.
	}

	y := make(CacheStore)
	_ = json.Unmarshal(bits, &y)

	return y
}

func (store *CacheStore) WriteToFile() error {
	b, err := json.Marshal(store)
	if err != nil {
		fmt.Printf("Could not save cache to file.")
	}
	ioutil.WriteFile(CachePath, b, 0644)
	return nil
}

func main() {
	opts := NewOptions()

	if opts.ClearCache {
		os.Remove(CachePath)
	}

	store := NewCacheStore()

	repo := NewRepo()
	branch_iterator := NewBranchIterator(repo)
	base_oid := LookupBaseOid(repo, opts.BaseBranch)

	comparisons := make(Comparisons, 0)

	// type BranchIteratorFunc func(*Branch, BranchType) error
	branch_iterator.ForEach(func(branch *git.Branch, btype git.BranchType) error {
		comp := NewComparison(repo, base_oid, branch, store)
		comparisons = append(comparisons, comp)
		return nil
	})

	sort.Sort(ComparisonsByWhen(comparisons))

	branch_length := comparisons.MaxBranchLength()

	for _, comp := range comparisons {
		comp.Execute()

		if comp.Name() == opts.BaseBranch {
			fmt.Printf(
				"%s%s * %-*s\n",
				comp.ColorCode(),
				comp.FormattedWhen(),
				branch_length, // http://stackoverflow.com/a/28870241
				comp.Name())
			continue
		}

		merged_string := ""
		if comp.IsMerged {
			merged_string = "(merged)"
		}

		if opts.Ahead != -1 && opts.Ahead != comp.Ahead {
			continue
		}

		if opts.Behind != -1 && opts.Behind != comp.Behind {
			continue
		}

		if opts.Merged && !comp.IsMerged {
			continue
		}

		if opts.NoMerged && comp.IsMerged {
			continue
		}

		fmt.Printf(
			"%s%s | %-*s | behind: %4d | ahead: %4d %s\n",
			comp.ColorCode(),
			comp.FormattedWhen(),
			branch_length, // http://stackoverflow.com/a/28870241
			comp.Name(),
			comp.Behind,
			comp.Ahead,
			merged_string)

		store[comp.CacheKey()] = comp
	}

	store.WriteToFile()

}
<commit_msg>Replace flag with codegansta/cli.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"strings"
	"time"
	"unicode/utf8"

	ioutil "io/ioutil"

	"github.com/codegangsta/cli"
	git "github.com/libgit2/git2go"
	"github.com/mgutz/ansi"
)

var (
	Red    string = ansi.ColorCode("red")
	Yellow        = ansi.ColorCode("yellow")
	Green         = ansi.ColorCode("green")
	Reset         = ansi.ColorCode("reset")
	Bold          = ansi.ColorCode("reset+b")
)

const CachePath = ".git/go_gb_cache.json"

func exit(msg string, args ...interface{}) {
	msg = fmt.Sprintf(msg, args...)
	fmt.Println(msg)
	os.Exit(1)
}

func NewRepo() *git.Repository {
	repo, err := git.OpenRepository(".")
	if err != nil {
		wd, _ := os.Getwd()
		exit("Could not open repository at '%s'", wd)
	}
	return repo
}

func NewBranchIterator(repo *git.Repository) *git.BranchIterator {
	i, err := repo.NewBranchIterator(git.BranchLocal)
	if err != nil {
		wd, _ := os.Getwd()
		exit("Failed to list branches for '%s'", wd)
	}
	return i
}

func LookupBaseOid(repo *git.Repository, baseBranchName string) *git.Oid {
	base_branch, err := repo.LookupBranch(baseBranchName, git.BranchLocal)
	if err != nil {
		exit("Error looking up base branch '%s'", baseBranchName)
	}

	return base_branch.Target()
}

type Comparison struct {
	Repo    *git.Repository
	BaseOid *git.Oid
	Branch  *git.Branch
	Oid     *git.Oid

	IsMerged bool
	Ahead    int
	Behind   int
}

func NewComparison(repo *git.Repository, base_oid *git.Oid, branch *git.Branch, store CacheStore) *Comparison {
	c := new(Comparison)

	c.Repo = repo
	c.BaseOid = base_oid

	c.Branch = branch
	c.Oid = branch.Target()

	cache := store[c.CacheKey()]

	if cache != nil {
		c.Ahead = cache.Ahead
		c.Behind = cache.Behind
		c.IsMerged = cache.IsMerged
	} else {
		c.IsMerged = false
		c.Ahead = -1
		c.Behind = -1
	}

	return c
}

func (c *Comparison) Name() string {
	name, err := c.Branch.Name()
	if err != nil {
		exit("Can't get branch name for '%s'", c.Oid)
	}
	return name
}

func (c *Comparison) IsHead() bool {
	head, err := c.Branch.IsHead()
	if err != nil {
		exit("Can't get IsHead for '%s'", c.Name())
	}
	return head
}

func (c *Comparison) Commit() *git.Commit {
	commit, err := c.Repo.LookupCommit(c.Oid)
	if err != nil {
		exit("Could not lookup commit '%s'.", c.Oid.String())
	}
	return commit
}

func (c *Comparison) ColorCode() string {
	hours, _ := time.ParseDuration("336h") // two weeks
	two_weeks := time.Now().Add(-hours)

	if c.IsHead() {
		return Green
	} else if c.When().Before(two_weeks) {
		return Red
	} else {
		return Yellow
	}
}

func (c *Comparison) When() time.Time {
	sig := c.Commit().Committer()
	return sig.When
}

func (c *Comparison) FormattedWhen() string {
	return c.When().Format("2006-01-02 15:04PM")
}

func (c *Comparison) CacheKey() string {
	strs := []string{c.BaseOid.String(), c.Oid.String()}
	return strings.Join(strs, "..")
}

func (c *Comparison) SetIsMerged() {
	if c.Oid.String() == c.BaseOid.String() {
		c.IsMerged = true
	} else {
		merged, err := c.Repo.DescendantOf(c.BaseOid, c.Oid)
		if err != nil {
			exit("Could not get descendant of '%s' and '%s'.", c.BaseOid.String(), c.Oid.String())
		}
		c.IsMerged = merged
	}
}

func (c *Comparison) SetAheadBehind() {
	var err error
	c.Ahead, c.Behind, err = c.Repo.AheadBehind(c.Oid, c.BaseOid)
	if err != nil {
		exit("Error getting ahead/behind", c.BaseOid.String())
	}
}

func (c *Comparison) Execute() {
	if c.Ahead > -1 && c.Behind > -1 {
		return
	}

	c.SetIsMerged()
	c.SetAheadBehind()
}

type Comparisons []*Comparison

func (cs Comparisons) MaxBranchLength() int {
	max := 30

	for _, comp := range cs {
		length := utf8.RuneCountInString(comp.Name())
		if length > max {
			max = length
		}
	}
	return max
}

type ComparisonsByWhen Comparisons

func (a ComparisonsByWhen) Len() int {
	return len(a)
}

func (a ComparisonsByWhen) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

func (a ComparisonsByWhen) Less(i, j int) bool {
	return a[i].When().Unix() < a[j].When().Unix()
}

type CacheStore map[string]*Comparison

func NewCacheStore() CacheStore {
	bits, err := ioutil.ReadFile(CachePath)
	if err != nil {
		// no-op: `cache.json` will be written on exit.
	}

	y := make(CacheStore)
	_ = json.Unmarshal(bits, &y)

	return y
}

func (store *CacheStore) WriteToFile() error {
	b, err := json.Marshal(store)
	if err != nil {
		fmt.Printf("Could not save cache to file.")
	}
	ioutil.WriteFile(CachePath, b, 0644)
	return nil
}

func run(ctx *cli.Context) {
	baseBranch := "master"
	if len(ctx.Args()) > 0 {
		baseBranch = ctx.Args().First()
	}

	if ctx.Bool("clear-cache") {
		os.Remove(CachePath)
	}

	store := NewCacheStore()

	repo := NewRepo()
	branch_iterator := NewBranchIterator(repo)
	base_oid := LookupBaseOid(repo, baseBranch)

	comparisons := make(Comparisons, 0)

	// type BranchIteratorFunc func(*Branch, BranchType) error
	branch_iterator.ForEach(func(branch *git.Branch, btype git.BranchType) error {
		comp := NewComparison(repo, base_oid, branch, store)
		comparisons = append(comparisons, comp)
		return nil
	})

	sort.Sort(ComparisonsByWhen(comparisons))

	branch_length := comparisons.MaxBranchLength()

	for _, comp := range comparisons {
		comp.Execute()

		if comp.Name() == baseBranch {
			fmt.Printf(
				"%s%s%s * %-*s\n",
				Bold,
				comp.ColorCode(),
				comp.FormattedWhen(),
				branch_length, // http://stackoverflow.com/a/28870241
				comp.Name())
			continue
		}

		merged_string := ""
		if comp.IsMerged {
			merged_string = "(merged)"
		}

		if ctx.Int("ahead") != -1 && ctx.Int("ahead") != comp.Ahead {
			continue
		}

		if ctx.Int("behind") != -1 && ctx.Int("behind") != comp.Behind {
			continue
		}

		if ctx.Bool("merged") && !comp.IsMerged {
			continue
		}

		if ctx.Bool("no-merged") && comp.IsMerged {
			continue
		}

		fmt.Printf(
			"%s%s%s | %-*s | behind: %4d | ahead: %4d %s\n",
			Reset,
			comp.ColorCode(),
			comp.FormattedWhen(),
			branch_length, // http://stackoverflow.com/a/28870241
			comp.Name(),
			comp.Behind,
			comp.Ahead,
			merged_string)

		store[comp.CacheKey()] = comp
	}

	store.WriteToFile()
}

func main() {
	app := cli.NewApp()
	app.Name = "gb"
	app.Usage = "A better way to list git branches in your terminal."
	app.Version = "1.0.0"
	app.Action = run

	app.Flags = []cli.Flag{
		cli.IntFlag{Name: "ahead", Value: -1, Usage: "only show branches that are <ahead> commits ahead."},
		cli.IntFlag{Name: "behind", Value: -1, Usage: "only show branches that are <behind> commits behind."},
		cli.BoolFlag{Name: "merged", Usage: "only show branches that are merged."},
		cli.BoolFlag{Name: "no-merged", Usage: "only show branches that are not merged."},
		cli.BoolFlag{Name: "clear-cache", Usage: "clear cache of comparisons."},
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2017 Intel Corporation
//
// 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 bat

import (
	"context"
	"strconv"

	"github.com/ciao-project/ciao/uuid"
)

// TenantSummary is a short form of Tenant
type TenantSummary struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

// TenantsListResponse stores a list of tenants retrieved by listTenants
type TenantsListResponse struct {
	Tenants []TenantSummary `json:"tenants"`
}

// TenantConfig stores the configurable attributes of a tenant.
type TenantConfig struct {
	Name       string `json:"name"`
	SubnetBits int    `json:"subnet_bits"`
}

// GetAllTenants retrieves a list of all tenants in the cluster by calling
// ciao-cli tenant list -all. An error will be returned if the following
// environment variables are not set; CIAO_ADMIN_CLIENT_CERT_FILE,
// CIAO_CONTROLLER.
func GetAllTenants(ctx context.Context) (TenantsListResponse, error) {
	var tenants TenantsListResponse

	args := []string{"tenant", "list", "-all", "-f", "{{tojson .}}"}
	err := RunCIAOCLIAsAdminJS(ctx, "", args, &tenants)
	if err != nil {
		return tenants, err
	}

	return tenants, nil
}

// CreateTenant creates a new tenant with the given config.
// It calls ciao-cli tenant create.
func CreateTenant(ctx context.Context, config TenantConfig) (TenantSummary, error) {
	var summary TenantSummary

	args := []string{"tenant", "create", "-tenant", uuid.Generate().String(), "-name", config.Name, "-cidr-prefix-size", strconv.Itoa(config.SubnetBits), "-f", "{{tojson .}}"}
	err := RunCIAOCLIAsAdminJS(ctx, "", args, &summary)

	return summary, err
}

// UpdateTenant updates a new tenant with the given config.
// It calls ciao-cli tenant update.
func UpdateTenant(ctx context.Context, ID string, config TenantConfig) error {
	args := []string{"tenant", "update", "-for-tenant", ID}
	if config.Name != "" {
		name := []string{"-name", config.Name}
		args = append(args, name...)
	}

	if config.SubnetBits != 0 {
		bits := []string{"-cidr-prefix-size", strconv.Itoa(config.SubnetBits)}
		args = append(args, bits...)
	}

	_, err := RunCIAOCLIAsAdmin(ctx, "", args)
	return err
}

// GetTenantConfig retrieves the configuration for the given tenant.
func GetTenantConfig(ctx context.Context, ID string) (TenantConfig, error) {
	var config TenantConfig

	args := []string{"tenant", "list", "-config", "-for-tenant", ID, "-f", "{{tojson .}}"}
	err := RunCIAOCLIAsAdminJS(ctx, "", args, &config)

	return config, err
}

// DeleteTenant will delete the given tenant.
// It calls ciao-cli tenant delete.
func DeleteTenant(ctx context.Context, ID string) error {
	args := []string{"tenant", "delete", "-tenant", ID}
	_, err := RunCIAOCLIAsAdmin(ctx, "", args)

	return err
}
<commit_msg>bat: Update for new permissions fields<commit_after>//
// Copyright (c) 2017 Intel Corporation
//
// 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 bat

import (
	"context"
	"strconv"

	"github.com/ciao-project/ciao/uuid"
)

// TenantSummary is a short form of Tenant
type TenantSummary struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

// TenantsListResponse stores a list of tenants retrieved by listTenants
type TenantsListResponse struct {
	Tenants []TenantSummary `json:"tenants"`
}

// TenantConfig stores the configurable attributes of a tenant.
type TenantConfig struct {
	Name        string `json:"name"`
	SubnetBits  int    `json:"subnet_bits"`
	Permissions struct {
		PrivilegedContainers bool `json:"privileged_containers"`
	} `json:"permissions"`
}

// GetAllTenants retrieves a list of all tenants in the cluster by calling
// ciao-cli tenant list -all. An error will be returned if the following
// environment variables are not set; CIAO_ADMIN_CLIENT_CERT_FILE,
// CIAO_CONTROLLER.
func GetAllTenants(ctx context.Context) (TenantsListResponse, error) {
	var tenants TenantsListResponse

	args := []string{"tenant", "list", "-all", "-f", "{{tojson .}}"}
	err := RunCIAOCLIAsAdminJS(ctx, "", args, &tenants)
	if err != nil {
		return tenants, err
	}

	return tenants, nil
}

// CreateTenant creates a new tenant with the given config.
// It calls ciao-cli tenant create.
func CreateTenant(ctx context.Context, config TenantConfig) (TenantSummary, error) {
	var summary TenantSummary

	args := []string{"tenant", "create", "-tenant", uuid.Generate().String(), "-name", config.Name, "-cidr-prefix-size", strconv.Itoa(config.SubnetBits), "-f", "{{tojson .}}"}

	if config.Permissions.PrivilegedContainers == true {
		args = append(args, "-create-privileged-containers")
	}

	err := RunCIAOCLIAsAdminJS(ctx, "", args, &summary)

	return summary, err
}

// UpdateTenant updates a new tenant with the given config.
// It calls ciao-cli tenant update.
func UpdateTenant(ctx context.Context, ID string, config TenantConfig) error {
	args := []string{"tenant", "update", "-for-tenant", ID}
	if config.Name != "" {
		name := []string{"-name", config.Name}
		args = append(args, name...)
	}

	if config.SubnetBits != 0 {
		bits := []string{"-cidr-prefix-size", strconv.Itoa(config.SubnetBits)}
		args = append(args, bits...)
	}

	if config.Permissions.PrivilegedContainers == true {
		args = append(args, "-create-privileged-containers")
	}

	_, err := RunCIAOCLIAsAdmin(ctx, "", args)
	return err
}

// GetTenantConfig retrieves the configuration for the given tenant.
func GetTenantConfig(ctx context.Context, ID string) (TenantConfig, error) {
	var config TenantConfig

	args := []string{"tenant", "list", "-config", "-for-tenant", ID, "-f", "{{tojson .}}"}
	err := RunCIAOCLIAsAdminJS(ctx, "", args, &config)

	return config, err
}

// DeleteTenant will delete the given tenant.
// It calls ciao-cli tenant delete.
func DeleteTenant(ctx context.Context, ID string) error {
	args := []string{"tenant", "delete", "-tenant", ID}
	_, err := RunCIAOCLIAsAdmin(ctx, "", args)

	return err
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"runtime"
)

const (
	queueLength       = 8000
	workerConcurrency = 64
	divider           = "--------------------------------------------------------"
)

var (
	logF *os.File
)

func init() {
	logF, err := os.OpenFile(
		getEnv("GR_BATCH_LOG", "batch.log"),
		os.O_RDWR|os.O_CREATE|os.O_APPEND,
		0644,
	)
	if err != nil {
		panic(err)
	}

	log.SetOutput(logF)
}

func main() {
	defer logF.Close()
	defer db.Close()

	runtime.GOMAXPROCS(workerConcurrency + 1)
	idQueue := make(chan int, queueLength)

	log.Println(divider)
	for i := 0; i < workerConcurrency; i++ {
		go workerLoop(i, idQueue)
	}
	schedulerLoop(idQueue)
}
<commit_msg>Concurrency must be 1<commit_after>package main

import (
	"log"
	"os"
	"runtime"
)

const (
	queueLength       = 10
	workerConcurrency = 1
	divider           = "--------------------------------------------------------"
)

var (
	logF *os.File
)

func init() {
	logF, err := os.OpenFile(
		getEnv("GR_BATCH_LOG", "batch.log"),
		os.O_RDWR|os.O_CREATE|os.O_APPEND,
		0644,
	)
	if err != nil {
		panic(err)
	}

	log.SetOutput(logF)
}

func main() {
	defer logF.Close()
	defer db.Close()

	runtime.GOMAXPROCS(workerConcurrency + 1)
	idQueue := make(chan int, queueLength)

	log.Println(divider)
	for i := 0; i < workerConcurrency; i++ {
		go workerLoop(i, idQueue)
	}
	schedulerLoop(idQueue)
}
<|endoftext|>


<commit_before>package bcast

/*
   bcast package for Go. Broadcasting on a set of channels.

   Copyright © 2013 Alexander I.Grafov <grafov@gmail.com>.
   All rights reserved.
   Use of this source code is governed by a BSD-style
   license that can be found in the LICENSE file.
*/

import (
	"testing"
	"time"
)

// Create new broadcast group.
// Join two members.
func TestNewGroupAndJoin(t *testing.T) {
	group := NewGroup()
	member1 := group.Join()
	member2 := group.Join()
	if member1.group != member2.group {
		t.Fatal("group for these members must be same")
	}
}

// Create new broadcast group.
// Join two members.
// Unjoin first member.
func TestUnjoin(t *testing.T) {
	group := NewGroup()
	member1 := group.Join()
	member2 := group.Join()
	if len(group.out) != 2 {
		t.Fatal("incorrect length of `out` slice")
	}

	go group.Broadcasting(0)

	member1.Close()
	time.Sleep(1 * time.Second) // because Broadcasting executed concurrently
	if len(group.out) > 1 || group.out[0] != member2.In {
		t.Fatal("unjoin member does not work")
	}
}

// Create new broadcast group.
// Join 12 members.
// Broadcast one integer from each member.
func TestBroadcast(t *testing.T) {
	var valcount int

	group := NewGroup()

	for i := 1; i <= 12; i++ {
		go func(i int, group *Group) {
			m := group.Join()
			m.Send(i)

			for {
				val := m.Recv()
				if val.(int) == i {
					t.Fatal("sent value was received by sender")
				}
				valcount++
			}
		}(i, group)
	}

	group.Broadcasting(100 * time.Millisecond)
	if valcount != 12*12-12 { // number of channels * number of messages - number of channels
		t.Fatal("not all messages broadcasted")
	}
}

// Create new broadcast group.
// Join 12 members.
// Broadcast one integer from each member.
func TestBroadcastFromMember(t *testing.T) {
	var valcount int

	group := NewGroup()

	for i := 1; i <= 12; i++ {
		go func(i int, group *Group) {
			m := group.Join()
			m.Send(i)

			for {
				val := m.Recv()
				if val.(int) == i {
					t.Fatal("sent value was received by sender")
				}
				valcount++
			}
		}(i, group)
	}

	group.Broadcasting(100 * time.Millisecond)
	if valcount != 12*12-12 { // number of channels * number of messages - number of channels
		t.Fatal("not all messages broadcasted")
	}
}

// Create new broadcast group.
// Join 12 members.
// Make group broadcast to all group members.
func TestGroupBroadcast(t *testing.T) {
	var valcount int

	group := NewGroup()

	for i := 1; i <= 12; i++ {
		go func(i int, group *Group) {
			m := group.Join()
			if val := m.Recv(); val != "group message" {
				t.Fatal("incorrect message received")
			}
			valcount++
		}(i, group)
	}

	go group.Broadcasting(200 * time.Millisecond)
	group.Send("group message")
	time.Sleep(150 * time.Millisecond)

	if valcount != 12 {
		t.Fatal("not all messages broadcasted")
	}
}

// Create new broadcast group.
// Join 128 members.
// Make group broadcast to all members.
func TestBroadcastOnLargeNumberOfMembers(t *testing.T) {
	const max = 128
	var valcount int

	group := NewGroup()
	for i := 1; i <= max; i++ {
		go func(i int, group *Group) {
			m := group.Join()
			m.Send(i)
			for {
				if val := m.Recv(); val.(int) == i {
					t.Fatal("sent value should not received by sender")
				}
				valcount++
			}
		}(i, group)
	}

	group.Broadcasting(1 * time.Second)
	if valcount != max*max-max { // number of channels * number of messages - number of channels
		t.Fatalf("not all messages broadcasted (%d/%d)", valcount, max*max-max)
	}
}
<commit_msg>making obvious the race condition without -race flag<commit_after>package bcast

/*
   bcast package for Go. Broadcasting on a set of channels.

   Copyright © 2013 Alexander I.Grafov <grafov@gmail.com>.
   All rights reserved.
   Use of this source code is governed by a BSD-style
   license that can be found in the LICENSE file.
*/

import (
	"testing"
	"time"
)

// Create new broadcast group.
// Join two members.
func TestNewGroupAndJoin(t *testing.T) {
	group := NewGroup()
	member1 := group.Join()
	member2 := group.Join()
	if member1.group != member2.group {
		t.Fatal("group for these members must be same")
	}
}

// Create new broadcast group.
// Join two members.
// Unjoin first member.
func TestUnjoin(t *testing.T) {
	group := NewGroup()
	member1 := group.Join()
	member2 := group.Join()
	if len(group.out) != 2 {
		t.Fatal("incorrect length of `out` slice")
	}

	go group.Broadcasting(0)

	member1.Close()
	time.Sleep(1 * time.Second) // because Broadcasting executed concurrently
	if len(group.out) > 1 || group.out[0] != member2.In {
		t.Fatal("unjoin member does not work")
	}
}

// Create new broadcast group.
// Join 12 members.
// Broadcast one integer from each member.
func TestBroadcast(t *testing.T) {
	var valcount int

	group := NewGroup()

	for i := 1; i <= 12; i++ {
		go func(i int, group *Group) {
			m := group.Join()
			m.Send(i)

			for {
				val := m.Recv()
				if val.(int) == i {
					t.Fatal("sent value was received by sender")
				}
				valcount++
			}
		}(i, group)
	}

	group.Broadcasting(100 * time.Millisecond)
	if valcount != 12*12-12 { // number of channels * number of messages - number of channels
		t.Fatal("not all messages broadcasted")
	}
}

// Create new broadcast group.
// Join 12 members.
// Broadcast one integer from each member.
func TestBroadcastFromMember(t *testing.T) {
	var valcount int

	group := NewGroup()

	for i := 1; i <= 12; i++ {
		go func(i int, group *Group) {
			m := group.Join()
			m.Send(i)

			for {
				val := m.Recv()
				if val.(int) == i {
					t.Fatal("sent value was received by sender")
				}
				valcount++
			}
		}(i, group)
	}

	group.Broadcasting(100 * time.Millisecond)
	if valcount != 12*12-12 { // number of channels * number of messages - number of channels
		t.Fatal("not all messages broadcasted")
	}
}

// Create new broadcast group.
// Join 12 members.
// Make group broadcast to all group members.
func TestGroupBroadcast(t *testing.T) {
	var valcount int

	group := NewGroup()

	for i := 1; i <= 12; i++ {
		go func(i int, group *Group) {
			m := group.Join()
			if val := m.Recv(); val != "group message" {
				t.Fatal("incorrect message received")
			}
			valcount++
		}(i, group)
	}

	go group.Broadcasting(200 * time.Millisecond)
	group.Send("group message")
	time.Sleep(150 * time.Millisecond)

	if valcount != 12 {
		t.Fatal("not all messages broadcasted")
	}
}

// Create new broadcast group.
// Join 128 members.
// Make group broadcast to all members.
func TestBroadcastOnLargeNumberOfMembers(t *testing.T) {
	const max = 1024
	var valcount int

	group := NewGroup()
	for i := 1; i <= max; i++ {
		go func(i int, group *Group) {
			m := group.Join()
			m.Send(i)
			for {
				if val := m.Recv(); val.(int) == i {
					t.Fatal("sent value should not received by sender")
				}
				valcount++
			}
		}(i, group)
	}

	group.Broadcasting(1 * time.Second)
	if valcount != max*max-max { // number of channels * number of messages - number of channels
		t.Fatalf("not all messages broadcasted (%d/%d)", valcount, max*max-max)
	}
}
<|endoftext|>


<commit_before>package golog

import "fmt"
import "strconv"
import "testing"
import "github.com/mndrix/golog/read"
import "github.com/mndrix/golog/term"

func BenchmarkTrue(b *testing.B) {
    m := NewMachine()
    g := read.Term_(`true.`)
    for i := 0; i < b.N; i++ {
        _ = m.ProveAll(g)
    }
}

func BenchmarkAppend(b *testing.B) {
    m := NewMachine().Consult(`
        append([], A, A).   % test same variable name as other clauses
        append([A|B], C, [A|D]) :-
            append(B, C, D).
    `)
    g := read.Term_(`append([a,b,c], [d,e], List).`)

    for i := 0; i < b.N; i++ {
        _ = m.ProveAll(g)
    }
}

// unify two compounds terms with deep structure. unification succeeds
func BenchmarkUnifyDeep(b *testing.B) {
    x := read.Term_(`a(b(c(d(e(f(g(h(i(j))))))))).`)
    y := read.Term_(`a(b(c(d(e(f(g(h(i(X))))))))).`)

    env := term.NewBindings()
    for i := 0; i < b.N; i++ {
        _, _ = x.Unify(env, y)
    }
}

// unify two compounds terms with deep structure. unification fails
func BenchmarkUnifyDeepFail(b *testing.B) {
    x := read.Term_(`a(b(c(d(e(f(g(h(i(j))))))))).`)
    y := read.Term_(`a(b(c(d(e(f(g(h(i(x))))))))).`)

    env := term.NewBindings()
    for i := 0; i < b.N; i++ {
        _, _ = x.Unify(env, y)
    }
}

func BenchmarkUnificationHash(b *testing.B) {
    x := read.Term_(`a(b(c(d(e(f(g(h(i(j))))))))).`)
    for i := 0; i < b.N; i++ {
        _ = term.UnificationHash([]term.Term{x}, 64, true)
    }
}

// test performance of a standard maplist implementation
func BenchmarkMaplist(b *testing.B) {
    m := NewMachine().Consult(`
        always_a(_, a).

        maplist(C, A, B) :-
            maplist_(A, B, C).

        maplist_([], [], _).
        maplist_([B|D], [C|E], A) :-
            call(A, B, C),
            maplist_(D, E, A).
    `)
    g := read.Term_(`maplist(always_a, [1,2,3,4,5], As).`)

    for i := 0; i < b.N; i++ {
        _ = m.ProveAll(g)
    }
}

// traditional, naive reverse benchmark
// The Art of Prolog by Sterling, etal says that reversing a 30 element
// list using this technique does 496 reductions.  From this we can
// calculate a rough measure of Golog's LIPS.
func BenchmarkNaiveReverse(b *testing.B) {
    m := NewMachine().Consult(`
        append([], A, A).
        append([A|B], C, [A|D]) :-
            append(B, C, D).

        reverse([],[]).
        reverse([X|Xs], Zs) :-
            reverse(Xs, Ys),
            append(Ys, [X], Zs).
    `)
    g := read.Term_(`reverse([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30], As).`)

    for i := 0; i < b.N; i++ {
        _ = m.ProveAll(g)
    }
}

func BenchmarkRead(b *testing.B) {
    for i := 0; i < b.N; i++ {
        _ = read.Term_(`reverse([1,2,3,4,5,6,7], Xs).`)
    }
}

// Low level benchmarks to test Go's implementation
func init() {   // avoid import errors when low level benchmarks comment out
    _ = fmt.Sprintf("")
    _ = strconv.Itoa(1)
}
/*
func BenchmarkLowLevelCompareUint64(b *testing.B) {
    var nintendo uint64 = 282429536481
    var other uint64 = 387429489
    for i := 0; i < b.N; i++ {
        if nintendo == other {
            // do nothing
        }
    }
}
func BenchmarkLowLevelCompareString(b *testing.B) {
    nintendo := "nintendo"
    other := "other"
    for i := 0; i < b.N; i++ {
        if nintendo == other {
            // do nothing
        }
    }
}
func BenchmarkLowLevelBitwise(b *testing.B) {
    var nintendo uint64 = 282429536481
    var other uint64 = 387429489
    for i := 0; i < b.N; i++ {
        if nintendo & other == nintendo {
            // do nothing
        }
    }
}
func BenchmarkLowLevelFloatBinaryExponent(b *testing.B) {
	f := 3.1415
    for i := 0; i < b.N; i++ {
        _ = strconv.FormatFloat(f, 'b', 0, 64)
    }
}
func BenchmarkLowLevelFloatDecimalExponent(b *testing.B) {
	f := 3.1415
    for i := 0; i < b.N; i++ {
        _ = strconv.FormatFloat(f, 'e', 64, 64)
    }
}
func BenchmarkLowLevelIntDecimal(b *testing.B) {
	var x uint64 = 1967
    for i := 0; i < b.N; i++ {
        _ = fmt.Sprintf("%d", x)
    }
}
func BenchmarkLowLevelIntHex(b *testing.B) {
	var x uint64 = 1967
    for i := 0; i < b.N; i++ {
        _ = fmt.Sprintf("%x", x)
    }
}
*/
<commit_msg>Benchmarks for Go interface variants<commit_after>package golog

import "fmt"
import "strconv"
import "testing"
import "github.com/mndrix/golog/read"
import "github.com/mndrix/golog/term"

func BenchmarkTrue(b *testing.B) {
    m := NewMachine()
    g := read.Term_(`true.`)
    for i := 0; i < b.N; i++ {
        _ = m.ProveAll(g)
    }
}

func BenchmarkAppend(b *testing.B) {
    m := NewMachine().Consult(`
        append([], A, A).   % test same variable name as other clauses
        append([A|B], C, [A|D]) :-
            append(B, C, D).
    `)
    g := read.Term_(`append([a,b,c], [d,e], List).`)

    for i := 0; i < b.N; i++ {
        _ = m.ProveAll(g)
    }
}

// unify two compounds terms with deep structure. unification succeeds
func BenchmarkUnifyDeep(b *testing.B) {
    x := read.Term_(`a(b(c(d(e(f(g(h(i(j))))))))).`)
    y := read.Term_(`a(b(c(d(e(f(g(h(i(X))))))))).`)

    env := term.NewBindings()
    for i := 0; i < b.N; i++ {
        _, _ = x.Unify(env, y)
    }
}

// unify two compounds terms with deep structure. unification fails
func BenchmarkUnifyDeepFail(b *testing.B) {
    x := read.Term_(`a(b(c(d(e(f(g(h(i(j))))))))).`)
    y := read.Term_(`a(b(c(d(e(f(g(h(i(x))))))))).`)

    env := term.NewBindings()
    for i := 0; i < b.N; i++ {
        _, _ = x.Unify(env, y)
    }
}

func BenchmarkUnificationHash(b *testing.B) {
    x := read.Term_(`a(b(c(d(e(f(g(h(i(j))))))))).`)
    for i := 0; i < b.N; i++ {
        _ = term.UnificationHash([]term.Term{x}, 64, true)
    }
}

// test performance of a standard maplist implementation
func BenchmarkMaplist(b *testing.B) {
    m := NewMachine().Consult(`
        always_a(_, a).

        maplist(C, A, B) :-
            maplist_(A, B, C).

        maplist_([], [], _).
        maplist_([B|D], [C|E], A) :-
            call(A, B, C),
            maplist_(D, E, A).
    `)
    g := read.Term_(`maplist(always_a, [1,2,3,4,5], As).`)

    for i := 0; i < b.N; i++ {
        _ = m.ProveAll(g)
    }
}

// traditional, naive reverse benchmark
// The Art of Prolog by Sterling, etal says that reversing a 30 element
// list using this technique does 496 reductions.  From this we can
// calculate a rough measure of Golog's LIPS.
func BenchmarkNaiveReverse(b *testing.B) {
    m := NewMachine().Consult(`
        append([], A, A).
        append([A|B], C, [A|D]) :-
            append(B, C, D).

        reverse([],[]).
        reverse([X|Xs], Zs) :-
            reverse(Xs, Ys),
            append(Ys, [X], Zs).
    `)
    g := read.Term_(`reverse([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30], As).`)

    for i := 0; i < b.N; i++ {
        _ = m.ProveAll(g)
    }
}

func BenchmarkRead(b *testing.B) {
    for i := 0; i < b.N; i++ {
        _ = read.Term_(`reverse([1,2,3,4,5,6,7], Xs).`)
    }
}

// Low level benchmarks to test Go's implementation
func init() {   // avoid import errors when low level benchmarks comment out
    _ = fmt.Sprintf("")
    _ = strconv.Itoa(1)
}
/*
func BenchmarkLowLevelCompareUint64(b *testing.B) {
    var nintendo uint64 = 282429536481
    var other uint64 = 387429489
    for i := 0; i < b.N; i++ {
        if nintendo == other {
            // do nothing
        }
    }
}
func BenchmarkLowLevelCompareString(b *testing.B) {
    nintendo := "nintendo"
    other := "other"
    for i := 0; i < b.N; i++ {
        if nintendo == other {
            // do nothing
        }
    }
}
func BenchmarkLowLevelBitwise(b *testing.B) {
    var nintendo uint64 = 282429536481
    var other uint64 = 387429489
    for i := 0; i < b.N; i++ {
        if nintendo & other == nintendo {
            // do nothing
        }
    }
}
func BenchmarkLowLevelFloatBinaryExponent(b *testing.B) {
	f := 3.1415
    for i := 0; i < b.N; i++ {
        _ = strconv.FormatFloat(f, 'b', 0, 64)
    }
}
func BenchmarkLowLevelFloatDecimalExponent(b *testing.B) {
	f := 3.1415
    for i := 0; i < b.N; i++ {
        _ = strconv.FormatFloat(f, 'e', 64, 64)
    }
}
func BenchmarkLowLevelIntDecimal(b *testing.B) {
	var x uint64 = 1967
    for i := 0; i < b.N; i++ {
        _ = fmt.Sprintf("%d", x)
    }
}
func BenchmarkLowLevelIntHex(b *testing.B) {
	var x uint64 = 1967
    for i := 0; i < b.N; i++ {
        _ = fmt.Sprintf("%x", x)
    }
}

// benchmarks to compare performance on interface-related code
type AnInterface interface {
    AMethod() int
}
type ImplementationOne int
func (*ImplementationOne) AMethod() int { return 1 }
type ImplementationTwo int
func (*ImplementationTwo) AMethod() int { return 2 }

func NotAMethod(x AnInterface) int {
    switch x.(type) {
        case *ImplementationOne:
            return 1
        case *ImplementationTwo:
            return 2
    }
    panic("impossible")
}

// how expensive is it to call a method?
func BenchmarkInterfaceMethod(b *testing.B) {
    var x AnInterface
    num := 100
    x = (*ImplementationOne)(&num)

    for i := 0; i < b.N; i++ {
        _ = x.AMethod()
    }
}

// how expensive is it to call a function that acts like a method?
func BenchmarkInterfaceFunction(b *testing.B) {
    var x AnInterface
    num := 100
    x = (*ImplementationOne)(&num)

    for i := 0; i < b.N; i++ {
        _ = NotAMethod(x)
    }
}

// how expensive is it to inline a type switch that acts like a method?
func BenchmarkInterfaceInlineTypeSwitch(b *testing.B) {
    var x AnInterface
    num := 100
    x = (*ImplementationOne)(&num)

    for i := 0; i < b.N; i++ {
        var y int
        switch x.(type) {
            case *ImplementationOne:
                y = 1
            case *ImplementationTwo:
                y = 2
        }
        _ = y
    }
}

// how expensive is it a manually-implemented type switch?
func BenchmarkInterfaceManualTypeSwitch(b *testing.B) {
    var x AnInterface
    num := 100
    x = (*ImplementationOne)(&num)

    for i := 0; i < b.N; i++ {
        var y int
        kind := x.AMethod()
        switch kind {
            case 1:
                y = 1
            case 2:
                y = 2
        }
        _ = y
    }
}

*/
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var maxCount uint64
	b.Run("read-random-badger", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			var val badger.KVItem
			for pb.Next() {
				key := newKey()
				if bdb.Get(key, &val); val.Value() != nil {
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("badger: Maximum no. of keys found: %d", maxCount)
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var maxCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if _, err := rdb.Get(key); err == nil {
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("rocks: Maximum no. of keys found: %d", maxCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var maxCount uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				_ = lmdbEnv.View(func(txn *lmdb.Txn) error {
					_, err := txn.Get(lmdbDBI, key)
					if err != nil {
						return err
					}
					count++
					return nil
				})
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("lmdb: Maximum no. of keys found: %d", maxCount)
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<commit_msg>Assert correct value size<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"sync/atomic"
	"testing"

	"github.com/bmatsuo/lmdb-go/lmdb"
	"github.com/dgraph-io/badger"
	"github.com/dgraph-io/badger-bench/store"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	ctx           = context.Background()
	numKeys       = flag.Float64("keys_mil", 10.0, "How many million keys to write.")
	flagDir       = flag.String("dir", "bench-tmp", "Where data is temporarily stored.")
	flagValueSize = flag.Int("valsz", 128, "Size of each value.")
)

const Mi int = 1000000
const Mf float64 = 1000000

func getBadger() (*badger.KV, error) {
	opt := badger.DefaultOptions
	opt.MapTablesTo = table.LoadToRAM
	opt.Dir = *flagDir + "/badger"
	opt.ValueDir = opt.Dir
	fmt.Println(opt.Dir)
	opt.DoNotCompact = true
	opt.ValueGCThreshold = 0.0
	return badger.NewKV(&opt)
}

func getRocks() *store.Store {
	rdb, err := store.NewReadOnlyStore(*flagDir + "/rocks")
	y.Check(err)
	return rdb
}

func getLmdb() *lmdb.Env {
	lmdbEnv, err := lmdb.NewEnv()
	y.Check(err)
	err = lmdbEnv.SetMaxDBs(1)
	y.Check(err)
	err = lmdbEnv.SetMapSize(1 << 38) // ~273Gb
	y.Check(err)

	err = lmdbEnv.Open(*flagDir+"/lmdb", lmdb.Readonly, 0777)
	y.Check(err)
	return lmdbEnv
}

func newKey() []byte {
	k := rand.Int() % int(*numKeys*Mf)
	key := fmt.Sprintf("vsz=%05d-k=%010d", *flagValueSize, k) // 22 bytes.
	return []byte(key)
}

func print(count int) {
	if count%100000 == 0 {
		fmt.Printf(".")
	} else if count%Mi == 0 {
		fmt.Printf("-")
	}
}

func BenchmarkReadRandomBadger(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	defer bdb.Close()

	var maxCount uint64
	b.Run("read-random-badger", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			var val badger.KVItem
			for pb.Next() {
				key := newKey()
				if bdb.Get(key, &val); val.Value() != nil {
					if len(val.Value()) != *flagValueSize {
						b.Logf("Value size %d != %d", len(val.Value()), *flagValueSize)
						b.FailNow()
					}
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("badger: Maximum no. of keys found: %d", maxCount)
}

func BenchmarkReadRandomRocks(b *testing.B) {
	rdb := getRocks()
	defer rdb.Close()

	var maxCount uint64
	b.Run("read-random-rocks", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				if v, err := rdb.Get(key); err == nil {

					if len(v.Data()) != *flagValueSize {
						b.Logf("Value size %d != %d", len(v.Data()), *flagValueSize)
						b.FailNow()
					}
					count++
				}
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("rocks: Maximum no. of keys found: %d", maxCount)
}

func BenchmarkReadRandomLmdb(b *testing.B) {
	lmdbEnv := getLmdb()
	defer lmdbEnv.Close()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	var maxCount uint64
	b.Run("read-random-lmdb", func(b *testing.B) {
		var totalCount uint64
		b.RunParallel(func(pb *testing.PB) {
			var count uint64
			for pb.Next() {
				key := newKey()
				_ = lmdbEnv.View(func(txn *lmdb.Txn) error {
					v, err := txn.Get(lmdbDBI, key)
					if err != nil {
						return err
					}
					if len(v) != *flagValueSize {
						b.Logf("Value size %d != %d", v, *flagValueSize)
						b.FailNow()
					}
					count++
					return nil
				})
			}
			atomic.AddUint64(&totalCount, count)
		})
		if totalCount > maxCount {
			maxCount = totalCount
		}
	})
	b.Logf("lmdb: Maximum no. of keys found: %d", maxCount)
}

func safecopy(dst []byte, src []byte) []byte {
	if cap(dst) < len(src) {
		dst = make([]byte, len(src))
	}
	dst = dst[0:len(src)]
	copy(dst, src)
	return dst
}

func BenchmarkIterateRocks(b *testing.B) {
	rdb := getRocks()
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("rocksdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			itr := rdb.NewIterator()
			var count int
			for itr.SeekToFirst(); itr.Valid(); itr.Next() {
				{
					// do some processing.
					k = safecopy(k, itr.Key().Data())
					v = safecopy(v, itr.Value().Data())
				}
				if len(v) != *flagValueSize {
					b.Logf("Value size %d != %d", v, *flagValueSize)
					b.FailNow()
				}

				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateLmdb(b *testing.B) {
	lmdbEnv := getLmdb()

	var lmdbDBI lmdb.DBI
	// Acquire handle
	err := lmdbEnv.View(func(txn *lmdb.Txn) error {
		var err error
		lmdbDBI, err = txn.OpenDBI("bench", 0)
		return err
	})
	y.Check(err)
	defer lmdbEnv.CloseDBI(lmdbDBI)

	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("lmdb-iterate", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			err = lmdbEnv.View(func(txn *lmdb.Txn) error {
				cur, err := txn.OpenCursor(lmdbDBI)
				if err != nil {
					return err
				}
				defer cur.Close()

				for {
					k1, v1, err := cur.Get(nil, nil, lmdb.Next)
					if lmdb.IsNotFound(err) {
						return nil
					}
					if err != nil {
						return err
					}
					if len(v1) != *flagValueSize {
						b.Logf("Value size %d != %d", v, *flagValueSize)
						b.FailNow()
					}

					//fmt.Printf("%s %s\n", k, v)

					// do some processing.
					k = safecopy(k, k1)
					v = safecopy(v, v1)

					count++
					print(count)
					if count > 2*Mi {
						break
					}
				}
				return nil
			})
			y.Check(err)
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerOnlyKeys(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	b.ResetTimer()

	b.Run("badger-iterate-onlykeys", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			// 100 = size, 0 = num workers, false = fwd direction.
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
				}
				count++
				print(count)
				if count > 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func BenchmarkIterateBadgerWithValues(b *testing.B) {
	bdb, err := getBadger()
	y.Check(err)
	k := make([]byte, 1024)
	v := make([]byte, Mi)
	b.ResetTimer()

	b.Run("badger-iterate-withvals", func(b *testing.B) {
		for j := 0; j < b.N; j++ {
			var count int
			opt := badger.IteratorOptions{}
			opt.PrefetchSize = 10000
			opt.FetchValues = true
			itr := bdb.NewIterator(opt)
			for itr.Rewind(); itr.Valid(); itr.Next() {
				item := itr.Item()
				{
					// do some processing.
					k = safecopy(k, item.Key())
					v = safecopy(v, item.Value())
				}
				if len(v) != *flagValueSize {
					b.Logf("Value size %d != %d", v, *flagValueSize)
					b.FailNow()
				}

				count++
				print(count)
				if count >= 2*Mi {
					break
				}
			}
			b.Logf("[%d] Counted %d keys\n", j, count)
		}
	})
}

func TestMain(m *testing.M) {
	flag.Parse()
	// call flag.Parse() here if TestMain uses flags
	go http.ListenAndServe(":8080", nil)
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>// This file contains code taken from the "go" source tree, governed
// by the go license:

// Copyright (c) 2012 The Go Authors. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//    * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//    * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//    * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package main

import (
	"crypto/tls"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

const buildV = false

// From go/src/cmd/go/vcs.go

// securityMode specifies whether a function should make network
// calls using insecure transports (eg, plain text HTTP).
// The zero value is "secure".
type securityMode int

const (
	secure securityMode = iota
	insecure
)

// repoRoot represents a version control system, a repo, and a root of
// where to put it on disk.
type repoRoot struct {
	vcs string

	// repo is the repository URL, including scheme
	repo string

	// root is the import path corresponding to the root of the
	// repository
	root string
}

// repoRootForImportDynamic finds a *repoRoot for a custom domain that's not
// statically known by repoRootForImportPathStatic.
//
// This handles custom import paths like "name.tld/pkg/foo" or just "name.tld".
func queryRepoRoot(importPath string, security securityMode) (*repoRoot, error) {
	slash := strings.Index(importPath, "/")
	if slash < 0 {
		slash = len(importPath)
	}
	host := importPath[:slash]
	if !strings.Contains(host, ".") {
		return nil, errors.New("import path does not begin with hostname")
	}
	urlStr, body, err := httpsOrHTTP(importPath, security)
	if err != nil {
		msg := "https fetch: %v"
		if security == insecure {
			msg = "http/" + msg
		}
		return nil, fmt.Errorf(msg, err)
	}
	defer body.Close()
	imports, err := parseMetaGoImports(body)
	if err != nil {
		return nil, fmt.Errorf("parsing %s: %v", importPath, err)
	}
	// Find the matched meta import.
	mmi, err := matchGoImport(imports, importPath)
	if err != nil {
		if err != errNoMatch {
			return nil, fmt.Errorf("parse %s: %v", urlStr, err)
		}
		return nil, fmt.Errorf("parse %s: no go-import meta tags", urlStr)
	}
	if buildV {
		log.Printf("get %q: found meta tag %#v at %s", importPath, mmi, urlStr)
	}
	// If the import was "uni.edu/bob/project", which said the
	// prefix was "uni.edu" and the RepoRoot was "evilroot.com",
	// make sure we don't trust Bob and check out evilroot.com to
	// "uni.edu" yet (possibly overwriting/preempting another
	// non-evil student).  Instead, first verify the root and see
	// if it matches Bob's claim.
	if mmi.Prefix != importPath {
		if buildV {
			log.Printf("get %q: verifying non-authoritative meta tag", importPath)
		}
		urlStr0 := urlStr
		var imports []metaImport
		urlStr, imports, err = metaImportsForPrefix(mmi.Prefix, security)
		if err != nil {
			return nil, err
		}
		metaImport2, err := matchGoImport(imports, importPath)
		if err != nil || mmi != metaImport2 {
			return nil, fmt.Errorf("%s and %s disagree about go-import for %s", urlStr0, urlStr, mmi.Prefix)
		}
	}

	if !strings.Contains(mmi.RepoRoot, "://") {
		return nil, fmt.Errorf("%s: invalid repo root %q; no scheme", urlStr, mmi.RepoRoot)
	}
	return &repoRoot{
		vcs:  mmi.VCS,
		repo: mmi.RepoRoot,
		root: mmi.Prefix,
	}, nil
}

var (
	fetchCacheMu sync.Mutex
	fetchCache   = map[string]fetchResult{} // key is metaImportsForPrefix's importPrefix
)

// metaImportsForPrefix takes a package's root import path as declared in a <meta> tag
// and returns its HTML discovery URL and the parsed metaImport lines
// found on the page.
//
// The importPath is of the form "golang.org/x/tools".
// It is an error if no imports are found.
// urlStr will still be valid if err != nil.
// The returned urlStr will be of the form "https://golang.org/x/tools?go-get=1"
func metaImportsForPrefix(importPrefix string, security securityMode) (urlStr string, imports []metaImport, err error) {
	setCache := func(res fetchResult) fetchResult {
		fetchCacheMu.Lock()
		defer fetchCacheMu.Unlock()
		fetchCache[importPrefix] = res
		return res
	}

	fetch := func() fetchResult {
		fetchCacheMu.Lock()
		if res, ok := fetchCache[importPrefix]; ok {
			fetchCacheMu.Unlock()
			return res
		}
		fetchCacheMu.Unlock()

		urlStr, body, err := httpsOrHTTP(importPrefix, security)
		if err != nil {
			return setCache(fetchResult{urlStr: urlStr, err: fmt.Errorf("fetch %s: %v", urlStr, err)})
		}
		imports, err := parseMetaGoImports(body)
		if err != nil {
			return setCache(fetchResult{urlStr: urlStr, err: fmt.Errorf("parsing %s: %v", urlStr, err)})
		}
		if len(imports) == 0 {
			err = fmt.Errorf("fetch %s: no go-import meta tag", urlStr)
		}
		return setCache(fetchResult{urlStr: urlStr, imports: imports, err: err})
	}

	res := fetch()
	return res.urlStr, res.imports, res.err
}

type fetchResult struct {
	urlStr  string // e.g. "https://foo.com/x/bar?go-get=1"
	imports []metaImport
	err     error
}

// metaImport represents the parsed <meta name="go-import"
// content="prefix vcs reporoot" /> tags from HTML files.
type metaImport struct {
	Prefix, VCS, RepoRoot string
}

// errNoMatch is returned from matchGoImport when there's no applicable match.
var errNoMatch = errors.New("no import match")

// matchGoImport returns the metaImport from imports matching importPath.
// An error is returned if there are multiple matches.
// errNoMatch is returned if none match.
func matchGoImport(imports []metaImport, importPath string) (_ metaImport, err error) {
	match := -1
	for i, im := range imports {
		if !strings.HasPrefix(importPath, im.Prefix) {
			continue
		}
		if match != -1 {
			err = fmt.Errorf("multiple meta tags match import path %q", importPath)
			return
		}
		match = i
	}
	if match == -1 {
		err = errNoMatch
		return
	}
	return imports[match], nil
}

// From go/src/cmd/go/http.go

// httpClient is the default HTTP client, but a variable so it can be
// changed by tests, without modifying http.DefaultClient.
var httpClient = http.DefaultClient

// impatientInsecureHTTPClient is used in -insecure mode,
// when we're connecting to https servers that might not be there
// or might be using self-signed certificates.
var impatientInsecureHTTPClient = &http.Client{
	Timeout: time.Duration(5 * time.Second),
	Transport: &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	},
}

type httpError struct {
	status     string
	statusCode int
	url        string
}

func (e *httpError) Error() string {
	return fmt.Sprintf("%s: %s", e.url, e.status)
}

// httpGET returns the data from an HTTP GET request for the given URL.
func httpGET(url string) ([]byte, error) {
	resp, err := httpClient.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		err := &httpError{status: resp.Status, statusCode: resp.StatusCode, url: url}

		return nil, err
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("%s: %v", url, err)
	}
	return b, nil
}

// httpsOrHTTP returns the body of either the importPath's
// https resource or, if unavailable, the http resource.
func httpsOrHTTP(importPath string, security securityMode) (urlStr string, body io.ReadCloser, err error) {
	fetch := func(scheme string) (urlStr string, res *http.Response, err error) {
		u, err := url.Parse(scheme + "://" + importPath)
		if err != nil {
			return "", nil, err
		}
		u.RawQuery = "go-get=1"
		urlStr = u.String()
		if buildV {
			log.Printf("Fetching %s", urlStr)
		}
		if security == insecure && scheme == "https" { // fail earlier
			res, err = impatientInsecureHTTPClient.Get(urlStr)
		} else {
			res, err = httpClient.Get(urlStr)
		}
		return
	}
	closeBody := func(res *http.Response) {
		if res != nil {
			res.Body.Close()
		}
	}
	urlStr, res, err := fetch("https")
	if err != nil {
		if buildV {
			log.Printf("https fetch failed: %v", err)
		}
		if security == insecure {
			closeBody(res)
			urlStr, res, err = fetch("http")
		}
	}
	if err != nil {
		closeBody(res)
		return "", nil, err
	}
	// Note: accepting a non-200 OK here, so people can serve a
	// meta import in their http 404 page.
	if buildV {
		log.Printf("Parsing meta tags from %s (status code %d)", urlStr, res.StatusCode)
	}
	return urlStr, res.Body, nil
}

// From go/src/cmd/go/discovery.go

// charsetReader returns a reader for the given charset. Currently
// it only supports UTF-8 and ASCII. Otherwise, it returns a meaningful
// error which is printed by go get, so the user can find why the package
// wasn't downloaded if the encoding is not supported. Note that, in
// order to reduce potential errors, ASCII is treated as UTF-8 (i.e. characters
// greater than 0x7f are not rejected).
func charsetReader(charset string, input io.Reader) (io.Reader, error) {
	switch strings.ToLower(charset) {
	case "ascii":
		return input, nil
	default:
		return nil, fmt.Errorf("can't decode XML document using charset %q", charset)
	}
}

// parseMetaGoImports returns meta imports from the HTML in r.
// Parsing ends at the end of the <head> section or the beginning of the <body>.
func parseMetaGoImports(r io.Reader) (imports []metaImport, err error) {
	d := xml.NewDecoder(r)
	d.CharsetReader = charsetReader
	d.Strict = false
	var t xml.Token
	for {
		t, err = d.RawToken()
		if err != nil {
			if err == io.EOF || len(imports) > 0 {
				err = nil
			}
			return
		}
		if e, ok := t.(xml.StartElement); ok && strings.EqualFold(e.Name.Local, "body") {
			return
		}
		if e, ok := t.(xml.EndElement); ok && strings.EqualFold(e.Name.Local, "head") {
			return
		}
		e, ok := t.(xml.StartElement)
		if !ok || !strings.EqualFold(e.Name.Local, "meta") {
			continue
		}
		if attrValue(e.Attr, "name") != "go-import" {
			continue
		}
		if f := strings.Fields(attrValue(e.Attr, "content")); len(f) == 3 {
			imports = append(imports, metaImport{
				Prefix:   f[0],
				VCS:      f[1],
				RepoRoot: f[2],
			})
		}
	}
}

// attrValue returns the attribute value for the case-insensitive key
// `name', or the empty string if nothing is found.
func attrValue(attrs []xml.Attr, name string) string {
	for _, a := range attrs {
		if strings.EqualFold(a.Name.Local, name) {
			return a.Value
		}
	}
	return ""
}
<commit_msg>attempt to guess host if no go meta-tag is exported<commit_after>// This file contains code taken from the "go" source tree, governed
// by the go license:

// Copyright (c) 2012 The Go Authors. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//    * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//    * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//    * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package main

import (
	"crypto/tls"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

const buildV = false

// From go/src/cmd/go/vcs.go

// securityMode specifies whether a function should make network
// calls using insecure transports (eg, plain text HTTP).
// The zero value is "secure".
type securityMode int

const (
	secure securityMode = iota
	insecure
)

// repoRoot represents a version control system, a repo, and a root of
// where to put it on disk.
type repoRoot struct {
	vcs string

	// repo is the repository URL, including scheme
	repo string

	// root is the import path corresponding to the root of the
	// repository
	root string
}

// repoRootForImportDynamic finds a *repoRoot for a custom domain that's not
// statically known by repoRootForImportPathStatic.
//
// This handles custom import paths like "name.tld/pkg/foo" or just "name.tld".
func queryRepoRoot(importPath string, security securityMode) (*repoRoot, error) {
	slash := strings.Index(importPath, "/")
	if slash < 0 {
		slash = len(importPath)
	}
	host := importPath[:slash]
	if !strings.Contains(host, ".") {
		return nil, errors.New("import path does not begin with hostname")
	}
	urlStr, body, err := httpsOrHTTP(importPath, security)
	if err != nil {
		msg := "https fetch: %v"
		if security == insecure {
			msg = "http/" + msg
		}
		return nil, fmt.Errorf(msg, err)
	}
	defer body.Close()
	imports, err := parseMetaGoImports(body)
	if err != nil {
		return nil, fmt.Errorf("parsing %s: %v", importPath, err)
	}
	// Find the matched meta import.
	mmi, err := matchGoImport(imports, importPath)
	if err != nil {
		if err != errNoMatch {
			return nil, fmt.Errorf("parse %s: %v", urlStr, err)
		}

		bestGuess := strings.Join(strings.Split(importPath, "/")[:3], "/")
		return &repoRoot{
			vcs:  "git",
			repo: fmt.Sprintf("https://%s.git", bestGuess),
			root: bestGuess,
		}, nil
	}
	if buildV {
		log.Printf("get %q: found meta tag %#v at %s", importPath, mmi, urlStr)
	}
	// If the import was "uni.edu/bob/project", which said the
	// prefix was "uni.edu" and the RepoRoot was "evilroot.com",
	// make sure we don't trust Bob and check out evilroot.com to
	// "uni.edu" yet (possibly overwriting/preempting another
	// non-evil student).  Instead, first verify the root and see
	// if it matches Bob's claim.
	if mmi.Prefix != importPath {
		if buildV {
			log.Printf("get %q: verifying non-authoritative meta tag", importPath)
		}
		urlStr0 := urlStr
		var imports []metaImport
		urlStr, imports, err = metaImportsForPrefix(mmi.Prefix, security)
		if err != nil {
			return nil, err
		}
		metaImport2, err := matchGoImport(imports, importPath)
		if err != nil || mmi != metaImport2 {
			return nil, fmt.Errorf("%s and %s disagree about go-import for %s", urlStr0, urlStr, mmi.Prefix)
		}
	}

	if !strings.Contains(mmi.RepoRoot, "://") {
		return nil, fmt.Errorf("%s: invalid repo root %q; no scheme", urlStr, mmi.RepoRoot)
	}
	return &repoRoot{
		vcs:  mmi.VCS,
		repo: mmi.RepoRoot,
		root: mmi.Prefix,
	}, nil
}

var (
	fetchCacheMu sync.Mutex
	fetchCache   = map[string]fetchResult{} // key is metaImportsForPrefix's importPrefix
)

// metaImportsForPrefix takes a package's root import path as declared in a <meta> tag
// and returns its HTML discovery URL and the parsed metaImport lines
// found on the page.
//
// The importPath is of the form "golang.org/x/tools".
// It is an error if no imports are found.
// urlStr will still be valid if err != nil.
// The returned urlStr will be of the form "https://golang.org/x/tools?go-get=1"
func metaImportsForPrefix(importPrefix string, security securityMode) (urlStr string, imports []metaImport, err error) {
	setCache := func(res fetchResult) fetchResult {
		fetchCacheMu.Lock()
		defer fetchCacheMu.Unlock()
		fetchCache[importPrefix] = res
		return res
	}

	fetch := func() fetchResult {
		fetchCacheMu.Lock()
		if res, ok := fetchCache[importPrefix]; ok {
			fetchCacheMu.Unlock()
			return res
		}
		fetchCacheMu.Unlock()

		urlStr, body, err := httpsOrHTTP(importPrefix, security)
		if err != nil {
			return setCache(fetchResult{urlStr: urlStr, err: fmt.Errorf("fetch %s: %v", urlStr, err)})
		}
		imports, err := parseMetaGoImports(body)
		if err != nil {
			return setCache(fetchResult{urlStr: urlStr, err: fmt.Errorf("parsing %s: %v", urlStr, err)})
		}
		if len(imports) == 0 {
			err = fmt.Errorf("fetch %s: no go-import meta tag", urlStr)
		}
		return setCache(fetchResult{urlStr: urlStr, imports: imports, err: err})
	}

	res := fetch()
	return res.urlStr, res.imports, res.err
}

type fetchResult struct {
	urlStr  string // e.g. "https://foo.com/x/bar?go-get=1"
	imports []metaImport
	err     error
}

// metaImport represents the parsed <meta name="go-import"
// content="prefix vcs reporoot" /> tags from HTML files.
type metaImport struct {
	Prefix, VCS, RepoRoot string
}

// errNoMatch is returned from matchGoImport when there's no applicable match.
var errNoMatch = errors.New("no import match")

// matchGoImport returns the metaImport from imports matching importPath.
// An error is returned if there are multiple matches.
// errNoMatch is returned if none match.
func matchGoImport(imports []metaImport, importPath string) (_ metaImport, err error) {
	match := -1
	for i, im := range imports {
		if !strings.HasPrefix(importPath, im.Prefix) {
			continue
		}
		if match != -1 {
			err = fmt.Errorf("multiple meta tags match import path %q", importPath)
			return
		}
		match = i
	}
	if match == -1 {
		err = errNoMatch
		return
	}
	return imports[match], nil
}

// From go/src/cmd/go/http.go

// httpClient is the default HTTP client, but a variable so it can be
// changed by tests, without modifying http.DefaultClient.
var httpClient = http.DefaultClient

// impatientInsecureHTTPClient is used in -insecure mode,
// when we're connecting to https servers that might not be there
// or might be using self-signed certificates.
var impatientInsecureHTTPClient = &http.Client{
	Timeout: time.Duration(5 * time.Second),
	Transport: &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	},
}

type httpError struct {
	status     string
	statusCode int
	url        string
}

func (e *httpError) Error() string {
	return fmt.Sprintf("%s: %s", e.url, e.status)
}

// httpGET returns the data from an HTTP GET request for the given URL.
func httpGET(url string) ([]byte, error) {
	resp, err := httpClient.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		err := &httpError{status: resp.Status, statusCode: resp.StatusCode, url: url}

		return nil, err
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("%s: %v", url, err)
	}
	return b, nil
}

// httpsOrHTTP returns the body of either the importPath's
// https resource or, if unavailable, the http resource.
func httpsOrHTTP(importPath string, security securityMode) (urlStr string, body io.ReadCloser, err error) {
	fetch := func(scheme string) (urlStr string, res *http.Response, err error) {
		u, err := url.Parse(scheme + "://" + importPath)
		if err != nil {
			return "", nil, err
		}
		u.RawQuery = "go-get=1"
		urlStr = u.String()
		if buildV {
			log.Printf("Fetching %s", urlStr)
		}
		if security == insecure && scheme == "https" { // fail earlier
			res, err = impatientInsecureHTTPClient.Get(urlStr)
		} else {
			res, err = httpClient.Get(urlStr)
		}
		return
	}
	closeBody := func(res *http.Response) {
		if res != nil {
			res.Body.Close()
		}
	}
	urlStr, res, err := fetch("https")
	if err != nil {
		if buildV {
			log.Printf("https fetch failed: %v", err)
		}
		if security == insecure {
			closeBody(res)
			urlStr, res, err = fetch("http")
		}
	}
	if err != nil {
		closeBody(res)
		return "", nil, err
	}
	// Note: accepting a non-200 OK here, so people can serve a
	// meta import in their http 404 page.
	if buildV {
		log.Printf("Parsing meta tags from %s (status code %d)", urlStr, res.StatusCode)
	}
	return urlStr, res.Body, nil
}

// From go/src/cmd/go/discovery.go

// charsetReader returns a reader for the given charset. Currently
// it only supports UTF-8 and ASCII. Otherwise, it returns a meaningful
// error which is printed by go get, so the user can find why the package
// wasn't downloaded if the encoding is not supported. Note that, in
// order to reduce potential errors, ASCII is treated as UTF-8 (i.e. characters
// greater than 0x7f are not rejected).
func charsetReader(charset string, input io.Reader) (io.Reader, error) {
	switch strings.ToLower(charset) {
	case "ascii":
		return input, nil
	default:
		return nil, fmt.Errorf("can't decode XML document using charset %q", charset)
	}
}

// parseMetaGoImports returns meta imports from the HTML in r.
// Parsing ends at the end of the <head> section or the beginning of the <body>.
func parseMetaGoImports(r io.Reader) (imports []metaImport, err error) {
	d := xml.NewDecoder(r)
	d.CharsetReader = charsetReader
	d.Strict = false
	var t xml.Token
	for {
		t, err = d.RawToken()
		if err != nil {
			if err == io.EOF || len(imports) > 0 {
				err = nil
			}
			return
		}
		if e, ok := t.(xml.StartElement); ok && strings.EqualFold(e.Name.Local, "body") {
			return
		}
		if e, ok := t.(xml.EndElement); ok && strings.EqualFold(e.Name.Local, "head") {
			return
		}
		e, ok := t.(xml.StartElement)
		if !ok || !strings.EqualFold(e.Name.Local, "meta") {
			continue
		}
		if attrValue(e.Attr, "name") != "go-import" {
			continue
		}
		if f := strings.Fields(attrValue(e.Attr, "content")); len(f) == 3 {
			imports = append(imports, metaImport{
				Prefix:   f[0],
				VCS:      f[1],
				RepoRoot: f[2],
			})
		}
	}
}

// attrValue returns the attribute value for the case-insensitive key
// `name', or the empty string if nothing is found.
func attrValue(attrs []xml.Attr, name string) string {
	for _, a := range attrs {
		if strings.EqualFold(a.Name.Local, name) {
			return a.Value
		}
	}
	return ""
}
<|endoftext|>