<commit_before>package routing

import (
	"math"

	"container/heap"

	"github.com/lightningnetwork/lnd/channeldb"
	"github.com/roasbeef/btcd/btcec"
	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcutil"
)

const (
	// HopLimit is the maximum number hops that is permissible as a route.
	// Any potential paths found that lie above this limit will be rejected
	// with an error. This value is computed using the current fixed-size
	// packet length of the Sphinx construction.
	HopLimit = 20

	// infinity is used as a starting distance in our shortest path search.
	infinity = math.MaxFloat64
)

// Route represents a path through the channel graph which runs over one or
// more channels in succession. This struct carries all the information
// required to craft the Sphinx onion packet, and send the payment along the
// first hop in the path. A route is only selected as valid if all the channels
// have sufficient capacity to carry the initial payment amount after fees are
// accounted for.
type Route struct {
	// TotalTimeLock is the cumulative (final) time lock across the entire
	// route. This is the CLTV value that should be extended to the first
	// hop in the route. All other hops will decrement the time-lock as
	// advertised, leaving enough time for all hops to wait for or present
	// the payment preimage to complete the payment.
	TotalTimeLock uint32

	// TotalFees is the sum of the fees paid at each hop within the final
	// route. In the case of a one-hop payment, this value will be zero as
	// we don't need to pay a fee it ourself.
	TotalFees btcutil.Amount

	// TotalAmount is the total amount of funds required to complete a
	// payment over this route. This value includes the cumulative fees at
	// each hop. As a result, the HTLC extended to the first-hop in the
	// route will need to have at least this many satoshis, otherwise the
	// route will fail at an intermediate node due to an insufficient
	// amount of fees.
	TotalAmount btcutil.Amount

	// Hops contains details concerning the specific forwarding details at
	// each hop.
	Hops []*Hop
}

// Hop represents the forwarding details at a particular position within the
// final route. This struct houses the values necessary to create the HTLC
// which will travel along this hop, and also encode the per-hop payload
// included within the Sphinx packet.
type Hop struct {
	// Channel is the active payment channel edge that this hop will travel
	// along.
	Channel *ChannelHop

	// TimeLockDelta is the delta that this hop will subtract from the HTLC
	// before extending it to the next hop in the route.
	TimeLockDelta uint16

	// AmtToForward is the amount that this hop will forward to the next
	// hop. This value is less than the value that the incoming HTLC
	// carries as a fee will be subtracted by the hop.
	AmtToForward btcutil.Amount

	// Fee is the total fee that this hop will subtract from the incoming
	// payment, this difference nets the hop fees for forwarding the
	// payment.
	Fee btcutil.Amount
}

// ChannelHop is an intermediate hop within the network with a greater
// multi-hop payment route. This struct contains the relevant routing policy of
// the particular edge, as well as the total capacity, and origin chain of the
// channel itself.
type ChannelHop struct {
	// Capacity is the total capacity of the channel being traversed. This
	// value is expressed for stability in satoshis.
	Capacity btcutil.Amount

	// Chain is a 32-byte has that denotes the base blockchain network of
	// the channel. The 32-byte hash is the "genesis" block of the
	// blockchain, or the very first block in the chain.
	//
	// TODO(roasbeef): store chain within edge info/policy in database.
	Chain chainhash.Hash

	*channeldb.ChannelEdgePolicy
}

// computeFee computes the fee to forward an HTLC of `amt` satoshis over the
// passed active payment channel. This value is currently computed as specified
// in BOLT07, but will likely change in the near future.
func computeFee(amt btcutil.Amount, edge *ChannelHop) btcutil.Amount {
	return edge.FeeBaseMSat + (amt*edge.FeeProportionalMillionths)/1000000
}

// newRoute returns a fully valid route between the source and target that's
// capable of supporting a payment of `amtToSend` after fees are fully
// computed. If the route is too long, or the selected path cannot support the
// fully payment including fees, then a non-nil error is returned.
//
// NOTE: The passed slice of ChannelHops MUST be sorted in reverse order: from
// the target to the source node of the path finding aattempt.
func newRoute(amtToSend btcutil.Amount, pathEdges []*ChannelHop) (*Route, error) {
	route := &Route{
		Hops: make([]*Hop, len(pathEdges)),
	}

	// The running amount is the total amount of satoshis required at this
	// point in the route. We start this value at the amount we want to
	// send to the destination. This value will then get successively
	// larger as we compute the fees going backwards.
	runningAmt := amtToSend
	pathLength := len(pathEdges)
	for i, edge := range pathEdges {
		// Now we create the hop struct for this point in the route.
		// The amount to forward is the running amount, and we compute
		// the required fee based on this amount.
		nextHop := &Hop{
			Channel:       edge,
			AmtToForward:  runningAmt,
			Fee:           computeFee(runningAmt, edge),
			TimeLockDelta: edge.TimeLockDelta,
		}
		edge.Node.PubKey.Curve = nil

		// As a sanity check, we ensure that the selected channel has
		// enough capacity to forward the required amount which
		// includes the fee dictated at each hop.
		if nextHop.AmtToForward > nextHop.Channel.Capacity {
			return nil, ErrInsufficientCapacity
		}

		// We don't pay any fees to ourselves on the first-hop channel,
		// so we don't tally up the running fee and amount.
		if i != len(pathEdges)-1 {
			// For a node to forward an HTLC, then following
			// inequality most hold true: amt_in - fee >=
			// amt_to_forward. Therefore we add the fee this node
			// consumes in order to calculate the amount that it
			// show be forwarded by the prior node which is the
			// next hop in our loop.
			runningAmt += nextHop.Fee

			// Next we tally the total fees (thus far) in the
			// route, and also accumulate the total timelock in the
			// route by adding the node's time lock delta which is
			// the amount of blocks it'll subtract from the
			// incoming time lock.
			route.TotalFees += nextHop.Fee
		} else {
			nextHop.Fee = 0
		}

		route.TotalTimeLock += uint32(nextHop.TimeLockDelta)

		// Finally, as we're currently talking the route backwards, we
		// reverse the index in order to place this hop at the proper
		// spot in the forward direction of the route.
		route.Hops[pathLength-1-i] = nextHop
	}

	// The total amount required for this route will be the value the
	// source extends to the first hop in the route.
	route.TotalAmount = runningAmt

	return route, nil
}

// vertex is a simple alias for the serialization of a compressed Bitcoin
// public key.
type vertex [33]byte

// newVertex returns a new vertex given a public key.
func newVertex(pub *btcec.PublicKey) vertex {
	var v vertex
	copy(v[:], pub.SerializeCompressed())
	return v
}

// edgeWithPrev is a helper struct used in path finding that couples an
// directional edge with the node's ID in the opposite direction.
type edgeWithPrev struct {
	edge     *ChannelHop
	prevNode *btcec.PublicKey
}

// edgeWeight computes the weight of an edge. This value is used when searching
// for the shortest path within the channel graph between two nodes. Currently
// this is just 1 + the cltv delta value required at this hop, this value
// should be tuned with experimental and empirical data.
//
// TODO(roasbeef): compute robust weight metric
func edgeWeight(e *channeldb.ChannelEdgePolicy) float64 {
	return float64(1 + e.TimeLockDelta)
}

// findRoute attempts to find a path from the source node within the
// ChannelGraph to the target node that's capable of supporting a payment of
// `amt` value. The current approach implemented is modified version of
// Dijkstra's algorithm to find a single shortest path between the source node
// and the destination. The distance metric used for edges is related to the
// time-lock+fee costs along a particular edge. If a path is found, this
// function returns a slice of ChannelHop structs which encoded the chosen path
// (backwards) from the target to the source.
func findRoute(graph *channeldb.ChannelGraph, sourceNode *channeldb.LightningNode,
	target *btcec.PublicKey, amt btcutil.Amount) ([]*ChannelHop, error) {

	// First we'll initialize an empty heap which'll help us to quickly
	// locate the next edge we should visit next during our graph
	// traversal.
	var nodeHeap distanceHeap

	// For each node/vertex the graph we create an entry in the distance
	// map for the node set with a distance of "infinity".  We also mark
	// add the node to our set of unvisited nodes.
	distance := make(map[vertex]nodeWithDist)
	if err := graph.ForEachNode(func(node *channeldb.LightningNode) error {
		// TODO(roasbeef): with larger graph can just use disk seeks
		// with a visited map
		distance[newVertex(node.PubKey)] = nodeWithDist{
			dist: infinity,
			node: node,
		}
		return nil
	}); err != nil {
		return nil, err
	}

	// To start, we add the source of our path finding attempt to the
	// distance map with with a distance of 0. This indicates our starting
	// point in the graph traversal.
	sourceVertex := newVertex(sourceNode.PubKey)
	distance[sourceVertex] = nodeWithDist{
		dist: 0,
		node: sourceNode,
	}

	// To start, our source node will the sole item within our distance
	// heap.
	heap.Push(&nodeHeap, distance[sourceVertex])

	// We'll use this map as a series of "previous" hop pointers. So to get
	// to `vertex` we'll take the edge that it's mapped to within `prev`.
	prev := make(map[vertex]edgeWithPrev)

	for nodeHeap.Len() != 0 {
		// Fetch the node within the smallest distance from our source
		// from the heap.
		bestNode := heap.Pop(&nodeHeap).(nodeWithDist).node

		// If we've reached our target (or we don't have any outgoing
		// edges), then we're done here and can exit the graph
		// traversal early.
		if bestNode.PubKey.IsEqual(target) {
			break
		}

		// Now that we've found the next potential step to take we'll
		// examine all the outgoing edge (channels) from this node to
		// further our graph traversal.
		pivot := newVertex(bestNode.PubKey)
		err := bestNode.ForEachChannel(nil, func(edgeInfo *channeldb.ChannelEdgeInfo,
			edge *channeldb.ChannelEdgePolicy) error {

			// Compute the tentative distance to this new
			// channel/edge which is the distance to our current
			// pivot node plus the weight of this edge.
			tempDist := distance[pivot].dist + edgeWeight(edge)

			// If this new tentative distance is better than the
			// current best known distance to this node, then we
			// record the new better distance, and also populate
			// our "next hop" map with this edge. We'll also shave
			// off irrelevant edges by adding the sufficient
			// capacity of an edge to our relaxation condition.
			v := newVertex(edge.Node.PubKey)
			if tempDist < distance[v].dist &&
				edgeInfo.Capacity >= amt {

				distance[v] = nodeWithDist{
					dist: tempDist,
					node: edge.Node,
				}
				prev[v] = edgeWithPrev{
					edge: &ChannelHop{
						ChannelEdgePolicy: edge,
						Capacity:          edgeInfo.Capacity,
					},
					prevNode: bestNode.PubKey,
				}

				// Add this new node to our heap as we'd like
				// to further explore down this edge.
				heap.Push(&nodeHeap, distance[v])
			}
			return nil
		})
		if err != nil {
			return nil, err
		}
	}

	// If the target node isn't found in the prev hop map, then a path
	// doesn't exist, so we terminate in an error.
	if _, ok := prev[newVertex(target)]; !ok {
		return nil, ErrNoPathFound
	}

	// If the potential route if below the max hop limit, then we'll use
	// the prevHop map to unravel the path. We end up with a list of edges
	// in the reverse direction which we'll use to properly calculate the
	// timelock and fee values.
	pathEdges := make([]*ChannelHop, 0, len(prev))
	prevNode := newVertex(target)
	for prevNode != sourceVertex { // TODO(roasbeef): assumes no cycles
		// Add the current hop to the limit of path edges then walk
		// backwards from this hop via the prev pointer for this hop
		// within the prevHop map.
		pathEdges = append(pathEdges, prev[prevNode].edge)
		prevNode = newVertex(prev[prevNode].prevNode)
	}

	// The route is invalid if it spans more than 20 hops. The current
	// Sphinx (onion routing) implementation can only encode up to 20 hops
	// as the entire packet is fixed size. If this route is more than 20
	// hops, then it's invalid.
	if len(pathEdges) > HopLimit {
		return nil, ErrMaxHopsExceeded
	}

	return pathEdges, nil
}
}
<commit_msg>routing: newRoute now expects the path in forwards order<commit_after>package routing

import (
	"math"

	"container/heap"

	"github.com/lightningnetwork/lnd/channeldb"
	"github.com/roasbeef/btcd/btcec"
	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcutil"
)

const (
	// HopLimit is the maximum number hops that is permissible as a route.
	// Any potential paths found that lie above this limit will be rejected
	// with an error. This value is computed using the current fixed-size
	// packet length of the Sphinx construction.
	HopLimit = 20

	// infinity is used as a starting distance in our shortest path search.
	infinity = math.MaxFloat64
)

// Route represents a path through the channel graph which runs over one or
// more channels in succession. This struct carries all the information
// required to craft the Sphinx onion packet, and send the payment along the
// first hop in the path. A route is only selected as valid if all the channels
// have sufficient capacity to carry the initial payment amount after fees are
// accounted for.
type Route struct {
	// TotalTimeLock is the cumulative (final) time lock across the entire
	// route. This is the CLTV value that should be extended to the first
	// hop in the route. All other hops will decrement the time-lock as
	// advertised, leaving enough time for all hops to wait for or present
	// the payment preimage to complete the payment.
	TotalTimeLock uint32

	// TotalFees is the sum of the fees paid at each hop within the final
	// route. In the case of a one-hop payment, this value will be zero as
	// we don't need to pay a fee it ourself.
	TotalFees btcutil.Amount

	// TotalAmount is the total amount of funds required to complete a
	// payment over this route. This value includes the cumulative fees at
	// each hop. As a result, the HTLC extended to the first-hop in the
	// route will need to have at least this many satoshis, otherwise the
	// route will fail at an intermediate node due to an insufficient
	// amount of fees.
	TotalAmount btcutil.Amount

	// Hops contains details concerning the specific forwarding details at
	// each hop.
	Hops []*Hop
}

// Hop represents the forwarding details at a particular position within the
// final route. This struct houses the values necessary to create the HTLC
// which will travel along this hop, and also encode the per-hop payload
// included within the Sphinx packet.
type Hop struct {
	// Channel is the active payment channel edge that this hop will travel
	// along.
	Channel *ChannelHop

	// TimeLockDelta is the delta that this hop will subtract from the HTLC
	// before extending it to the next hop in the route.
	TimeLockDelta uint16

	// AmtToForward is the amount that this hop will forward to the next
	// hop. This value is less than the value that the incoming HTLC
	// carries as a fee will be subtracted by the hop.
	AmtToForward btcutil.Amount

	// Fee is the total fee that this hop will subtract from the incoming
	// payment, this difference nets the hop fees for forwarding the
	// payment.
	Fee btcutil.Amount
}

// ChannelHop is an intermediate hop within the network with a greater
// multi-hop payment route. This struct contains the relevant routing policy of
// the particular edge, as well as the total capacity, and origin chain of the
// channel itself.
type ChannelHop struct {
	// Capacity is the total capacity of the channel being traversed. This
	// value is expressed for stability in satoshis.
	Capacity btcutil.Amount

	// Chain is a 32-byte has that denotes the base blockchain network of
	// the channel. The 32-byte hash is the "genesis" block of the
	// blockchain, or the very first block in the chain.
	//
	// TODO(roasbeef): store chain within edge info/policy in database.
	Chain chainhash.Hash

	*channeldb.ChannelEdgePolicy
}

// computeFee computes the fee to forward an HTLC of `amt` satoshis over the
// passed active payment channel. This value is currently computed as specified
// in BOLT07, but will likely change in the near future.
func computeFee(amt btcutil.Amount, edge *ChannelHop) btcutil.Amount {
	return edge.FeeBaseMSat + (amt*edge.FeeProportionalMillionths)/1000000
}

// newRoute returns a fully valid route between the source and target that's
// capable of supporting a payment of `amtToSend` after fees are fully
// computed. If the route is too long, or the selected path cannot support the
// fully payment including fees, then a non-nil error is returned.
//
// NOTE: The passed slice of ChannelHops MUST be sorted in forward order: from
// the source to the target node of the path finding attempt.
func newRoute(amtToSend btcutil.Amount, pathEdges []*ChannelHop) (*Route, error) {
	route := &Route{
		Hops: make([]*Hop, len(pathEdges)),
	}

	// The running amount is the total amount of satoshis required at this
	// point in the route. We start this value at the amount we want to
	// send to the destination. This value will then get successively
	// larger as we compute the fees going backwards.
	runningAmt := amtToSend
	pathLength := len(pathEdges)
	for i := pathLength - 1; i >= 0; i-- {
		edge := pathEdges[i]

		// Now we create the hop struct for this point in the route.
		// The amount to forward is the running amount, and we compute
		// the required fee based on this amount.
		nextHop := &Hop{
			Channel:       edge,
			AmtToForward:  runningAmt,
			Fee:           computeFee(runningAmt, edge),
			TimeLockDelta: edge.TimeLockDelta,
		}
		edge.Node.PubKey.Curve = nil

		// As a sanity check, we ensure that the selected channel has
		// enough capacity to forward the required amount which
		// includes the fee dictated at each hop.
		if nextHop.AmtToForward > nextHop.Channel.Capacity {
			return nil, ErrInsufficientCapacity
		}

		// We don't pay any fees to ourselves on the first-hop channel,
		// so we don't tally up the running fee and amount.
		if i != len(pathEdges)-1 {
			// For a node to forward an HTLC, then following
			// inequality most hold true: amt_in - fee >=
			// amt_to_forward. Therefore we add the fee this node
			// consumes in order to calculate the amount that it
			// show be forwarded by the prior node which is the
			// next hop in our loop.
			runningAmt += nextHop.Fee

			// Next we tally the total fees (thus far) in the
			// route, and also accumulate the total timelock in the
			// route by adding the node's time lock delta which is
			// the amount of blocks it'll subtract from the
			// incoming time lock.
			route.TotalFees += nextHop.Fee
		} else {
			nextHop.Fee = 0
		}

		route.TotalTimeLock += uint32(nextHop.TimeLockDelta)

		route.Hops[i] = nextHop
	}

	// The total amount required for this route will be the value the
	// source extends to the first hop in the route.
	route.TotalAmount = runningAmt

	return route, nil
}

// vertex is a simple alias for the serialization of a compressed Bitcoin
// public key.
type vertex [33]byte

// newVertex returns a new vertex given a public key.
func newVertex(pub *btcec.PublicKey) vertex {
	var v vertex
	copy(v[:], pub.SerializeCompressed())
	return v
}

// edgeWithPrev is a helper struct used in path finding that couples an
// directional edge with the node's ID in the opposite direction.
type edgeWithPrev struct {
	edge     *ChannelHop
	prevNode *btcec.PublicKey
}

// edgeWeight computes the weight of an edge. This value is used when searching
// for the shortest path within the channel graph between two nodes. Currently
// this is just 1 + the cltv delta value required at this hop, this value
// should be tuned with experimental and empirical data.
//
// TODO(roasbeef): compute robust weight metric
func edgeWeight(e *channeldb.ChannelEdgePolicy) float64 {
	return float64(1 + e.TimeLockDelta)
}

// findRoute attempts to find a path from the source node within the
// ChannelGraph to the target node that's capable of supporting a payment of
// `amt` value. The current approach implemented is modified version of
// Dijkstra's algorithm to find a single shortest path between the source node
// and the destination. The distance metric used for edges is related to the
// time-lock+fee costs along a particular edge. If a path is found, this
// function returns a slice of ChannelHop structs which encoded the chosen path
// (backwards) from the target to the source.
func findRoute(graph *channeldb.ChannelGraph, sourceNode *channeldb.LightningNode,
	target *btcec.PublicKey, amt btcutil.Amount) ([]*ChannelHop, error) {

	// First we'll initialize an empty heap which'll help us to quickly
	// locate the next edge we should visit next during our graph
	// traversal.
	var nodeHeap distanceHeap

	// For each node/vertex the graph we create an entry in the distance
	// map for the node set with a distance of "infinity".  We also mark
	// add the node to our set of unvisited nodes.
	distance := make(map[vertex]nodeWithDist)
	if err := graph.ForEachNode(func(node *channeldb.LightningNode) error {
		// TODO(roasbeef): with larger graph can just use disk seeks
		// with a visited map
		distance[newVertex(node.PubKey)] = nodeWithDist{
			dist: infinity,
			node: node,
		}
		return nil
	}); err != nil {
		return nil, err
	}

	// To start, we add the source of our path finding attempt to the
	// distance map with with a distance of 0. This indicates our starting
	// point in the graph traversal.
	sourceVertex := newVertex(sourceNode.PubKey)
	distance[sourceVertex] = nodeWithDist{
		dist: 0,
		node: sourceNode,
	}

	// To start, our source node will the sole item within our distance
	// heap.
	heap.Push(&nodeHeap, distance[sourceVertex])

	// We'll use this map as a series of "previous" hop pointers. So to get
	// to `vertex` we'll take the edge that it's mapped to within `prev`.
	prev := make(map[vertex]edgeWithPrev)

	for nodeHeap.Len() != 0 {
		// Fetch the node within the smallest distance from our source
		// from the heap.
		bestNode := heap.Pop(&nodeHeap).(nodeWithDist).node

		// If we've reached our target (or we don't have any outgoing
		// edges), then we're done here and can exit the graph
		// traversal early.
		if bestNode.PubKey.IsEqual(target) {
			break
		}

		// Now that we've found the next potential step to take we'll
		// examine all the outgoing edge (channels) from this node to
		// further our graph traversal.
		pivot := newVertex(bestNode.PubKey)
		err := bestNode.ForEachChannel(nil, func(edgeInfo *channeldb.ChannelEdgeInfo,
			edge *channeldb.ChannelEdgePolicy) error {

			// Compute the tentative distance to this new
			// channel/edge which is the distance to our current
			// pivot node plus the weight of this edge.
			tempDist := distance[pivot].dist + edgeWeight(edge)

			// If this new tentative distance is better than the
			// current best known distance to this node, then we
			// record the new better distance, and also populate
			// our "next hop" map with this edge. We'll also shave
			// off irrelevant edges by adding the sufficient
			// capacity of an edge to our relaxation condition.
			v := newVertex(edge.Node.PubKey)
			if tempDist < distance[v].dist &&
				edgeInfo.Capacity >= amt {

				distance[v] = nodeWithDist{
					dist: tempDist,
					node: edge.Node,
				}
				prev[v] = edgeWithPrev{
					edge: &ChannelHop{
						ChannelEdgePolicy: edge,
						Capacity:          edgeInfo.Capacity,
					},
					prevNode: bestNode.PubKey,
				}

				// Add this new node to our heap as we'd like
				// to further explore down this edge.
				heap.Push(&nodeHeap, distance[v])
			}
			return nil
		})
		if err != nil {
			return nil, err
		}
	}

	// If the target node isn't found in the prev hop map, then a path
	// doesn't exist, so we terminate in an error.
	if _, ok := prev[newVertex(target)]; !ok {
		return nil, ErrNoPathFound
	}

	// If the potential route if below the max hop limit, then we'll use
	// the prevHop map to unravel the path. We end up with a list of edges
	// in the reverse direction which we'll use to properly calculate the
	// timelock and fee values.
	pathEdges := make([]*ChannelHop, 0, len(prev))
	prevNode := newVertex(target)
	for prevNode != sourceVertex { // TODO(roasbeef): assumes no cycles
		// Add the current hop to the limit of path edges then walk
		// backwards from this hop via the prev pointer for this hop
		// within the prevHop map.
		pathEdges = append(pathEdges, prev[prevNode].edge)
		prevNode = newVertex(prev[prevNode].prevNode)
	}

	// The route is invalid if it spans more than 20 hops. The current
	// Sphinx (onion routing) implementation can only encode up to 20 hops
	// as the entire packet is fixed size. If this route is more than 20
	// hops, then it's invalid.
	if len(pathEdges) > HopLimit {
		return nil, ErrMaxHopsExceeded
	}

	return pathEdges, nil
}
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.
package libkbfs

import (
	"crypto/rand"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/keybase/kbfs/kbfscodec"
	"github.com/keybase/kbfs/kbfscrypto"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

// blockReturner contains a block value to copy into requested blocks, and a
// channel to synchronize on with the worker.
type blockReturner struct {
	block      Block
	continueCh chan error
	startCh    chan struct{}
}

// fakeBlockGetter allows specifying and obtaining fake blocks.
type fakeBlockGetter struct {
	mtx           sync.RWMutex
	blockMap      map[BlockPointer]blockReturner
	codec         kbfscodec.Codec
	respectCancel bool
}

// newFakeBlockGetter returns a fakeBlockGetter.
func newFakeBlockGetter(respectCancel bool) *fakeBlockGetter {
	return &fakeBlockGetter{
		blockMap:      make(map[BlockPointer]blockReturner),
		codec:         kbfscodec.NewMsgpack(),
		respectCancel: respectCancel,
	}
}

// setBlockToReturn sets the block that will be returned for a given
// BlockPointer. Returns a writeable channel that getBlock will wait on, to
// allow synchronization of tests.
func (bg *fakeBlockGetter) setBlockToReturn(blockPtr BlockPointer,
	block Block) (startCh <-chan struct{}, continueCh chan<- error) {
	bg.mtx.Lock()
	defer bg.mtx.Unlock()
	sCh, cCh := make(chan struct{}), make(chan error)
	bg.blockMap[blockPtr] = blockReturner{
		block:      block,
		startCh:    sCh,
		continueCh: cCh,
	}
	return sCh, cCh
}

// getBlock implements the interface for realBlockGetter.
func (bg *fakeBlockGetter) getBlock(ctx context.Context, kmd KeyMetadata,
	blockPtr BlockPointer, block Block) error {
	bg.mtx.RLock()
	defer bg.mtx.RUnlock()
	source, ok := bg.blockMap[blockPtr]
	if !ok {
		return errors.New("Block doesn't exist in fake block map")
	}
	cancelCh := make(chan struct{})
	if bg.respectCancel {
		go func() {
			<-ctx.Done()
			close(cancelCh)
		}()
	}
	// Wait until the caller tells us to continue
	for {
		select {
		case source.startCh <- struct{}{}:
		case err := <-source.continueCh:
			if err != nil {
				return err
			}
			return kbfscodec.Update(bg.codec, block, source.block)
		case <-cancelCh:
			return ctx.Err()
		}
	}
}

func (bg *fakeBlockGetter) assembleBlock(ctx context.Context,
	kmd KeyMetadata, ptr BlockPointer, block Block, buf []byte,
	serverHalf kbfscrypto.BlockCryptKeyServerHalf) error {
	source, ok := bg.blockMap[ptr]
	if !ok {
		return errors.New("Block doesn't exist in fake block map")
	}
	block.Set(source.block)
	return nil
}

func makeFakeFileBlock(t *testing.T, doHash bool) *FileBlock {
	buf := make([]byte, 16)
	_, err := rand.Read(buf)
	require.NoError(t, err)
	block := &FileBlock{
		Contents: buf,
	}
	if doHash {
		_ = block.GetHash()
	}
	return block
}

func TestBlockRetrievalWorkerBasic(t *testing.T) {
	t.Log("Test the basic ability of a worker to return a block.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	ptr1 := makeRandomBlockPointer(t)
	block1 := makeFakeFileBlock(t, false)
	_, continueCh1 := bg.setBlockToReturn(ptr1, block1)

	block := &FileBlock{}
	ch := q.Request(context.Background(), 1, makeKMD(), ptr1, block,
		NoCacheEntry)
	continueCh1 <- nil
	err := <-ch
	require.NoError(t, err)
	require.Equal(t, block1, block)
}

func TestBlockRetrievalWorkerMultipleWorkers(t *testing.T) {
	t.Log("Test the ability of multiple workers to retrieve concurrently.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 2, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	ptr1, ptr2 := makeRandomBlockPointer(t), makeRandomBlockPointer(t)
	block1, block2 := makeFakeFileBlock(t, false), makeFakeFileBlock(t, false)
	_, continueCh1 := bg.setBlockToReturn(ptr1, block1)
	_, continueCh2 := bg.setBlockToReturn(ptr2, block2)

	t.Log("Make 2 requests for 2 different blocks")
	block := &FileBlock{}
	req1Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, block,
		NoCacheEntry)
	req2Ch := q.Request(context.Background(), 1, makeKMD(), ptr2, block,
		NoCacheEntry)

	t.Log("Allow the second request to complete before the first")
	continueCh2 <- nil
	err := <-req2Ch
	require.NoError(t, err)
	require.Equal(t, block2, block)

	t.Log("Make another request for ptr2")
	req2Ch = q.Request(context.Background(), 1, makeKMD(), ptr2, block,
		NoCacheEntry)
	continueCh2 <- nil
	err = <-req2Ch
	require.NoError(t, err)
	require.Equal(t, block2, block)

	t.Log("Complete the ptr1 request")
	continueCh1 <- nil
	err = <-req1Ch
	require.NoError(t, err)
	require.Equal(t, block1, block)
}

func TestBlockRetrievalWorkerWithQueue(t *testing.T) {
	t.Log("Test the ability of a worker and queue to work correctly together.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	ptr1, ptr2, ptr3 := makeRandomBlockPointer(t), makeRandomBlockPointer(t),
		makeRandomBlockPointer(t)
	block1, block2, block3 := makeFakeFileBlock(t, false),
		makeFakeFileBlock(t, false), makeFakeFileBlock(t, false)
	startCh1, continueCh1 := bg.setBlockToReturn(ptr1, block1)
	_, continueCh2 := bg.setBlockToReturn(ptr2, block2)
	_, continueCh3 := bg.setBlockToReturn(ptr3, block3)

	t.Log("Make 3 retrievals for 3 different blocks. All retrievals after " +
		"the first should be queued.")
	block := &FileBlock{}
	testBlock1 := &FileBlock{}
	testBlock2 := &FileBlock{}
	req1Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, block,
		NoCacheEntry)
	req2Ch := q.Request(context.Background(), 1, makeKMD(), ptr2, block,
		NoCacheEntry)
	req3Ch := q.Request(context.Background(), 1, makeKMD(), ptr3, testBlock1,
		NoCacheEntry)
	// Ensure the worker picks up the first request
	<-startCh1
	t.Log("Make a high priority request for the third block, which should " +
		"complete next.")
	req4Ch := q.Request(context.Background(), 2, makeKMD(), ptr3, testBlock2,
		NoCacheEntry)

	t.Log("Allow the ptr1 retrieval to complete.")
	continueCh1 <- nil
	err := <-req1Ch
	require.NoError(t, err)
	require.Equal(t, block1, block)

	t.Log("Allow the ptr3 retrieval to complete. Both waiting requests " +
		"should complete.")
	continueCh3 <- nil
	err1 := <-req3Ch
	err2 := <-req4Ch
	require.NoError(t, err1)
	require.NoError(t, err2)
	require.Equal(t, block3, testBlock1)
	require.Equal(t, block3, testBlock2)

	t.Log("Complete the ptr2 retrieval.")
	continueCh2 <- nil
	err = <-req2Ch
	require.NoError(t, err)
	require.Equal(t, block2, block)
}

func TestBlockRetrievalWorkerCancel(t *testing.T) {
	t.Log("Test the ability of a worker to handle a request cancelation.")
	bg := newFakeBlockGetter(true)
	q := newBlockRetrievalQueue(0, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	ptr1 := makeRandomBlockPointer(t)
	block1 := makeFakeFileBlock(t, false)
	_, _ = bg.setBlockToReturn(ptr1, block1)

	block := &FileBlock{}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	ch := q.Request(ctx, 1, makeKMD(), ptr1, block, NoCacheEntry)
	err := <-ch
	require.EqualError(t, err, context.Canceled.Error())
}

func TestBlockRetrievalWorkerShutdown(t *testing.T) {
	t.Log("Test that worker shutdown works.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 0, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	w := newBlockRetrievalWorker(bg, q, onDemandWorker)
	require.NotNil(t, w)

	ptr1 := makeRandomBlockPointer(t)
	block1 := makeFakeFileBlock(t, false)
	_, continueCh := bg.setBlockToReturn(ptr1, block1)

	w.Shutdown()
	block := &FileBlock{}
	ctx, cancel := context.WithCancel(context.Background())
	// Ensure the context loop is stopped so the test doesn't leak goroutines
	defer cancel()
	ch := q.Request(ctx, 1, makeKMD(), ptr1, block, NoCacheEntry)
	shutdown := false
	select {
	case <-ch:
		t.Fatal("Expected not to retrieve a result from the Request.")
	case continueCh <- nil:
		t.Fatal("Expected the block getter not to be receiving.")
	default:
		shutdown = true
	}
	require.True(t, shutdown)

	// Ensure the test completes in a reasonable time.
	timer := time.NewTimer(10 * time.Second)
	doneCh := make(chan struct{})
	go func() {
		w.Shutdown()
		close(doneCh)
	}()
	select {
	case <-timer.C:
		t.Fatal("Expected another Shutdown not to block.")
	case <-doneCh:
	}
}

func TestBlockRetrievalWorkerMultipleBlockTypes(t *testing.T) {
	t.Log("Test that we can retrieve the same block into different block " +
		"types.")
	codec := kbfscodec.NewMsgpack()
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	t.Log("Setup source blocks")
	ptr1 := makeRandomBlockPointer(t)
	block1 := makeFakeFileBlock(t, false)
	_, continueCh1 := bg.setBlockToReturn(ptr1, block1)
	testCommonBlock := &CommonBlock{}
	err := kbfscodec.Update(codec, testCommonBlock, block1)
	require.NoError(t, err)

	t.Log("Make a retrieval for the same block twice, but with a different " +
		"target block type.")
	testBlock1 := &FileBlock{}
	testBlock2 := &CommonBlock{}
	req1Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, testBlock1,
		NoCacheEntry)
	req2Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, testBlock2,
		NoCacheEntry)

	t.Log("Allow the first ptr1 retrieval to complete.")
	continueCh1 <- nil
	err = <-req1Ch
	require.NoError(t, err)
	require.Equal(t, testBlock1, block1)

	t.Log("Allow the second ptr1 retrieval to complete.")
	continueCh1 <- nil
	err = <-req2Ch
	require.NoError(t, err)
	require.Equal(t, testBlock2, testCommonBlock)
}

func TestBlockRetrievalWorkerPrefetchedPriorityElevation(t *testing.T) {
	t.Log("TODO")
}
<commit_msg>block_retrieval_worker: Fix test and add new test for prefetching workers.<commit_after>// Copyright 2016 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.
package libkbfs

import (
	"crypto/rand"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/keybase/kbfs/kbfscodec"
	"github.com/keybase/kbfs/kbfscrypto"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context"
)

// blockReturner contains a block value to copy into requested blocks, and a
// channel to synchronize on with the worker.
type blockReturner struct {
	block      Block
	continueCh chan error
	startCh    chan struct{}
}

// fakeBlockGetter allows specifying and obtaining fake blocks.
type fakeBlockGetter struct {
	mtx           sync.RWMutex
	blockMap      map[BlockPointer]blockReturner
	codec         kbfscodec.Codec
	respectCancel bool
}

// newFakeBlockGetter returns a fakeBlockGetter.
func newFakeBlockGetter(respectCancel bool) *fakeBlockGetter {
	return &fakeBlockGetter{
		blockMap:      make(map[BlockPointer]blockReturner),
		codec:         kbfscodec.NewMsgpack(),
		respectCancel: respectCancel,
	}
}

// setBlockToReturn sets the block that will be returned for a given
// BlockPointer. Returns a writeable channel that getBlock will wait on, to
// allow synchronization of tests.
func (bg *fakeBlockGetter) setBlockToReturn(blockPtr BlockPointer,
	block Block) (startCh <-chan struct{}, continueCh chan<- error) {
	bg.mtx.Lock()
	defer bg.mtx.Unlock()
	sCh, cCh := make(chan struct{}), make(chan error)
	bg.blockMap[blockPtr] = blockReturner{
		block:      block,
		startCh:    sCh,
		continueCh: cCh,
	}
	return sCh, cCh
}

// getBlock implements the interface for realBlockGetter.
func (bg *fakeBlockGetter) getBlock(ctx context.Context, kmd KeyMetadata,
	blockPtr BlockPointer, block Block) error {
	bg.mtx.RLock()
	defer bg.mtx.RUnlock()
	source, ok := bg.blockMap[blockPtr]
	if !ok {
		return errors.New("Block doesn't exist in fake block map")
	}
	cancelCh := make(chan struct{})
	if bg.respectCancel {
		go func() {
			<-ctx.Done()
			close(cancelCh)
		}()
	}
	// Wait until the caller tells us to continue
	for {
		select {
		case source.startCh <- struct{}{}:
		case err := <-source.continueCh:
			if err != nil {
				return err
			}
			block.Set(source.block)
			return nil
		case <-cancelCh:
			return ctx.Err()
		}
	}
}

func (bg *fakeBlockGetter) assembleBlock(ctx context.Context,
	kmd KeyMetadata, ptr BlockPointer, block Block, buf []byte,
	serverHalf kbfscrypto.BlockCryptKeyServerHalf) error {
	source, ok := bg.blockMap[ptr]
	if !ok {
		return errors.New("Block doesn't exist in fake block map")
	}
	block.Set(source.block)
	return nil
}

func makeFakeFileBlock(t *testing.T, doHash bool) *FileBlock {
	buf := make([]byte, 16)
	_, err := rand.Read(buf)
	require.NoError(t, err)
	block := &FileBlock{
		Contents: buf,
	}
	if doHash {
		_ = block.GetHash()
	}
	return block
}

func TestBlockRetrievalWorkerBasic(t *testing.T) {
	t.Log("Test the basic ability of a worker to return a block.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	ptr1 := makeRandomBlockPointer(t)
	block1 := makeFakeFileBlock(t, false)
	_, continueCh1 := bg.setBlockToReturn(ptr1, block1)

	block := &FileBlock{}
	ch := q.Request(context.Background(), 1, makeKMD(), ptr1, block,
		NoCacheEntry)
	continueCh1 <- nil
	err := <-ch
	require.NoError(t, err)
	require.Equal(t, block1, block)
}

func TestBlockRetrievalWorkerMultipleWorkers(t *testing.T) {
	t.Log("Test the ability of multiple workers to retrieve concurrently.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 2, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	ptr1, ptr2 := makeRandomBlockPointer(t), makeRandomBlockPointer(t)
	block1, block2 := makeFakeFileBlock(t, false), makeFakeFileBlock(t, false)
	_, continueCh1 := bg.setBlockToReturn(ptr1, block1)
	_, continueCh2 := bg.setBlockToReturn(ptr2, block2)

	t.Log("Make 2 requests for 2 different blocks")
	block := &FileBlock{}
	req1Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, block,
		NoCacheEntry)
	req2Ch := q.Request(context.Background(), 1, makeKMD(), ptr2, block,
		NoCacheEntry)

	t.Log("Allow the second request to complete before the first")
	continueCh2 <- nil
	err := <-req2Ch
	require.NoError(t, err)
	require.Equal(t, block2, block)

	t.Log("Make another request for ptr2")
	req2Ch = q.Request(context.Background(), 1, makeKMD(), ptr2, block,
		NoCacheEntry)
	continueCh2 <- nil
	err = <-req2Ch
	require.NoError(t, err)
	require.Equal(t, block2, block)

	t.Log("Complete the ptr1 request")
	continueCh1 <- nil
	err = <-req1Ch
	require.NoError(t, err)
	require.Equal(t, block1, block)
}

func TestBlockRetrievalWorkerWithQueue(t *testing.T) {
	t.Log("Test the ability of a worker and queue to work correctly together.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	ptr1, ptr2, ptr3 := makeRandomBlockPointer(t), makeRandomBlockPointer(t),
		makeRandomBlockPointer(t)
	block1, block2, block3 := makeFakeFileBlock(t, false),
		makeFakeFileBlock(t, false), makeFakeFileBlock(t, false)
	startCh1, continueCh1 := bg.setBlockToReturn(ptr1, block1)
	_, continueCh2 := bg.setBlockToReturn(ptr2, block2)
	_, continueCh3 := bg.setBlockToReturn(ptr3, block3)

	t.Log("Make 3 retrievals for 3 different blocks. All retrievals after " +
		"the first should be queued.")
	block := &FileBlock{}
	testBlock1 := &FileBlock{}
	testBlock2 := &FileBlock{}
	req1Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, block,
		NoCacheEntry)
	req2Ch := q.Request(context.Background(), 1, makeKMD(), ptr2, block,
		NoCacheEntry)
	req3Ch := q.Request(context.Background(), 1, makeKMD(), ptr3, testBlock1,
		NoCacheEntry)
	// Ensure the worker picks up the first request
	<-startCh1
	t.Log("Make a high priority request for the third block, which should " +
		"complete next.")
	req4Ch := q.Request(context.Background(), 2, makeKMD(), ptr3, testBlock2,
		NoCacheEntry)

	t.Log("Allow the ptr1 retrieval to complete.")
	continueCh1 <- nil
	err := <-req1Ch
	require.NoError(t, err)
	require.Equal(t, block1, block)

	t.Log("Allow the ptr3 retrieval to complete. Both waiting requests " +
		"should complete.")
	continueCh3 <- nil
	err1 := <-req3Ch
	err2 := <-req4Ch
	require.NoError(t, err1)
	require.NoError(t, err2)
	require.Equal(t, block3, testBlock1)
	require.Equal(t, block3, testBlock2)

	t.Log("Complete the ptr2 retrieval.")
	continueCh2 <- nil
	err = <-req2Ch
	require.NoError(t, err)
	require.Equal(t, block2, block)
}

func TestBlockRetrievalWorkerCancel(t *testing.T) {
	t.Log("Test the ability of a worker to handle a request cancelation.")
	bg := newFakeBlockGetter(true)
	q := newBlockRetrievalQueue(0, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	ptr1 := makeRandomBlockPointer(t)
	block1 := makeFakeFileBlock(t, false)
	_, _ = bg.setBlockToReturn(ptr1, block1)

	block := &FileBlock{}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	ch := q.Request(ctx, 1, makeKMD(), ptr1, block, NoCacheEntry)
	err := <-ch
	require.EqualError(t, err, context.Canceled.Error())
}

func TestBlockRetrievalWorkerShutdown(t *testing.T) {
	t.Log("Test that worker shutdown works.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 0, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	w := newBlockRetrievalWorker(bg, q, onDemandWorker)
	require.NotNil(t, w)

	ptr1 := makeRandomBlockPointer(t)
	block1 := makeFakeFileBlock(t, false)
	_, continueCh := bg.setBlockToReturn(ptr1, block1)

	w.Shutdown()
	block := &FileBlock{}
	ctx, cancel := context.WithCancel(context.Background())
	// Ensure the context loop is stopped so the test doesn't leak goroutines
	defer cancel()
	ch := q.Request(ctx, 1, makeKMD(), ptr1, block, NoCacheEntry)
	shutdown := false
	select {
	case <-ch:
		t.Fatal("Expected not to retrieve a result from the Request.")
	case continueCh <- nil:
		t.Fatal("Expected the block getter not to be receiving.")
	default:
		shutdown = true
	}
	require.True(t, shutdown)

	// Ensure the test completes in a reasonable time.
	timer := time.NewTimer(10 * time.Second)
	doneCh := make(chan struct{})
	go func() {
		w.Shutdown()
		close(doneCh)
	}()
	select {
	case <-timer.C:
		t.Fatal("Expected another Shutdown not to block.")
	case <-doneCh:
	}
}

func TestBlockRetrievalWorkerMultipleBlockTypes(t *testing.T) {
	t.Log("Test that we can retrieve the same block into different block " +
		"types.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(0, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	t.Log("Setup source blocks")
	ptr1 := makeRandomBlockPointer(t)
	block1 := makeFakeFileBlock(t, false)
	_, continueCh1 := bg.setBlockToReturn(ptr1, block1)
	testCommonBlock := &CommonBlock{}
	testCommonBlock.Set(block1)
	require.Equal(t, &CommonBlock{}, testCommonBlock)

	t.Log("Make a retrieval for the same block twice, but with a different " +
		"target block type.")
	testBlock1 := &FileBlock{}
	testBlock2 := &CommonBlock{}
	req1Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, testBlock1,
		NoCacheEntry)
	req2Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, testBlock2,
		NoCacheEntry)

	t.Log("Allow the first ptr1 retrieval to complete.")
	continueCh1 <- nil
	err := <-req1Ch
	require.NoError(t, err)
	require.Equal(t, testBlock1, block1)

	t.Log("Allow the second ptr1 retrieval to complete.")
	continueCh1 <- nil
	err = <-req2Ch
	require.NoError(t, err)
	require.Equal(t, testBlock2, testCommonBlock)
}

func TestBlockRetrievalWorkerPrefetchedPriorityElevation(t *testing.T) {
	t.Log("Test that we can escalate the priority of a request and it " +
		"correctly switches workers.")
	bg := newFakeBlockGetter(false)
	q := newBlockRetrievalQueue(1, 1, newTestBlockRetrievalConfig(t, bg, nil))
	require.NotNil(t, q)
	defer q.Shutdown()

	t.Log("Setup source blocks")
	ptr1, ptr2 := makeRandomBlockPointer(t), makeRandomBlockPointer(t)
	block1, block2 := makeFakeFileBlock(t, false), makeFakeFileBlock(t, false)
	_, continueCh1 := bg.setBlockToReturn(ptr1, block1)
	_, continueCh2 := bg.setBlockToReturn(ptr2, block2)

	t.Log("Make a low-priority request. This will get to the worker.")
	testBlock1 := &FileBlock{}
	req1Ch := q.Request(context.Background(), 1, makeKMD(), ptr1, testBlock1,
		NoCacheEntry)

	t.Log("Make another low-priority request. This will block.")
	testBlock2 := &FileBlock{}
	req2Ch := q.Request(context.Background(), 1, makeKMD(), ptr2, testBlock2,
		NoCacheEntry)

	t.Log("Make an on-demand request for the same block as the blocked " +
		"request.")
	testBlock3 := &FileBlock{}
	req3Ch := q.Request(context.Background(), defaultOnDemandRequestPriority,
		makeKMD(), ptr2, testBlock3, NoCacheEntry)

	t.Log("Release the requests for the second block first. " +
		"Since the prefetch worker is still blocked, this confirms that the " +
		"escalation to an on-demand worker was successful.")
	continueCh2 <- nil
	err := <-req3Ch
	require.NoError(t, err)
	require.Equal(t, testBlock3, block2)
	err = <-req2Ch
	require.NoError(t, err)
	require.Equal(t, testBlock2, block2)

	t.Log("Allow the initial ptr1 request to complete.")
	continueCh1 <- nil
	err = <-req1Ch
	require.NoError(t, err)
	require.Equal(t, testBlock1, block1)
}
<|endoftext|>


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

package work

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"
	"testing"

	"cmd/go/internal/base"
	"cmd/go/internal/cfg"
	"cmd/go/internal/load"
)

func TestRemoveDevNull(t *testing.T) {
	fi, err := os.Lstat(os.DevNull)
	if err != nil {
		t.Skip(err)
	}
	if fi.Mode().IsRegular() {
		t.Errorf("Lstat(%s).Mode().IsRegular() = true; expected false", os.DevNull)
	}
	mayberemovefile(os.DevNull)
	_, err = os.Lstat(os.DevNull)
	if err != nil {
		t.Errorf("mayberemovefile(%s) did remove it; oops", os.DevNull)
	}
}

func TestSplitPkgConfigOutput(t *testing.T) {
	for _, test := range []struct {
		in   []byte
		want []string
	}{
		{[]byte(`-r:foo -L/usr/white\ space/lib -lfoo\ bar -lbar\ baz`), []string{"-r:foo", "-L/usr/white space/lib", "-lfoo bar", "-lbar baz"}},
		{[]byte(`-lextra\ fun\ arg\\`), []string{`-lextra fun arg\`}},
		{[]byte(`broken flag\`), []string{"broken", "flag"}},
		{[]byte("\textra     whitespace\r\n"), []string{"extra", "whitespace"}},
		{[]byte("     \r\n      "), nil},
	} {
		got := splitPkgConfigOutput(test.in)
		if !reflect.DeepEqual(got, test.want) {
			t.Errorf("splitPkgConfigOutput(%v) = %v; want %v", test.in, got, test.want)
		}
	}
}

func TestSharedLibName(t *testing.T) {
	// TODO(avdva) - make these values platform-specific
	prefix := "lib"
	suffix := ".so"
	testData := []struct {
		args      []string
		pkgs      []*load.Package
		expected  string
		expectErr bool
		rootedAt  string
	}{
		{
			args:     []string{"std"},
			pkgs:     []*load.Package{},
			expected: "std",
		},
		{
			args:     []string{"std", "cmd"},
			pkgs:     []*load.Package{},
			expected: "std,cmd",
		},
		{
			args:     []string{},
			pkgs:     []*load.Package{pkgImportPath("gopkg.in/somelib")},
			expected: "gopkg.in-somelib",
		},
		{
			args:     []string{"./..."},
			pkgs:     []*load.Package{pkgImportPath("somelib")},
			expected: "somelib",
			rootedAt: "somelib",
		},
		{
			args:     []string{"../somelib", "../somelib"},
			pkgs:     []*load.Package{pkgImportPath("somelib")},
			expected: "somelib",
		},
		{
			args:     []string{"../lib1", "../lib2"},
			pkgs:     []*load.Package{pkgImportPath("gopkg.in/lib1"), pkgImportPath("gopkg.in/lib2")},
			expected: "gopkg.in-lib1,gopkg.in-lib2",
		},
		{
			args: []string{"./..."},
			pkgs: []*load.Package{
				pkgImportPath("gopkg.in/dir/lib1"),
				pkgImportPath("gopkg.in/lib2"),
				pkgImportPath("gopkg.in/lib3"),
			},
			expected: "gopkg.in",
			rootedAt: "gopkg.in",
		},
		{
			args:      []string{"std", "../lib2"},
			pkgs:      []*load.Package{},
			expectErr: true,
		},
		{
			args:      []string{"all", "./"},
			pkgs:      []*load.Package{},
			expectErr: true,
		},
		{
			args:      []string{"cmd", "fmt"},
			pkgs:      []*load.Package{},
			expectErr: true,
		},
	}
	for _, data := range testData {
		func() {
			if data.rootedAt != "" {
				tmpGopath, err := ioutil.TempDir("", "gopath")
				if err != nil {
					t.Fatal(err)
				}
				oldGopath := cfg.BuildContext.GOPATH
				defer func() {
					cfg.BuildContext.GOPATH = oldGopath
					os.Chdir(base.Cwd)
					err := os.RemoveAll(tmpGopath)
					if err != nil {
						t.Error(err)
					}
				}()
				root := filepath.Join(tmpGopath, "src", data.rootedAt)
				err = os.MkdirAll(root, 0755)
				if err != nil {
					t.Fatal(err)
				}
				cfg.BuildContext.GOPATH = tmpGopath
				os.Chdir(root)
			}
			computed, err := libname(data.args, data.pkgs)
			if err != nil {
				if !data.expectErr {
					t.Errorf("libname returned an error %q, expected a name", err.Error())
				}
			} else if data.expectErr {
				t.Errorf("libname returned %q, expected an error", computed)
			} else {
				expected := prefix + data.expected + suffix
				if expected != computed {
					t.Errorf("libname returned %q, expected %q", computed, expected)
				}
			}
		}()
	}
}

func pkgImportPath(pkgpath string) *load.Package {
	return &load.Package{
		PackagePublic: load.PackagePublic{
			ImportPath: pkgpath,
		},
	}
}

// When installing packages, the installed package directory should
// respect the group sticky bit and group name of the destination
// directory.
// See https://golang.org/issue/18878.
func TestRespectGroupSticky(t *testing.T) {
	if runtime.GOOS == "nacl" {
		t.Skip("can't set group sticky bit with chmod on nacl")
	}

	var b Builder

	// Check that `cp` is called instead of `mv` by looking at the output
	// of `(*Builder).ShowCmd` afterwards as a sanity check.
	cfg.BuildX = true
	var cmdBuf bytes.Buffer
	b.Print = func(a ...interface{}) (int, error) {
		return cmdBuf.WriteString(fmt.Sprint(a...))
	}

	stickydir, err := ioutil.TempDir("", "GroupSticky")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(stickydir)

	testdir, err := ioutil.TempDir(stickydir, "testdir")
	if err != nil {
		t.Fatal(err)
	}

	// Change testdir's permissions to include group sticky bit.
	if err := os.Chmod(testdir, 0755|os.ModeSetgid); err != nil {
		t.Fatal(err)
	}

	pkgfile, err := ioutil.TempFile("", "pkgfile")
	if err != nil {
		t.Fatalf("ioutil.TempFile(\"\", \"pkgfile\"): %v", err)
	}
	defer os.Remove(pkgfile.Name())
	defer pkgfile.Close()

	stickyFile := filepath.Join(testdir, "sticky")
	if err := b.moveOrCopyFile(nil, stickyFile, pkgfile.Name(), 0666, true); err != nil {
		t.Fatalf("moveOrCopyFile: %v", err)
	}

	got := strings.TrimSpace(cmdBuf.String())
	want := b.fmtcmd("", "cp %s %s", pkgfile.Name(), stickyFile)
	if got != want {
		t.Fatalf("moveOrCopyFile(%q, %q): want %q, got %q", stickyFile, pkgfile.Name(), want, got)
	}
}
<commit_msg>cmd/go/internal/work: more TestRespectSetgidDir fixes<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package work

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"
	"testing"

	"cmd/go/internal/base"
	"cmd/go/internal/cfg"
	"cmd/go/internal/load"
)

func TestRemoveDevNull(t *testing.T) {
	fi, err := os.Lstat(os.DevNull)
	if err != nil {
		t.Skip(err)
	}
	if fi.Mode().IsRegular() {
		t.Errorf("Lstat(%s).Mode().IsRegular() = true; expected false", os.DevNull)
	}
	mayberemovefile(os.DevNull)
	_, err = os.Lstat(os.DevNull)
	if err != nil {
		t.Errorf("mayberemovefile(%s) did remove it; oops", os.DevNull)
	}
}

func TestSplitPkgConfigOutput(t *testing.T) {
	for _, test := range []struct {
		in   []byte
		want []string
	}{
		{[]byte(`-r:foo -L/usr/white\ space/lib -lfoo\ bar -lbar\ baz`), []string{"-r:foo", "-L/usr/white space/lib", "-lfoo bar", "-lbar baz"}},
		{[]byte(`-lextra\ fun\ arg\\`), []string{`-lextra fun arg\`}},
		{[]byte(`broken flag\`), []string{"broken", "flag"}},
		{[]byte("\textra     whitespace\r\n"), []string{"extra", "whitespace"}},
		{[]byte("     \r\n      "), nil},
	} {
		got := splitPkgConfigOutput(test.in)
		if !reflect.DeepEqual(got, test.want) {
			t.Errorf("splitPkgConfigOutput(%v) = %v; want %v", test.in, got, test.want)
		}
	}
}

func TestSharedLibName(t *testing.T) {
	// TODO(avdva) - make these values platform-specific
	prefix := "lib"
	suffix := ".so"
	testData := []struct {
		args      []string
		pkgs      []*load.Package
		expected  string
		expectErr bool
		rootedAt  string
	}{
		{
			args:     []string{"std"},
			pkgs:     []*load.Package{},
			expected: "std",
		},
		{
			args:     []string{"std", "cmd"},
			pkgs:     []*load.Package{},
			expected: "std,cmd",
		},
		{
			args:     []string{},
			pkgs:     []*load.Package{pkgImportPath("gopkg.in/somelib")},
			expected: "gopkg.in-somelib",
		},
		{
			args:     []string{"./..."},
			pkgs:     []*load.Package{pkgImportPath("somelib")},
			expected: "somelib",
			rootedAt: "somelib",
		},
		{
			args:     []string{"../somelib", "../somelib"},
			pkgs:     []*load.Package{pkgImportPath("somelib")},
			expected: "somelib",
		},
		{
			args:     []string{"../lib1", "../lib2"},
			pkgs:     []*load.Package{pkgImportPath("gopkg.in/lib1"), pkgImportPath("gopkg.in/lib2")},
			expected: "gopkg.in-lib1,gopkg.in-lib2",
		},
		{
			args: []string{"./..."},
			pkgs: []*load.Package{
				pkgImportPath("gopkg.in/dir/lib1"),
				pkgImportPath("gopkg.in/lib2"),
				pkgImportPath("gopkg.in/lib3"),
			},
			expected: "gopkg.in",
			rootedAt: "gopkg.in",
		},
		{
			args:      []string{"std", "../lib2"},
			pkgs:      []*load.Package{},
			expectErr: true,
		},
		{
			args:      []string{"all", "./"},
			pkgs:      []*load.Package{},
			expectErr: true,
		},
		{
			args:      []string{"cmd", "fmt"},
			pkgs:      []*load.Package{},
			expectErr: true,
		},
	}
	for _, data := range testData {
		func() {
			if data.rootedAt != "" {
				tmpGopath, err := ioutil.TempDir("", "gopath")
				if err != nil {
					t.Fatal(err)
				}
				oldGopath := cfg.BuildContext.GOPATH
				defer func() {
					cfg.BuildContext.GOPATH = oldGopath
					os.Chdir(base.Cwd)
					err := os.RemoveAll(tmpGopath)
					if err != nil {
						t.Error(err)
					}
				}()
				root := filepath.Join(tmpGopath, "src", data.rootedAt)
				err = os.MkdirAll(root, 0755)
				if err != nil {
					t.Fatal(err)
				}
				cfg.BuildContext.GOPATH = tmpGopath
				os.Chdir(root)
			}
			computed, err := libname(data.args, data.pkgs)
			if err != nil {
				if !data.expectErr {
					t.Errorf("libname returned an error %q, expected a name", err.Error())
				}
			} else if data.expectErr {
				t.Errorf("libname returned %q, expected an error", computed)
			} else {
				expected := prefix + data.expected + suffix
				if expected != computed {
					t.Errorf("libname returned %q, expected %q", computed, expected)
				}
			}
		}()
	}
}

func pkgImportPath(pkgpath string) *load.Package {
	return &load.Package{
		PackagePublic: load.PackagePublic{
			ImportPath: pkgpath,
		},
	}
}

// When installing packages, the installed package directory should
// respect the SetGID bit and group name of the destination
// directory.
// See https://golang.org/issue/18878.
func TestRespectSetgidDir(t *testing.T) {
	if runtime.GOOS == "nacl" {
		t.Skip("can't set SetGID bit with chmod on nacl")
	}

	var b Builder

	// Check that `cp` is called instead of `mv` by looking at the output
	// of `(*Builder).ShowCmd` afterwards as a sanity check.
	cfg.BuildX = true
	var cmdBuf bytes.Buffer
	b.Print = func(a ...interface{}) (int, error) {
		return cmdBuf.WriteString(fmt.Sprint(a...))
	}

	setgiddir, err := ioutil.TempDir("", "SetGroupID")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(setgiddir)

	if runtime.GOOS == "freebsd" {
		err = os.Chown(setgiddir, os.Getuid(), os.Getgid())
		if err != nil {
			t.Fatal(err)
		}
	}

	// Change setgiddir's permissions to include the SetGID bit.
	if err := os.Chmod(setgiddir, 0755|os.ModeSetgid); err != nil {
		t.Fatal(err)
	}

	pkgfile, err := ioutil.TempFile("", "pkgfile")
	if err != nil {
		t.Fatalf("ioutil.TempFile(\"\", \"pkgfile\"): %v", err)
	}
	defer os.Remove(pkgfile.Name())
	defer pkgfile.Close()

	dirGIDFile := filepath.Join(setgiddir, "setgid")
	if err := b.moveOrCopyFile(nil, dirGIDFile, pkgfile.Name(), 0666, true); err != nil {
		t.Fatalf("moveOrCopyFile: %v", err)
	}

	got := strings.TrimSpace(cmdBuf.String())
	want := b.fmtcmd("", "cp %s %s", pkgfile.Name(), dirGIDFile)
	if got != want {
		t.Fatalf("moveOrCopyFile(%q, %q): want %q, got %q", dirGIDFile, pkgfile.Name(), want, got)
	}
}
<|endoftext|>


<commit_before>package byteexec

import (
	"io/ioutil"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/getlantern/testify/assert"
)

const (
	program = "helloworld"

	concurrency = 10
)

func TestExec(t *testing.T) {
	data, err := Asset(program)
	if err != nil {
		t.Fatalf("Unable to read helloworld program: %s", err)
	}
	be := createByteExec(t, data)

	// Concurrently create some other BEs and make sure they don't get errors
	var wg sync.WaitGroup
	wg.Add(concurrency)
	for i := 0; i < concurrency; i++ {
		_, err := New(data, program)
		assert.NoError(t, err, "Concurrent New should have succeeded")
		wg.Done()
	}
	wg.Wait()

	originalInfo := testByteExec(t, be)

	// Recreate be and make sure file is reused
	be = createByteExec(t, data)
	updatedInfo := testByteExec(t, be)
	assert.Equal(t, originalInfo.ModTime(), updatedInfo.ModTime(), "File modification time should be unchanged after creating new ByteExec")

	// Now mess with the file permissions and make sure that we can still run
	err = os.Chmod(be.filename, 0655)
	if err != nil {
		t.Fatalf("Unable to chmod test executable %s: %s", be.filename, err)
	}
	be = createByteExec(t, data)
	updatedInfo = testByteExec(t, be)
	assert.Equal(t, fileMode, updatedInfo.Mode(), "File mode is changed back to %v", fileMode)

	// Now mess with the file contents and make sure it gets overwritten on next
	// ByteExec
	ioutil.WriteFile(be.filename, []byte("Junk"), 0755)
	be = createByteExec(t, data)
	updatedInfo = testByteExec(t, be)
	assert.NotEqual(t, originalInfo.ModTime(), updatedInfo.ModTime(), "File modification time should be changed after creating new ByteExec on bad data")
}

func createByteExec(t *testing.T, data []byte) *Exec {
	// Sleep 1 second to give file timestamp a chance to increase
	time.Sleep(1 * time.Second)

	be, err := New(data, program)
	if err != nil {
		t.Fatalf("Unable to create new ByteExec: %s", err)
	}
	return be
}

func testByteExec(t *testing.T, be *Exec) os.FileInfo {
	cmd := be.Command()
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Errorf("Unable to run helloworld program: %s", err)
	}
	assert.Equal(t, "Hello world\n", string(out), "Should receive expected output from helloworld program")

	fileInfo, err := os.Stat(be.filename)
	if err != nil {
		t.Fatalf("Unable to re-stat file %s: %s", be.filename, err)
	}
	return fileInfo
}
<commit_msg>getlantern/lantern#2092 Fixed test for filemode<commit_after>package byteexec

import (
	"io/ioutil"
	"os"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/getlantern/testify/assert"
)

const (
	program = "helloworld"

	concurrency = 10
)

func TestExec(t *testing.T) {
	data, err := Asset(program)
	if err != nil {
		t.Fatalf("Unable to read helloworld program: %s", err)
	}
	be := createByteExec(t, data)

	// Concurrently create some other BEs and make sure they don't get errors
	var wg sync.WaitGroup
	wg.Add(concurrency)
	for i := 0; i < concurrency; i++ {
		_, err := New(data, program)
		assert.NoError(t, err, "Concurrent New should have succeeded")
		wg.Done()
	}
	wg.Wait()

	originalInfo := testByteExec(t, be)

	// Recreate be and make sure file is reused
	be = createByteExec(t, data)
	updatedInfo := testByteExec(t, be)
	assert.Equal(t, originalInfo.ModTime(), updatedInfo.ModTime(), "File modification time should be unchanged after creating new ByteExec")

	// Now mess with the file permissions and make sure that we can still run
	err = os.Chmod(be.filename, 0655)
	if err != nil {
		t.Fatalf("Unable to chmod test executable %s: %s", be.filename, err)
	}
	be = createByteExec(t, data)
	updatedInfo = testByteExec(t, be)
	assert.True(t, fileMode == updatedInfo.Mode(), fmt.Sprintf("File mode was %v instead of %v", updatedInfo.Mode(), fileMode))

	// Now mess with the file contents and make sure it gets overwritten on next
	// ByteExec
	ioutil.WriteFile(be.filename, []byte("Junk"), 0755)
	be = createByteExec(t, data)
	updatedInfo = testByteExec(t, be)
	assert.NotEqual(t, originalInfo.ModTime(), updatedInfo.ModTime(), "File modification time should be changed after creating new ByteExec on bad data")
}

func createByteExec(t *testing.T, data []byte) *Exec {
	// Sleep 1 second to give file timestamp a chance to increase
	time.Sleep(1 * time.Second)

	be, err := New(data, program)
	if err != nil {
		t.Fatalf("Unable to create new ByteExec: %s", err)
	}
	return be
}

func testByteExec(t *testing.T, be *Exec) os.FileInfo {
	cmd := be.Command()
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Errorf("Unable to run helloworld program: %s", err)
	}
	assert.Equal(t, "Hello world\n", string(out), "Should receive expected output from helloworld program")

	fileInfo, err := os.Stat(be.filename)
	if err != nil {
		t.Fatalf("Unable to re-stat file %s: %s", be.filename, err)
	}
	return fileInfo
}
<|endoftext|>


<commit_before>package twitter

import (
	"fmt"
	"net/url"
	"os"
	"time"

	"github.com/pkg/errors"
	"github.com/utahta/go-twitter/types"
	"github.com/utahta/momoclo-channel/appengine/lib/log"
	"github.com/utahta/momoclo-channel/appengine/model"
	"github.com/utahta/momoclo-crawler"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

// Tweet text message
func TweetMessage(ctx context.Context, text string) error {
	if disabled() {
		return nil
	}

	c, err := newClient(ctx)
	if err != nil {
		return err
	}

	if _, err := c.Tweet(text, nil); err != nil {
		return err
	}
	return nil
}

// Tweet channel
func TweetChannel(ctx context.Context, ch *crawler.Channel) error {
	if disabled() {
		return nil
	}

	reqCtx, cancel := context.WithTimeout(ctx, 540*time.Second)
	defer cancel()
	glog := log.GaeLog(ctx)
	eg := new(errgroup.Group)
	for _, item := range ch.Items {
		item := item
		eg.Go(func() error {
			if err := model.NewTweetItem(item).Put(ctx); err != nil {
				return nil
			}

			if err := tweetChannelItem(reqCtx, ch.Title, item); err != nil {
				glog.Error(err)
				return err
			}
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		return errors.Wrap(err, "errors occurred in TweetChannel.")
	}
	return nil
}

func tweetChannelItem(ctx context.Context, title string, item *crawler.ChannelItem) error {
	glog := log.GaeLog(ctx)
	c, err := newClient(ctx)
	if err != nil {
		return err
	}

	const maxUploadMediaLen = 4
	var images [][]string
	var tmp []string
	for _, image := range item.Images {
		tmp = append(tmp, image.Url)
		if len(tmp) == maxUploadMediaLen {
			images = append(images, tmp)
			tmp = nil
		}
	}
	if len(tmp) > 0 {
		images = append(images, tmp)
	}
	videos := item.Videos
	text := truncateText(title, item)

	var tweets *types.Tweets
	if len(images) > 0 {
		tweets, err = c.TweetImageURLs(text, images[0], nil)
		images = images[1:]
	} else if len(videos) > 0 {
		tweets, err = c.TweetVideoURL(text, videos[0].Url, "video/mp4", nil)
		videos = videos[1:]
	} else {
		tweets, err = c.Tweet(text, nil)
	}

	if err != nil {
		glog.Errorf("failed to post tweet. text:%s err:%v", text, err)
		return err
	}
	glog.Infof("Post tweet. text:%s images:%v videos:%v", text, len(item.Images), len(item.Videos))

	if len(images) > 0 {
		for _, urlsStr := range images {
			v := url.Values{}
			v.Set("in_reply_to_status_id", tweets.IDStr)
			tweets, err = c.TweetImageURLs("", urlsStr, v)
			if err != nil {
				glog.Errorf("failed to post images. urls:%v err:%v", urlsStr, err)
			}
		}
	}

	if len(videos) > 0 {
		for _, video := range videos {
			v := url.Values{}
			v.Set("in_reply_to_status_id", tweets.IDStr)
			tweets, err = c.TweetVideoURL("", video.Url, "video/mp4", v)
			if err != nil {
				glog.Errorf("failed to post video. url:%v err:%v", video.Url, err)
			}
		}
	}
	return nil
}

func truncateText(channelTitle string, item *crawler.ChannelItem) string {
	const maxTweetTextLen = 77 // ハッシュタグや URL や画像を除いて投稿可能な文字数

	title := []rune(fmt.Sprintf("%s %s", channelTitle, item.Title))
	if len(title) >= maxTweetTextLen {
		title = append(title[0:maxTweetTextLen-3], []rune("...")...)
	}
	return fmt.Sprintf("%s %s #momoclo #ももクロ", string(title), item.Url)
}

// if true disable tweet
func disabled() bool {
	e := os.Getenv("TWEET_DISABLE")
	if e != "" {
		return true
	}
	return false
}
<commit_msg>Fix disabled position<commit_after>package twitter

import (
	"fmt"
	"net/url"
	"os"
	"time"

	"github.com/pkg/errors"
	"github.com/utahta/go-twitter/types"
	"github.com/utahta/momoclo-channel/appengine/lib/log"
	"github.com/utahta/momoclo-channel/appengine/model"
	"github.com/utahta/momoclo-crawler"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

// Tweet text message
func TweetMessage(ctx context.Context, text string) error {
	if disabled() {
		return nil
	}

	c, err := newClient(ctx)
	if err != nil {
		return err
	}

	if _, err := c.Tweet(text, nil); err != nil {
		return err
	}
	return nil
}

// Tweet channel
func TweetChannel(ctx context.Context, ch *crawler.Channel) error {
	reqCtx, cancel := context.WithTimeout(ctx, 540*time.Second)
	defer cancel()

	glog := log.GaeLog(ctx)
	eg := new(errgroup.Group)
	for _, item := range ch.Items {
		item := item
		eg.Go(func() error {
			if err := model.NewTweetItem(item).Put(ctx); err != nil {
				return nil
			}

			if err := tweetChannelItem(reqCtx, ch.Title, item); err != nil {
				glog.Error(err)
				return err
			}
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		return errors.Wrap(err, "errors occurred in TweetChannel.")
	}
	return nil
}

func tweetChannelItem(ctx context.Context, title string, item *crawler.ChannelItem) error {
	if disabled() {
		return nil
	}

	glog := log.GaeLog(ctx)
	c, err := newClient(ctx)
	if err != nil {
		return err
	}

	const maxUploadMediaLen = 4
	var images [][]string
	var tmp []string
	for _, image := range item.Images {
		tmp = append(tmp, image.Url)
		if len(tmp) == maxUploadMediaLen {
			images = append(images, tmp)
			tmp = nil
		}
	}
	if len(tmp) > 0 {
		images = append(images, tmp)
	}
	videos := item.Videos
	text := truncateText(title, item)

	var tweets *types.Tweets
	if len(images) > 0 {
		tweets, err = c.TweetImageURLs(text, images[0], nil)
		images = images[1:]
	} else if len(videos) > 0 {
		tweets, err = c.TweetVideoURL(text, videos[0].Url, "video/mp4", nil)
		videos = videos[1:]
	} else {
		tweets, err = c.Tweet(text, nil)
	}

	if err != nil {
		glog.Errorf("failed to post tweet. text:%s err:%v", text, err)
		return err
	}
	glog.Infof("Post tweet. text:%s images:%v videos:%v", text, len(item.Images), len(item.Videos))

	if len(images) > 0 {
		for _, urlsStr := range images {
			v := url.Values{}
			v.Set("in_reply_to_status_id", tweets.IDStr)
			tweets, err = c.TweetImageURLs("", urlsStr, v)
			if err != nil {
				glog.Errorf("failed to post images. urls:%v err:%v", urlsStr, err)
			}
		}
	}

	if len(videos) > 0 {
		for _, video := range videos {
			v := url.Values{}
			v.Set("in_reply_to_status_id", tweets.IDStr)
			tweets, err = c.TweetVideoURL("", video.Url, "video/mp4", v)
			if err != nil {
				glog.Errorf("failed to post video. url:%v err:%v", video.Url, err)
			}
		}
	}
	return nil
}

func truncateText(channelTitle string, item *crawler.ChannelItem) string {
	const maxTweetTextLen = 77 // ハッシュタグや URL や画像を除いて投稿可能な文字数

	title := []rune(fmt.Sprintf("%s %s", channelTitle, item.Title))
	if len(title) >= maxTweetTextLen {
		title = append(title[0:maxTweetTextLen-3], []rune("...")...)
	}
	return fmt.Sprintf("%s %s #momoclo #ももクロ", string(title), item.Url)
}

// if true disable tweet
func disabled() bool {
	e := os.Getenv("TWEET_DISABLE")
	if e != "" {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package mpfluentd

import (
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.fluentd")

func metricName(names ...string) string {
	return strings.Join(names, ".")
}

// FluentdPlugin mackerel plugin for Fluentd
type FluentdPlugin struct {
	Target          string
	Prefix          string
	Tempfile        string
	pluginType      string
	pluginIDPattern *regexp.Regexp
	extendedMetrics []string

	plugins []FluentdPluginMetrics
}

// MetricKeyPrefix interface for PluginWithPrefix
func (f FluentdPlugin) MetricKeyPrefix() string {
	if f.Prefix == "" {
		f.Prefix = "fluentd"
	}
	return f.Prefix
}

// FluentdPluginMetrics metrics
type FluentdPluginMetrics struct {
	RetryCount            uint64 `json:"retry_count"`
	BufferQueueLength     uint64 `json:"buffer_queue_length"`
	BufferTotalQueuedSize uint64 `json:"buffer_total_queued_size"`
	OutputPlugin          bool   `json:"output_plugin"`
	Type                  string `json:"type"`
	PluginCategory        string `json:"plugin_category"`
	PluginID              string `json:"plugin_id"`
	normalizedPluginID    string

	// extended metrics fluentd >= 1.6
	// https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released
	EmitRecords                      uint64  `json:"emit_records"`
	EmitCount                        uint64  `json:"emit_count"`
	WriteCount                       uint64  `json:"write_count"`
	RollbackCount                    uint64  `json:"rollback_count"`
	SlowFlushCount                   uint64  `json:"slow_flush_count"`
	FlushTimeCount                   uint64  `json:"flush_time_count"`
	BufferStageLength                uint64  `json:"buffer_stage_length"`
	BufferStageByteSize              uint64  `json:"buffer_stage_byte_size"`
	BufferQueueByteSize              uint64  `json:"buffer_queue_byte_size"`
	BufferAvailableBufferSpaceRatios float64 `json:"buffer_available_buffer_space_ratios"`
}

func (fpm FluentdPluginMetrics) getExtended(name string) float64 {
	switch name {
	case "emit_records":
		return float64(fpm.EmitRecords)
	case "emit_count":
		return float64(fpm.EmitCount)
	case "write_count":
		return float64(fpm.WriteCount)
	case "rollback_count":
		return float64(fpm.RollbackCount)
	case "slow_flush_count":
		return float64(fpm.SlowFlushCount)
	case "flush_time_count":
		return float64(fpm.FlushTimeCount)
	case "buffer_stage_length":
		return float64(fpm.BufferStageLength)
	case "buffer_stage_byte_size":
		return float64(fpm.BufferStageByteSize)
	case "buffer_queue_byte_size":
		return float64(fpm.BufferQueueByteSize)
	case "buffer_available_buffer_space_ratios":
		return fpm.BufferAvailableBufferSpaceRatios
	}
	return 0
}

// FluentMonitorJSON monitor json
type FluentMonitorJSON struct {
	Plugins []FluentdPluginMetrics `json:"plugins"`
}

var normalizePluginIDRe = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizePluginID(in string) string {
	return normalizePluginIDRe.ReplaceAllString(in, "_")
}

func (fpm FluentdPluginMetrics) getNormalizedPluginID() string {
	if fpm.normalizedPluginID == "" {
		fpm.normalizedPluginID = normalizePluginID(fpm.PluginID)
	}
	return fpm.normalizedPluginID
}

func (f *FluentdPlugin) parseStats(body []byte) (map[string]interface{}, error) {
	var j FluentMonitorJSON
	err := json.Unmarshal(body, &j)
	f.plugins = j.Plugins

	metrics := make(map[string]interface{})
	for _, p := range f.plugins {
		if f.nonTargetPlugin(p) {
			continue
		}
		pid := p.getNormalizedPluginID()
		metrics[metricName("retry_count", pid)] = float64(p.RetryCount)
		metrics[metricName("buffer_queue_length", pid)] = float64(p.BufferQueueLength)
		metrics[metricName("buffer_total_queued_size", pid)] = float64(p.BufferTotalQueuedSize)
		for _, name := range f.extendedMetrics {
			metrics[metricName(name, pid)] = p.getExtended(name)
		}
	}
	return metrics, err
}

func (f *FluentdPlugin) nonTargetPlugin(plugin FluentdPluginMetrics) bool {
	if plugin.PluginCategory != "output" {
		return true
	}
	if f.pluginType != "" && f.pluginType != plugin.Type {
		return true
	}
	if f.pluginIDPattern != nil && !f.pluginIDPattern.MatchString(plugin.PluginID) {
		return true
	}
	return false
}

// FetchMetrics interface for mackerelplugin
func (f FluentdPlugin) FetchMetrics() (map[string]interface{}, error) {
	req, err := http.NewRequest(http.MethodGet, f.Target, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "mackerel-plugin-fluentd")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return f.parseStats(body)
}

var defaultGraphs = map[string]mp.Graphs{
	"retry_count": {
		Label: "retry count",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_queue_length": {
		Label: "queue length",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_total_queued_size": {
		Label: "buffer total queued size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

var extendedGraphs = map[string]mp.Graphs{
	"emit_records": {
		Label: "emitted records",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"emit_count": {
		Label: "emit calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"write_count": {
		Label: "write/try_write calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"rollback_count": {
		Label: "rollbacks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"slow_flush_count": {
		Label: "slow flushes",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"flush_time_count": {
		Label: "buffer flush time in msec",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"buffer_stage_length": {
		Label: "length of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_stage_byte_size": {
		Label: "bytesize of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_queue_byte_size": {
		Label: "bytesize of queued buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_available_buffer_space_ratios": {
		Label: "available space for buffer",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

// GraphDefinition interface for mackerelplugin
func (f FluentdPlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(f.Prefix)
	graphs := make(map[string]mp.Graphs, len(defaultGraphs))
	for key, g := range defaultGraphs {
		graphs[key] = mp.Graphs{
			Label:   (labelPrefix + " " + g.Label),
			Unit:    g.Unit,
			Metrics: g.Metrics,
		}
	}
	for _, name := range f.extendedMetrics {
		fullName := metricName(name)
		if g, ok := extendedGraphs[fullName]; ok {
			graphs[fullName] = mp.Graphs{
				Label:   (labelPrefix + " " + g.Label),
				Unit:    g.Unit,
				Metrics: g.Metrics,
			}
		}
	}
	return graphs
}

// Do the plugin
func Do() {
	host := flag.String("host", "localhost", "fluentd monitor_agent host")
	port := flag.String("port", "24220", "fluentd monitor_agent port")
	pluginType := flag.String("plugin-type", "", "Gets the metric that matches this plugin type")
	pluginIDPatternString := flag.String("plugin-id-pattern", "", "Gets the metric that matches this plugin id pattern")
	prefix := flag.String("metric-key-prefix", "fluentd", "Metric key prefix")
	tempFile := flag.String("tempfile", "", "Temp file name")
	extendedMetricNames := flag.String("extended_metrics", "", "extended metric names joind with ',' or 'all' (fluentd >= v1.6.0)")
	flag.Parse()

	var pluginIDPattern *regexp.Regexp
	var err error
	if *pluginIDPatternString != "" {
		pluginIDPattern, err = regexp.Compile(*pluginIDPatternString)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to exec mackerel-plugin-fluentd: invalid plugin-id-pattern: %s\n", err)
			os.Exit(1)
		}
	}

	var extendedMetrics []string
	switch *extendedMetricNames {
	case "all":
		for key := range extendedGraphs {
			extendedMetrics = append(extendedMetrics, key)
		}
	case "":
	default:
		for _, name := range strings.Split(*extendedMetricNames, ",") {
			fullName := metricName(name)
			if _, exists := extendedGraphs[fullName]; !exists {
				fmt.Fprintf(os.Stderr, "extended_metrics %s is not supported. See also https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released\n", name)
				os.Exit(1)
			}
			extendedMetrics = append(extendedMetrics, name)
		}
	}
	f := FluentdPlugin{
		Target:          fmt.Sprintf("http://%s:%s/api/plugins.json", *host, *port),
		Prefix:          *prefix,
		Tempfile:        *tempFile,
		pluginType:      *pluginType,
		pluginIDPattern: pluginIDPattern,
		extendedMetrics: extendedMetrics,
	}

	helper := mp.NewMackerelPlugin(f)

	helper.Tempfile = *tempFile
	if *tempFile == "" {
		tempFileSuffix := []string{*host, *port}
		if *pluginType != "" {
			tempFileSuffix = append(tempFileSuffix, *pluginType)
		}
		if *pluginIDPatternString != "" {
			tempFileSuffix = append(tempFileSuffix, fmt.Sprintf("%x", md5.Sum([]byte(*pluginIDPatternString))))
		}
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-fluentd-%s", strings.Join(tempFileSuffix, "-")))
	}

	helper.Run()
}
<commit_msg>[fluentd] add alias for backward compatibility.<commit_after>package mpfluentd

import (
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.fluentd")

func metricName(names ...string) string {
	return strings.Join(names, ".")
}

// FluentdPlugin mackerel plugin for Fluentd
type FluentdPlugin struct {
	Target          string
	Prefix          string
	Tempfile        string
	pluginType      string
	pluginIDPattern *regexp.Regexp
	extendedMetrics []string

	plugins []FluentdPluginMetrics
}

type FluentdMetrics = FluentdPlugin

// MetricKeyPrefix interface for PluginWithPrefix
func (f FluentdPlugin) MetricKeyPrefix() string {
	if f.Prefix == "" {
		f.Prefix = "fluentd"
	}
	return f.Prefix
}

// FluentdPluginMetrics metrics
type FluentdPluginMetrics struct {
	RetryCount            uint64 `json:"retry_count"`
	BufferQueueLength     uint64 `json:"buffer_queue_length"`
	BufferTotalQueuedSize uint64 `json:"buffer_total_queued_size"`
	OutputPlugin          bool   `json:"output_plugin"`
	Type                  string `json:"type"`
	PluginCategory        string `json:"plugin_category"`
	PluginID              string `json:"plugin_id"`
	normalizedPluginID    string

	// extended metrics fluentd >= 1.6
	// https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released
	EmitRecords                      uint64  `json:"emit_records"`
	EmitCount                        uint64  `json:"emit_count"`
	WriteCount                       uint64  `json:"write_count"`
	RollbackCount                    uint64  `json:"rollback_count"`
	SlowFlushCount                   uint64  `json:"slow_flush_count"`
	FlushTimeCount                   uint64  `json:"flush_time_count"`
	BufferStageLength                uint64  `json:"buffer_stage_length"`
	BufferStageByteSize              uint64  `json:"buffer_stage_byte_size"`
	BufferQueueByteSize              uint64  `json:"buffer_queue_byte_size"`
	BufferAvailableBufferSpaceRatios float64 `json:"buffer_available_buffer_space_ratios"`
}

func (fpm FluentdPluginMetrics) getExtended(name string) float64 {
	switch name {
	case "emit_records":
		return float64(fpm.EmitRecords)
	case "emit_count":
		return float64(fpm.EmitCount)
	case "write_count":
		return float64(fpm.WriteCount)
	case "rollback_count":
		return float64(fpm.RollbackCount)
	case "slow_flush_count":
		return float64(fpm.SlowFlushCount)
	case "flush_time_count":
		return float64(fpm.FlushTimeCount)
	case "buffer_stage_length":
		return float64(fpm.BufferStageLength)
	case "buffer_stage_byte_size":
		return float64(fpm.BufferStageByteSize)
	case "buffer_queue_byte_size":
		return float64(fpm.BufferQueueByteSize)
	case "buffer_available_buffer_space_ratios":
		return fpm.BufferAvailableBufferSpaceRatios
	}
	return 0
}

// FluentMonitorJSON monitor json
type FluentMonitorJSON struct {
	Plugins []FluentdPluginMetrics `json:"plugins"`
}

var normalizePluginIDRe = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizePluginID(in string) string {
	return normalizePluginIDRe.ReplaceAllString(in, "_")
}

func (fpm FluentdPluginMetrics) getNormalizedPluginID() string {
	if fpm.normalizedPluginID == "" {
		fpm.normalizedPluginID = normalizePluginID(fpm.PluginID)
	}
	return fpm.normalizedPluginID
}

func (f *FluentdPlugin) parseStats(body []byte) (map[string]interface{}, error) {
	var j FluentMonitorJSON
	err := json.Unmarshal(body, &j)
	f.plugins = j.Plugins

	metrics := make(map[string]interface{})
	for _, p := range f.plugins {
		if f.nonTargetPlugin(p) {
			continue
		}
		pid := p.getNormalizedPluginID()
		metrics[metricName("retry_count", pid)] = float64(p.RetryCount)
		metrics[metricName("buffer_queue_length", pid)] = float64(p.BufferQueueLength)
		metrics[metricName("buffer_total_queued_size", pid)] = float64(p.BufferTotalQueuedSize)
		for _, name := range f.extendedMetrics {
			metrics[metricName(name, pid)] = p.getExtended(name)
		}
	}
	return metrics, err
}

func (f *FluentdPlugin) nonTargetPlugin(plugin FluentdPluginMetrics) bool {
	if plugin.PluginCategory != "output" {
		return true
	}
	if f.pluginType != "" && f.pluginType != plugin.Type {
		return true
	}
	if f.pluginIDPattern != nil && !f.pluginIDPattern.MatchString(plugin.PluginID) {
		return true
	}
	return false
}

// FetchMetrics interface for mackerelplugin
func (f FluentdPlugin) FetchMetrics() (map[string]interface{}, error) {
	req, err := http.NewRequest(http.MethodGet, f.Target, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "mackerel-plugin-fluentd")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return f.parseStats(body)
}

var defaultGraphs = map[string]mp.Graphs{
	"retry_count": {
		Label: "retry count",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_queue_length": {
		Label: "queue length",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_total_queued_size": {
		Label: "buffer total queued size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

var extendedGraphs = map[string]mp.Graphs{
	"emit_records": {
		Label: "emitted records",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"emit_count": {
		Label: "emit calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"write_count": {
		Label: "write/try_write calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"rollback_count": {
		Label: "rollbacks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"slow_flush_count": {
		Label: "slow flushes",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"flush_time_count": {
		Label: "buffer flush time in msec",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"buffer_stage_length": {
		Label: "length of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_stage_byte_size": {
		Label: "bytesize of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_queue_byte_size": {
		Label: "bytesize of queued buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"buffer_available_buffer_space_ratios": {
		Label: "available space for buffer",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

// GraphDefinition interface for mackerelplugin
func (f FluentdPlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(f.Prefix)
	graphs := make(map[string]mp.Graphs, len(defaultGraphs))
	for key, g := range defaultGraphs {
		graphs[key] = mp.Graphs{
			Label:   (labelPrefix + " " + g.Label),
			Unit:    g.Unit,
			Metrics: g.Metrics,
		}
	}
	for _, name := range f.extendedMetrics {
		fullName := metricName(name)
		if g, ok := extendedGraphs[fullName]; ok {
			graphs[fullName] = mp.Graphs{
				Label:   (labelPrefix + " " + g.Label),
				Unit:    g.Unit,
				Metrics: g.Metrics,
			}
		}
	}
	return graphs
}

// Do the plugin
func Do() {
	host := flag.String("host", "localhost", "fluentd monitor_agent host")
	port := flag.String("port", "24220", "fluentd monitor_agent port")
	pluginType := flag.String("plugin-type", "", "Gets the metric that matches this plugin type")
	pluginIDPatternString := flag.String("plugin-id-pattern", "", "Gets the metric that matches this plugin id pattern")
	prefix := flag.String("metric-key-prefix", "fluentd", "Metric key prefix")
	tempFile := flag.String("tempfile", "", "Temp file name")
	extendedMetricNames := flag.String("extended_metrics", "", "extended metric names joind with ',' or 'all' (fluentd >= v1.6.0)")
	flag.Parse()

	var pluginIDPattern *regexp.Regexp
	var err error
	if *pluginIDPatternString != "" {
		pluginIDPattern, err = regexp.Compile(*pluginIDPatternString)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to exec mackerel-plugin-fluentd: invalid plugin-id-pattern: %s\n", err)
			os.Exit(1)
		}
	}

	var extendedMetrics []string
	switch *extendedMetricNames {
	case "all":
		for key := range extendedGraphs {
			extendedMetrics = append(extendedMetrics, key)
		}
	case "":
	default:
		for _, name := range strings.Split(*extendedMetricNames, ",") {
			fullName := metricName(name)
			if _, exists := extendedGraphs[fullName]; !exists {
				fmt.Fprintf(os.Stderr, "extended_metrics %s is not supported. See also https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released\n", name)
				os.Exit(1)
			}
			extendedMetrics = append(extendedMetrics, name)
		}
	}
	f := FluentdPlugin{
		Target:          fmt.Sprintf("http://%s:%s/api/plugins.json", *host, *port),
		Prefix:          *prefix,
		Tempfile:        *tempFile,
		pluginType:      *pluginType,
		pluginIDPattern: pluginIDPattern,
		extendedMetrics: extendedMetrics,
	}

	helper := mp.NewMackerelPlugin(f)

	helper.Tempfile = *tempFile
	if *tempFile == "" {
		tempFileSuffix := []string{*host, *port}
		if *pluginType != "" {
			tempFileSuffix = append(tempFileSuffix, *pluginType)
		}
		if *pluginIDPatternString != "" {
			tempFileSuffix = append(tempFileSuffix, fmt.Sprintf("%x", md5.Sum([]byte(*pluginIDPatternString))))
		}
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-fluentd-%s", strings.Join(tempFileSuffix, "-")))
	}

	helper.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"io"
	"path/filepath"
	"time"

	"github.com/google/go-github/github"
	"github.com/pkg/errors"
	"golang.org/x/sync/errgroup"
)

type GHR struct {
	GitHub GitHub

	outStream io.Writer
}

func (g *GHR) CreateRelease(ctx context.Context, req *github.RepositoryRelease, recreate bool) (*github.RepositoryRelease, error) {

	// When draft release creation is requested,
	// create it witout any check (it can).
	if *req.Draft {
		fmt.Fprintln(g.outStream, "==> Create a draft release")
		return g.GitHub.CreateRelease(ctx, req)
	}

	// Check release is exist or not.
	// If release is not found, then create a new release.
	release, err := g.GitHub.GetRelease(ctx, *req.TagName)
	if err != nil {
		if err != RelaseNotFound {
			return nil, errors.Wrap(err, "failed to get release")
		}
		Debugf("Release (with tag %s) is not found: create a new one",
			*req.TagName)

		if recreate {
			fmt.Fprintf(g.outStream,
				"WARNING: '-recreate' is specified but release (%s) is not found",
				*req.TagName)
		}

		fmt.Fprintln(g.outStream, "==> Create a new release")
		return g.GitHub.CreateRelease(ctx, req)
	}

	// recreae is not true. Then use that exiting release.
	if !recreate {
		Debugf("Release (with tag %s) exists: use exsiting one",
			*req.TagName)

		fmt.Fprintf(g.outStream, "WARNING: found release (%s). Use existing one.",
			*req.TagName)
		return release, nil
	}

	// When recreate is requested, delete exsiting release
	// and create a new release.
	fmt.Fprintln(g.outStream, "==> Recreate a release")
	if err := g.DeleteRelease(ctx, *release.ID, *req.TagName); err != nil {
		return nil, err
	}

	return g.GitHub.CreateRelease(ctx, req)
}

func (g *GHR) DeleteRelease(ctx context.Context, ID int, tag string) error {

	err := g.GitHub.DeleteRelease(ctx, ID)
	if err != nil {
		return err
	}

	err = g.GitHub.DeleteTag(ctx, tag)
	if err != nil {
		return err
	}

	return nil
}

func (g *GHR) UploadAssets(ctx context.Context, releaseID int, localAssets []string, parallel int) error {
	start := time.Now()
	defer func() {
		Debugf("UploadAssets: time: %d ms", int(time.Since(start).Seconds()*1000))
	}()

	eg, ctx := errgroup.WithContext(ctx)
	semaphore := make(chan struct{}, parallel)
	for _, localAsset := range localAssets {
		localAsset := localAsset
		eg.Go(func() error {
			semaphore <- struct{}{}
			defer func() {
				<-semaphore
			}()

			fmt.Fprintf(g.outStream, "--> Uploading: %15s\n", filepath.Base(localAsset))
			_, err := g.GitHub.UploadAsset(ctx, releaseID, localAsset)
			if err != nil {
				return errors.Wrapf(err,
					"failed to upload asset: %s", localAsset)
			}
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		return errors.Wrap(err, "one of goroutines is failed")
	}

	return nil
}

func (g *GHR) DeleteAssets(ctx context.Context, releaseID int, localAssets []string, parallel int) error {
	start := time.Now()
	defer func() {
		Debugf("DeleteAssets: time: %d ms", int(time.Since(start).Seconds()*1000))
	}()

	eg, ctx := errgroup.WithContext(ctx)

	assets, err := g.GitHub.ListAssets(ctx, releaseID)
	if err != nil {
		return errors.Wrap(err, "failed to list assets")
	}

	semaphore := make(chan struct{}, parallel)
	for _, localAsset := range localAssets {
		for _, asset := range assets {
			// https://golang.org/doc/faq#closures_and_goroutines
			localAsset, asset := localAsset, asset

			// Uploaded asset name is same as basename of local file
			if *asset.Name == filepath.Base(localAsset) {
				eg.Go(func() error {
					semaphore <- struct{}{}
					defer func() {
						<-semaphore
					}()

					fmt.Fprintf(g.outStream, "--> Deleting: %15s\n", *asset.Name)
					if err := g.GitHub.DeleteAsset(ctx, *asset.ID); err != nil {
						return errors.Wrapf(err,
							"failed to delete asset: %s", *asset.Name)
					}
					return nil
				})
			}
		}
	}

	if err := eg.Wait(); err != nil {
		return errors.Wrap(err, "one of goroutines is failed")
	}

	return nil
}
<commit_msg>Fix overlapping of DeleteRelease and CreateRelease<commit_after>package main

import (
	"context"
	"fmt"
	"io"
	"path/filepath"
	"time"

	"github.com/google/go-github/github"
	"github.com/pkg/errors"
	"golang.org/x/sync/errgroup"
)

type GHR struct {
	GitHub GitHub

	outStream io.Writer
}

func (g *GHR) CreateRelease(ctx context.Context, req *github.RepositoryRelease, recreate bool) (*github.RepositoryRelease, error) {

	// When draft release creation is requested,
	// create it witout any check (it can).
	if *req.Draft {
		fmt.Fprintln(g.outStream, "==> Create a draft release")
		return g.GitHub.CreateRelease(ctx, req)
	}

	// Check release is exist or not.
	// If release is not found, then create a new release.
	release, err := g.GitHub.GetRelease(ctx, *req.TagName)
	if err != nil {
		if err != RelaseNotFound {
			return nil, errors.Wrap(err, "failed to get release")
		}
		Debugf("Release (with tag %s) is not found: create a new one",
			*req.TagName)

		if recreate {
			fmt.Fprintf(g.outStream,
				"WARNING: '-recreate' is specified but release (%s) is not found",
				*req.TagName)
		}

		fmt.Fprintln(g.outStream, "==> Create a new release")
		return g.GitHub.CreateRelease(ctx, req)
	}

	// recreae is not true. Then use that exiting release.
	if !recreate {
		Debugf("Release (with tag %s) exists: use exsiting one",
			*req.TagName)

		fmt.Fprintf(g.outStream, "WARNING: found release (%s). Use existing one.",
			*req.TagName)
		return release, nil
	}

	// When recreate is requested, delete exsiting release
	// and create a new release.
	fmt.Fprintln(g.outStream, "==> Recreate a release")
	if err := g.DeleteRelease(ctx, *release.ID, *req.TagName); err != nil {
		return nil, err
	}

	return g.GitHub.CreateRelease(ctx, req)
}

func (g *GHR) DeleteRelease(ctx context.Context, ID int, tag string) error {

	err := g.GitHub.DeleteRelease(ctx, ID)
	if err != nil {
		return err
	}

	err = g.GitHub.DeleteTag(ctx, tag)
	if err != nil {
		return err
	}

	// This is because sometimes process of creating release on GitHub is more
	// fast than deleting tag.
	time.Sleep(5 * time.Second)

	return nil
}

func (g *GHR) UploadAssets(ctx context.Context, releaseID int, localAssets []string, parallel int) error {
	start := time.Now()
	defer func() {
		Debugf("UploadAssets: time: %d ms", int(time.Since(start).Seconds()*1000))
	}()

	eg, ctx := errgroup.WithContext(ctx)
	semaphore := make(chan struct{}, parallel)
	for _, localAsset := range localAssets {
		localAsset := localAsset
		eg.Go(func() error {
			semaphore <- struct{}{}
			defer func() {
				<-semaphore
			}()

			fmt.Fprintf(g.outStream, "--> Uploading: %15s\n", filepath.Base(localAsset))
			_, err := g.GitHub.UploadAsset(ctx, releaseID, localAsset)
			if err != nil {
				return errors.Wrapf(err,
					"failed to upload asset: %s", localAsset)
			}
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		return errors.Wrap(err, "one of goroutines is failed")
	}

	return nil
}

func (g *GHR) DeleteAssets(ctx context.Context, releaseID int, localAssets []string, parallel int) error {
	start := time.Now()
	defer func() {
		Debugf("DeleteAssets: time: %d ms", int(time.Since(start).Seconds()*1000))
	}()

	eg, ctx := errgroup.WithContext(ctx)

	assets, err := g.GitHub.ListAssets(ctx, releaseID)
	if err != nil {
		return errors.Wrap(err, "failed to list assets")
	}

	semaphore := make(chan struct{}, parallel)
	for _, localAsset := range localAssets {
		for _, asset := range assets {
			// https://golang.org/doc/faq#closures_and_goroutines
			localAsset, asset := localAsset, asset

			// Uploaded asset name is same as basename of local file
			if *asset.Name == filepath.Base(localAsset) {
				eg.Go(func() error {
					semaphore <- struct{}{}
					defer func() {
						<-semaphore
					}()

					fmt.Fprintf(g.outStream, "--> Deleting: %15s\n", *asset.Name)
					if err := g.GitHub.DeleteAsset(ctx, *asset.ID); err != nil {
						return errors.Wrapf(err,
							"failed to delete asset: %s", *asset.Name)
					}
					return nil
				})
			}
		}
	}

	if err := eg.Wait(); err != nil {
		return errors.Wrap(err, "one of goroutines is failed")
	}

	return nil
}
<|endoftext|>