<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

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

package tchannel

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/uber/tchannel-go/relay"

	"github.com/uber-go/atomic"
)

// _maxRelayTombs is the maximum number of tombs we'll accumulate in a single
// relayItems.
const _maxRelayTombs = 1e4

// _relayTombTTL is the length of time we'll keep a tomb before GC'ing it.
const _relayTombTTL = time.Second

type relayItem struct {
	*time.Timer

	stats       relay.CallStats
	remapID     uint32
	destination *Relayer
	tomb        bool
}

type relayItems struct {
	sync.RWMutex

	logger Logger
	tombs  uint64
	items  map[uint32]relayItem
}

func newRelayItems(logger Logger) *relayItems {
	return &relayItems{
		items:  make(map[uint32]relayItem),
		logger: logger,
	}
}

// Count returns the number of non-tombstone items in the relay.
func (r *relayItems) Count() int {
	r.RLock()
	n := len(r.items) - int(r.tombs)
	r.RUnlock()
	return n
}

// Get checks for a relay item by ID, returning the item and a bool indicating
// whether the item was found.
func (r *relayItems) Get(id uint32) (relayItem, bool) {
	r.RLock()
	item, ok := r.items[id]
	r.RUnlock()

	return item, ok
}

// Add adds a relay item.
func (r *relayItems) Add(id uint32, item relayItem) {
	r.Lock()
	r.items[id] = item
	r.Unlock()
}

// Delete removes a relayItem completely (without leaving a tombstone). It
// returns the deleted item, along with a bool indicating whether we completed a
// relayed call.
func (r *relayItems) Delete(id uint32) (relayItem, bool) {
	r.Lock()
	item, ok := r.items[id]
	if !ok {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Attempted to delete non-existent relay item.")
		return item, false
	}
	delete(r.items, id)
	if item.tomb {
		r.tombs--
	}
	r.Unlock()

	item.Stop()
	return item, !item.tomb
}

// Entomb sets the tomb bit on a relayItem and schedules a garbage collection. It
// returns the entombed item, along with a bool indicating whether we completed
// a relayed call.
func (r *relayItems) Entomb(id uint32, deleteAfter time.Duration) (relayItem, bool) {
	r.Lock()
	if r.tombs > _maxRelayTombs {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Too many tombstones, deleting relay item immediately.")
		return r.Delete(id)
	}
	item, ok := r.items[id]
	if !ok {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Can't find relay item to entomb.")
		return item, false
	}
	if item.tomb {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Re-entombing a tombstone.")
		return item, false
	}
	r.tombs++
	item.tomb = true
	r.items[id] = item
	r.Unlock()

	// TODO: We should be clearing these out in batches, rather than creating
	// individual timers for each item.
	time.AfterFunc(deleteAfter, func() { r.Delete(id) })
	return item, true
}

type frameType int

const (
	requestFrame  frameType = 0
	responseFrame frameType = 1
)

// A Relayer forwards frames.
type Relayer struct {
	stats relay.Stats
	hosts relay.Hosts

	// outbound is the remapping for requests that originated on this
	// connection, and are outbound towards some other connection.
	// It stores remappings for all request frames read on this connection.
	outbound *relayItems

	// inbound is the remapping for requests that originated on some other
	// connection which was directed to this connection.
	// It stores remappings for all response frames read on this connection.
	inbound *relayItems

	peers   *PeerList
	conn    *Connection
	logger  Logger
	pending atomic.Uint32
}

// NewRelayer constructs a Relayer.
func NewRelayer(ch *Channel, conn *Connection) *Relayer {
	return &Relayer{
		stats:    ch.relayStats,
		hosts:    ch.RelayHosts(),
		outbound: newRelayItems(ch.Logger().WithFields(LogField{"relay", "outbound"})),
		inbound:  newRelayItems(ch.Logger().WithFields(LogField{"relay", "inbound"})),
		peers:    ch.Peers(),
		conn:     conn,
		logger:   conn.log,
	}
}

// Hosts returns the RelayHosts guiding peer selection.
func (r *Relayer) Hosts() relay.Hosts {
	return r.hosts
}

// Relay is called for each frame that is read on the connection.
func (r *Relayer) Relay(f *Frame) error {
	if f.messageType() != messageTypeCallReq {
		return r.handleNonCallReq(f)
	}
	return r.handleCallReq(newLazyCallReq(f))
}

// Receive receives frames intended for this connection.
func (r *Relayer) Receive(f *Frame, fType frameType) {
	// If we receive a response frame, we expect to find that ID in our outbound.
	// If we receive a request frame, we expect to find that ID in our inbound.
	items := r.receiverItems(fType)
	isOriginator := fType == responseFrame

	item, ok := items.Get(f.Header.ID)
	if !ok {
		r.logger.WithFields(
			LogField{"ID", f.Header.ID},
		).Warn("Received a frame without a RelayItem.")
	} else if isOriginator {
		if msg := relayErrMsg(f); msg != "" {
			item.stats.Failed(msg)
		} else if indicatesSuccess(f) {
			item.stats.Succeeded()
		}
	}

	// TODO: Add some sort of timeout here to avoid blocking forever on a
	// stalled connection.
	r.conn.sendCh <- f
	if finishesCall(f) {
		items := r.receiverItems(fType)
		r.finishRelayItem(items, f.Header.ID)
	}
}

func (r *Relayer) canHandleNewCall() bool {
	var canHandle bool
	r.conn.withStateRLock(func() error {
		canHandle = r.conn.state == connectionActive
		if canHandle {
			r.pending.Inc()
		}
		return nil
	})
	return canHandle
}

func (r *Relayer) getDestination(f lazyCallReq, cs relay.CallStats) (*Connection, bool, error) {
	if _, ok := r.outbound.Get(f.Header.ID); ok {
		r.logger.WithFields(LogField{"id", f.Header.ID}).Warn("received duplicate callReq")
		cs.Failed("relay-" + ErrCodeProtocol.MetricsKey())
		// TODO: this is a protocol error, kill the connection.
		return nil, false, errors.New("callReq with already active ID")
	}

	// Get the destination
	hostPort := r.hosts.Get(f)
	if hostPort == "" {
		// TODO: What is the span in the error frame actually used for, and do we need it?
		r.conn.SendSystemError(f.Header.ID, nil, errUnknownGroup(f.Service()))
		cs.Failed("relay-" + ErrCodeDeclined.MetricsKey())
		return nil, false, nil
	}
	peer := r.peers.GetOrAdd(hostPort)

	// TODO: Should connections use the call timeout? Or a separate timeout?
	remoteConn, err := peer.getConnectionTimeout(f.TTL())
	if err != nil {
		r.logger.WithFields(
			ErrField(err),
			LogField{"hostPort", hostPort},
		).Warn("Failed to connect to relay host.")
		cs.Failed("relay-connection-failed")
		// TODO: Same as above, do we need span here?
		r.conn.SendSystemError(f.Header.ID, nil, NewWrappedSystemError(ErrCodeNetwork, err))
		return nil, false, nil
	}

	return remoteConn, true, nil
}

func (r *Relayer) handleCallReq(f lazyCallReq) error {
	callStats := r.stats.Begin(f)
	if !r.canHandleNewCall() {
		callStats.Failed("relay-channel-closed")
		callStats.End()
		return ErrChannelClosed
	}

	// Get a remote connection and check whether it can handle this call.
	remoteConn, ok, err := r.getDestination(f, callStats)
	if err == nil && ok {
		if !remoteConn.relay.canHandleNewCall() {
			err = NewSystemError(ErrCodeNetwork, "selected closed connection, retry")
			callStats.Failed("relay-connection-closed")
		}
	}
	if err != nil || !ok {
		// Failed to get a remote connection, or the connection is not in the right
		// state to handle this call. Since we already incremented pending on
		// the current relay, we need to decrement it.
		r.decrementPending()
		callStats.End()
		return err
	}

	destinationID := remoteConn.NextMessageID()
	ttl := f.TTL()
	// The remote side of the relay doesn't need to track stats.
	remoteConn.relay.addRelayItem(false /* isOriginator */, destinationID, f.Header.ID, r, ttl, nil)
	relayToDest := r.addRelayItem(true /* isOriginator */, f.Header.ID, destinationID, remoteConn.relay, ttl, callStats)

	f.Header.ID = destinationID
	relayToDest.destination.Receive(f.Frame, requestFrame)
	return nil
}

// Handle all frames except messageTypeCallReq.
func (r *Relayer) handleNonCallReq(f *Frame) error {
	frameType := frameTypeFor(f)

	// If we read a request frame, we need to use the outbound map to decide
	// the destination. Otherwise, we use the inbound map.
	items := r.outbound
	if frameType == responseFrame {
		items = r.inbound
	}

	item, ok := items.Get(f.Header.ID)
	if !ok {
		return errors.New("non-callReq for inactive ID")
	}
	if item.tomb {
		// Call timed out, ignore this frame. (We've already handled stats.)
		// TODO: metrics for late-arriving frames.
		return nil
	}
	originalID := f.Header.ID
	f.Header.ID = item.remapID
	item.destination.Receive(f, frameType)

	if finishesCall(f) {
		r.finishRelayItem(items, originalID)
	}
	return nil
}

// addRelayItem adds a relay item to either outbound or inbound.
func (r *Relayer) addRelayItem(isOriginator bool, id, remapID uint32, destination *Relayer, ttl time.Duration, cs relay.CallStats) relayItem {
	item := relayItem{
		stats:       cs,
		remapID:     remapID,
		destination: destination,
	}

	items := r.inbound
	if isOriginator {
		items = r.outbound
	}
	item.Timer = time.AfterFunc(ttl, func() { r.timeoutRelayItem(isOriginator, items, id) })
	items.Add(id, item)
	return item
}

func (r *Relayer) timeoutRelayItem(isOriginator bool, items *relayItems, id uint32) {
	item, ok := items.Entomb(id, _relayTombTTL)
	if !ok {
		return
	}
	if isOriginator {
		// TODO: As above. What's the span in the error frame for?
		r.conn.SendSystemError(id, nil, ErrTimeout)
		item.stats.Failed("timeout")
		item.stats.End()
	}

	r.decrementPending()
}

func (r *Relayer) finishRelayItem(items *relayItems, id uint32) {
	item, ok := items.Delete(id)
	if !ok {
		return
	}
	if item.stats != nil {
		item.stats.End()
	}
	r.decrementPending()
}

func (r *Relayer) decrementPending() {
	r.pending.Dec()
	r.conn.checkExchanges()
}

func (r *Relayer) canClose() bool {
	if r == nil {
		return true
	}
	return r.countPending() == 0
}

func (r *Relayer) countPending() uint32 {
	return r.pending.Load()
}

func (r *Relayer) receiverItems(fType frameType) *relayItems {
	if fType == requestFrame {
		return r.inbound
	}
	return r.outbound
}

func frameTypeFor(f *Frame) frameType {
	switch t := f.Header.messageType; t {
	case messageTypeCallRes, messageTypeCallResContinue, messageTypeError, messageTypePingRes:
		return responseFrame
	case messageTypeCallReq, messageTypeCallReqContinue, messageTypePingReq:
		return requestFrame
	default:
		panic(fmt.Sprintf("unsupported frame type: %v", t))
	}
}

func errUnknownGroup(group string) error {
	return NewSystemError(ErrCodeDeclined, "no peers for %q", group)
}

// relayErrMsg checks whether a relayer should mark this RPC as failed, and if
// so, what metrics key to use when reporting the failure.
func relayErrMsg(f *Frame) string {
	switch f.messageType() {
	case messageTypeError:
		return newLazyError(f).Code().MetricsKey()
	case messageTypeCallRes:
		if ok := newLazyCallRes(f).OK(); !ok {
			return "application-error"
		}
		return ""
	default:
		return ""
	}
}

// indicatesSuccess checks whether receipt of a frame, on the originating
// connection, indicates that the RPC succeeded.
func indicatesSuccess(f *Frame) bool {
	switch f.messageType() {
	case messageTypeCallRes:
		return newLazyCallRes(f).OK()
	default:
		return false
	}
}
<commit_msg>Clean up success/failure code in relay<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

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

package tchannel

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/uber/tchannel-go/relay"

	"github.com/uber-go/atomic"
)

// _maxRelayTombs is the maximum number of tombs we'll accumulate in a single
// relayItems.
const _maxRelayTombs = 1e4

// _relayTombTTL is the length of time we'll keep a tomb before GC'ing it.
const _relayTombTTL = time.Second

type relayItem struct {
	*time.Timer

	stats       relay.CallStats
	remapID     uint32
	destination *Relayer
	tomb        bool
}

type relayItems struct {
	sync.RWMutex

	logger Logger
	tombs  uint64
	items  map[uint32]relayItem
}

func newRelayItems(logger Logger) *relayItems {
	return &relayItems{
		items:  make(map[uint32]relayItem),
		logger: logger,
	}
}

// Count returns the number of non-tombstone items in the relay.
func (r *relayItems) Count() int {
	r.RLock()
	n := len(r.items) - int(r.tombs)
	r.RUnlock()
	return n
}

// Get checks for a relay item by ID, returning the item and a bool indicating
// whether the item was found.
func (r *relayItems) Get(id uint32) (relayItem, bool) {
	r.RLock()
	item, ok := r.items[id]
	r.RUnlock()

	return item, ok
}

// Add adds a relay item.
func (r *relayItems) Add(id uint32, item relayItem) {
	r.Lock()
	r.items[id] = item
	r.Unlock()
}

// Delete removes a relayItem completely (without leaving a tombstone). It
// returns the deleted item, along with a bool indicating whether we completed a
// relayed call.
func (r *relayItems) Delete(id uint32) (relayItem, bool) {
	r.Lock()
	item, ok := r.items[id]
	if !ok {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Attempted to delete non-existent relay item.")
		return item, false
	}
	delete(r.items, id)
	if item.tomb {
		r.tombs--
	}
	r.Unlock()

	item.Stop()
	return item, !item.tomb
}

// Entomb sets the tomb bit on a relayItem and schedules a garbage collection. It
// returns the entombed item, along with a bool indicating whether we completed
// a relayed call.
func (r *relayItems) Entomb(id uint32, deleteAfter time.Duration) (relayItem, bool) {
	r.Lock()
	if r.tombs > _maxRelayTombs {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Too many tombstones, deleting relay item immediately.")
		return r.Delete(id)
	}
	item, ok := r.items[id]
	if !ok {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Can't find relay item to entomb.")
		return item, false
	}
	if item.tomb {
		r.Unlock()
		r.logger.WithFields(LogField{"id", id}).Warn("Re-entombing a tombstone.")
		return item, false
	}
	r.tombs++
	item.tomb = true
	r.items[id] = item
	r.Unlock()

	// TODO: We should be clearing these out in batches, rather than creating
	// individual timers for each item.
	time.AfterFunc(deleteAfter, func() { r.Delete(id) })
	return item, true
}

type frameType int

const (
	requestFrame  frameType = 0
	responseFrame frameType = 1
)

// A Relayer forwards frames.
type Relayer struct {
	stats relay.Stats
	hosts relay.Hosts

	// outbound is the remapping for requests that originated on this
	// connection, and are outbound towards some other connection.
	// It stores remappings for all request frames read on this connection.
	outbound *relayItems

	// inbound is the remapping for requests that originated on some other
	// connection which was directed to this connection.
	// It stores remappings for all response frames read on this connection.
	inbound *relayItems

	peers   *PeerList
	conn    *Connection
	logger  Logger
	pending atomic.Uint32
}

// NewRelayer constructs a Relayer.
func NewRelayer(ch *Channel, conn *Connection) *Relayer {
	return &Relayer{
		stats:    ch.relayStats,
		hosts:    ch.RelayHosts(),
		outbound: newRelayItems(ch.Logger().WithFields(LogField{"relay", "outbound"})),
		inbound:  newRelayItems(ch.Logger().WithFields(LogField{"relay", "inbound"})),
		peers:    ch.Peers(),
		conn:     conn,
		logger:   conn.log,
	}
}

// Hosts returns the RelayHosts guiding peer selection.
func (r *Relayer) Hosts() relay.Hosts {
	return r.hosts
}

// Relay is called for each frame that is read on the connection.
func (r *Relayer) Relay(f *Frame) error {
	if f.messageType() != messageTypeCallReq {
		return r.handleNonCallReq(f)
	}
	return r.handleCallReq(newLazyCallReq(f))
}

// Receive receives frames intended for this connection.
func (r *Relayer) Receive(f *Frame, fType frameType) {
	// If we receive a response frame, we expect to find that ID in our outbound.
	// If we receive a request frame, we expect to find that ID in our inbound.
	items := r.receiverItems(fType)

	item, ok := items.Get(f.Header.ID)
	if !ok {
		r.logger.WithFields(
			LogField{"ID", f.Header.ID},
		).Warn("Received a frame without a RelayItem.")
		return
	}

	// call res frames don't include the OK bit, so we can't wait until the last
	// frame of a relayed RPC to determine if the call succeeded.
	isOriginator := fType == responseFrame
	if isOriginator {
		if succeeded, failed, failMsg := determinesCallSuccess(f); succeeded {
			item.stats.Succeeded()
		} else if failed {
			item.stats.Failed(failMsg)
		}
	}

	// TODO: Add some sort of timeout here to avoid blocking forever on a
	// stalled connection.
	r.conn.sendCh <- f
	if finishesCall(f) {
		items := r.receiverItems(fType)
		r.finishRelayItem(items, f.Header.ID)
	}
}

func (r *Relayer) canHandleNewCall() bool {
	var canHandle bool
	r.conn.withStateRLock(func() error {
		canHandle = r.conn.state == connectionActive
		if canHandle {
			r.pending.Inc()
		}
		return nil
	})
	return canHandle
}

func (r *Relayer) getDestination(f lazyCallReq, cs relay.CallStats) (*Connection, bool, error) {
	if _, ok := r.outbound.Get(f.Header.ID); ok {
		r.logger.WithFields(LogField{"id", f.Header.ID}).Warn("received duplicate callReq")
		cs.Failed("relay-" + ErrCodeProtocol.MetricsKey())
		// TODO: this is a protocol error, kill the connection.
		return nil, false, errors.New("callReq with already active ID")
	}

	// Get the destination
	hostPort := r.hosts.Get(f)
	if hostPort == "" {
		// TODO: What is the span in the error frame actually used for, and do we need it?
		r.conn.SendSystemError(f.Header.ID, nil, errUnknownGroup(f.Service()))
		cs.Failed("relay-" + ErrCodeDeclined.MetricsKey())
		return nil, false, nil
	}
	peer := r.peers.GetOrAdd(hostPort)

	// TODO: Should connections use the call timeout? Or a separate timeout?
	remoteConn, err := peer.getConnectionTimeout(f.TTL())
	if err != nil {
		r.logger.WithFields(
			ErrField(err),
			LogField{"hostPort", hostPort},
		).Warn("Failed to connect to relay host.")
		cs.Failed("relay-connection-failed")
		// TODO: Same as above, do we need span here?
		r.conn.SendSystemError(f.Header.ID, nil, NewWrappedSystemError(ErrCodeNetwork, err))
		return nil, false, nil
	}

	return remoteConn, true, nil
}

func (r *Relayer) handleCallReq(f lazyCallReq) error {
	callStats := r.stats.Begin(f)
	if !r.canHandleNewCall() {
		callStats.Failed("relay-channel-closed")
		callStats.End()
		return ErrChannelClosed
	}

	// Get a remote connection and check whether it can handle this call.
	remoteConn, ok, err := r.getDestination(f, callStats)
	if err == nil && ok {
		if !remoteConn.relay.canHandleNewCall() {
			err = NewSystemError(ErrCodeNetwork, "selected closed connection, retry")
			callStats.Failed("relay-connection-closed")
		}
	}
	if err != nil || !ok {
		// Failed to get a remote connection, or the connection is not in the right
		// state to handle this call. Since we already incremented pending on
		// the current relay, we need to decrement it.
		r.decrementPending()
		callStats.End()
		return err
	}

	destinationID := remoteConn.NextMessageID()
	ttl := f.TTL()
	// The remote side of the relay doesn't need to track stats.
	remoteConn.relay.addRelayItem(false /* isOriginator */, destinationID, f.Header.ID, r, ttl, nil)
	relayToDest := r.addRelayItem(true /* isOriginator */, f.Header.ID, destinationID, remoteConn.relay, ttl, callStats)

	f.Header.ID = destinationID
	relayToDest.destination.Receive(f.Frame, requestFrame)
	return nil
}

// Handle all frames except messageTypeCallReq.
func (r *Relayer) handleNonCallReq(f *Frame) error {
	frameType := frameTypeFor(f)

	// If we read a request frame, we need to use the outbound map to decide
	// the destination. Otherwise, we use the inbound map.
	items := r.outbound
	if frameType == responseFrame {
		items = r.inbound
	}

	item, ok := items.Get(f.Header.ID)
	if !ok {
		return errors.New("non-callReq for inactive ID")
	}
	if item.tomb {
		// Call timed out, ignore this frame. (We've already handled stats.)
		// TODO: metrics for late-arriving frames.
		return nil
	}
	originalID := f.Header.ID
	f.Header.ID = item.remapID
	item.destination.Receive(f, frameType)

	if finishesCall(f) {
		r.finishRelayItem(items, originalID)
	}
	return nil
}

// addRelayItem adds a relay item to either outbound or inbound.
func (r *Relayer) addRelayItem(isOriginator bool, id, remapID uint32, destination *Relayer, ttl time.Duration, cs relay.CallStats) relayItem {
	item := relayItem{
		stats:       cs,
		remapID:     remapID,
		destination: destination,
	}

	items := r.inbound
	if isOriginator {
		items = r.outbound
	}
	item.Timer = time.AfterFunc(ttl, func() { r.timeoutRelayItem(isOriginator, items, id) })
	items.Add(id, item)
	return item
}

func (r *Relayer) timeoutRelayItem(isOriginator bool, items *relayItems, id uint32) {
	item, ok := items.Entomb(id, _relayTombTTL)
	if !ok {
		return
	}
	if isOriginator {
		// TODO: As above. What's the span in the error frame for?
		r.conn.SendSystemError(id, nil, ErrTimeout)
		item.stats.Failed("timeout")
		item.stats.End()
	}

	r.decrementPending()
}

func (r *Relayer) finishRelayItem(items *relayItems, id uint32) {
	item, ok := items.Delete(id)
	if !ok {
		return
	}
	if item.stats != nil {
		item.stats.End()
	}
	r.decrementPending()
}

func (r *Relayer) decrementPending() {
	r.pending.Dec()
	r.conn.checkExchanges()
}

func (r *Relayer) canClose() bool {
	if r == nil {
		return true
	}
	return r.countPending() == 0
}

func (r *Relayer) countPending() uint32 {
	return r.pending.Load()
}

func (r *Relayer) receiverItems(fType frameType) *relayItems {
	if fType == requestFrame {
		return r.inbound
	}
	return r.outbound
}

func frameTypeFor(f *Frame) frameType {
	switch t := f.Header.messageType; t {
	case messageTypeCallRes, messageTypeCallResContinue, messageTypeError, messageTypePingRes:
		return responseFrame
	case messageTypeCallReq, messageTypeCallReqContinue, messageTypePingReq:
		return requestFrame
	default:
		panic(fmt.Sprintf("unsupported frame type: %v", t))
	}
}

func errUnknownGroup(group string) error {
	return NewSystemError(ErrCodeDeclined, "no peers for %q", group)
}

func determinesCallSuccess(f *Frame) (succeeded bool, failed bool, failMsg string) {
	ftype := f.messageType()

	if ftype == messageTypeError {
		// The call failed unexpectedly.
		msg := newLazyError(f).Code().MetricsKey()
		return false, true, msg
	}

	if ftype != messageTypeCallRes {
		// The call hasn't succeeded or failed yet.
		return false, false, ""
	}

	// The frame is a callRes, which could be either a success or an
	// application error.
	if ok := newLazyCallRes(f).OK(); !ok {
		return false, true, "application-error"
	}

	return true, false, ""
}
<|endoftext|>


<commit_before>package main

var resetFn CmdFn = func(c *Cmd, args []string) int {
	verbose := c.Flags.Bool("v", false, "verbose output")
	state := c.Flags.String("s", STATELABELS[NEW], "state to set")

	NEWSTATE, found := JOBSTATEIDS[*state]
	if !found {
		panic("State not found:" + *state)
	}

	store := NewStore(*root)
	journal := NewJournal(*verbose, *root+"/journal.log")

	defer journal.Close()
	defer store.Close()

	jobIDs := c.Flags.Args()

	jobList := NewJobList(store, journal)
	resetable := JobSpecific(jobIDs)

	for _, job := range jobList {
		if resetable(job) {
			job.SetState(NEWSTATE)
		}
	}
	return 0
}

func init() {
	cmd := NewCmd("reset", "show the status of commands", "tsinkf reset [-v] [taskIDs]", resetFn)
	cmdList[cmd.Name] = cmd
}
<commit_msg>Fix reset:<commit_after>package main

import (
	"fmt"
)

func contains(needle string, heystack []string) int {
	for i, possible := range heystack {
		if needle == possible {
			return i
		}
	}

	return -1
}

var resetFn CmdFn = func(c *Cmd, args []string) int {
	verbose := c.Flags.Bool("v", false, "verbose output")
	state := c.Flags.String("state", STATELABELS[NEW], "state to set")
	all := c.Flags.Bool("all", false, "operate on all jobs")

	c.Flags.Parse(args)

	NEWSTATE, found := JOBSTATEIDS[*state]
	if !found {
		panic("State not found:" + *state)
	}

	store := NewStore(*root)
	journal := NewJournal(*verbose, *root+"/journal.log")

	defer journal.Close()
	defer store.Close()

	jobIDs := c.Flags.Args()

	jobList := NewJobList(store, journal)

	if *all {

		for _, job := range jobList {
			fmt.Println(job.ToString())
		}

		fmt.Printf("Resetting all jobs to %s, type \"yes\" to confirm: ", STATELABELS[NEWSTATE])
		var input string
		_, err := fmt.Scanf("%s", &input)
		if err != nil {
			panic(err)
		}

		if input != "yes" {
			fmt.Println("ABORT")
			return 1
		}
	}

	for _, job := range jobList {
		if *all || contains(job.id, jobIDs) > -1 {
			job.SetState(NEWSTATE)
		}
	}

	return 0
}

func init() {
	cmd := NewCmd("reset", "show the status of commands", "tsinkf reset [-v] [taskIDs]", resetFn)
	cmdList[cmd.Name] = cmd
}
<|endoftext|>


<commit_before>package backoff

import "time"

// Retry the function f until it does not return error or BackOff stops.
// f is guaranteed to be run at least once.
//
// Example:
// 	operation := func() error {
// 		// An operation that may fail
// 	}
//
// 	err := backoff.Retry(operation, backoff.NewExponentialBackoff())
// 	if err != nil {
// 		// handle error
// 	}
//
// 	// operation is successfull
func Retry(f func() error, b BackOff) error {
	var err error
	var next time.Duration

	b.Reset()
	for {
		if err = f(); err == nil {
			return nil
		}

		if next = b.NextBackOff(); next == Stop {
			return err
		}

		time.Sleep(next)
	}
}
<commit_msg>update comment on Retry()<commit_after>package backoff

import "time"

// Retry the function f until it does not return error or BackOff stops.
// f is guaranteed to be run at least once.
// It is the caller's responsibility to reset b after Retry returns.
//
// Example:
// 	operation := func() error {
// 		// An operation that may fail
// 	}
//
// 	err := backoff.Retry(operation, backoff.NewExponentialBackoff())
// 	if err != nil {
// 		// handle error
// 	}
//
// 	// operation is successfull
func Retry(f func() error, b BackOff) error {
	var err error
	var next time.Duration

	b.Reset()
	for {
		if err = f(); err == nil {
			return nil
		}

		if next = b.NextBackOff(); next == Stop {
			return err
		}

		time.Sleep(next)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os/exec"
	"path"
	"math/rand"
	"net/http"
	"time"
)

type song struct {
	Path string            `json:"path"`
	Tags map[string]string `json:"tags"`
}

type river struct {
	Songs    map[string]*song `json:"songs"`
	library  string
	port     uint16
	convCmd  string
	probeCmd string
	json     []byte
}

func chooseCmd(a string, b string) (string, error) {
	if _, err := exec.LookPath(a); err != nil {
		if _, err := exec.LookPath(b); err != nil {
			return "", fmt.Errorf("could not find %s or %s executable", a, b)
		}
		return b, nil
	}
	return a, nil
}

func (s *song) readTags(container map[string]interface{}) {
	tagsRaw, ok := container["tags"]
	if !ok {
		return
	}
	for key, value := range tagsRaw.(map[string]interface{}) {
		s.Tags[key] = value.(string)
	}
}

func (r river) newSong(relPath string) (s *song, err error) {
	absPath := path.Join(r.library, relPath)
	cmd := exec.Command(r.probeCmd, "-print_format", "json", "-show_streams", absPath)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return
	}
	if err = cmd.Start(); err != nil {
		return
	}
	var streams struct {
		Streams []map[string]interface{}
	}
	if err = json.NewDecoder(stdout).Decode(&streams); err != nil {
		return
	}
	if err = cmd.Wait(); err != nil {
		return
	}
	audio := false
	for _, stream := range streams.Streams {
		if stream["codec_type"] == "audio" {
			audio = true
			break
		}
	}
	if !audio {
		return nil, fmt.Errorf("'%s' does not contain an audio stream", relPath)
	}
	cmd = exec.Command(r.probeCmd, "-print_format", "json", "-show_format", absPath)
	if stdout, err = cmd.StdoutPipe(); err != nil {
		return
	}
	if err = cmd.Start(); err != nil {
		return
	}
	var format struct {
		Format map[string]interface{}
	}
	if err = json.NewDecoder(stdout).Decode(&format); err != nil {
		return
	}
	if err = cmd.Wait(); err != nil {
		return
	}
	s = &song{Path: relPath, Tags: make(map[string]string)}
	s.readTags(format.Format)
	for _, stream := range streams.Streams {
		s.readTags(stream)
	}
	return
}

func id() string {
	letters := []byte("abcdefghijklmnopqrstuvwxyz")
	rand.Seed(time.Now().UnixNano())
	idBytes := make([]byte, 0, 8)
	for i := 0; i < cap(idBytes); i++ {
		idBytes = append(idBytes, letters[rand.Intn(len(letters))])
	}
	return string(idBytes)
}

func (r *river) readDir(relDir string) (err error) {
	absDir := path.Join(r.library, relDir)
	fis, err := ioutil.ReadDir(absDir)
	if err != nil {
		return
	}
	for _, fi := range fis {
		relPath := path.Join(relDir, fi.Name())
		if fi.IsDir() {
			if err = r.readDir(relPath); err != nil {
				return
			}
		} else {
			s, err := r.newSong(relPath)
			if err != nil {
				continue
			}
			r.Songs[id()] = s
		}
	}
	return
}

func newRiver(library string, port uint16) (r *river, err error) {
	r = &river{library: library, port: port}
	convCmd, err := chooseCmd("ffmpeg", "avconv")
	if err != nil {
		return nil, err
	}
	r.convCmd = convCmd
	probeCmd, err := chooseCmd("ffprobe", "avprobe")
	if err != nil {
		return nil, err
	}
	r.probeCmd = probeCmd
	r.Songs = make(map[string]*song)
	if err = r.readDir(""); err != nil {
		return nil, err
	}
	return
}

type songsHandler struct {
	river
}

func (sh songsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if err := json.NewEncoder(w).Encode(sh); err != nil {
		http.Error(w, "unable to encode song list", 500)
		return
	}
}

func (r river) serve() {
	http.Handle("/songs", songsHandler{r})
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", r.port), nil))
}

func main() {
	var flagLibrary = flag.String("library", "", "the music library")
	var flagPort = flag.Uint("port", 8080, "the port to listen on")
	flag.Parse()
	if *flagLibrary == "" {
		log.Fatal("no library path specified")
	}
	r, err := newRiver(*flagLibrary, uint16(*flagPort))
	if err != nil {
		log.Fatal(err)
	}
	r.serve()
}
<commit_msg>Add encoding/song access<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os/exec"
	"path"
	"math/rand"
	"net/http"
	"strings"
	"time"
)

type song struct {
	Path string            `json:"path"`
	Tags map[string]string `json:"tags"`
}

type river struct {
	Songs    map[string]*song `json:"songs"`
	library  string
	port     uint16
	convCmd  string
	probeCmd string
	json     []byte
}

func chooseCmd(a string, b string) (string, error) {
	if _, err := exec.LookPath(a); err != nil {
		if _, err := exec.LookPath(b); err != nil {
			return "", fmt.Errorf("could not find %s or %s executable", a, b)
		}
		return b, nil
	}
	return a, nil
}

func (s *song) readTags(container map[string]interface{}) {
	tagsRaw, ok := container["tags"]
	if !ok {
		return
	}
	for key, value := range tagsRaw.(map[string]interface{}) {
		s.Tags[key] = value.(string)
	}
}

func (r river) newSong(relPath string) (s *song, err error) {
	absPath := path.Join(r.library, relPath)
	cmd := exec.Command(r.probeCmd, "-print_format", "json", "-show_streams", absPath)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return
	}
	if err = cmd.Start(); err != nil {
		return
	}
	var streams struct {
		Streams []map[string]interface{}
	}
	if err = json.NewDecoder(stdout).Decode(&streams); err != nil {
		return
	}
	if err = cmd.Wait(); err != nil {
		return
	}
	audio := false
	for _, stream := range streams.Streams {
		if stream["codec_type"] == "audio" {
			audio = true
			break
		}
	}
	if !audio {
		return nil, fmt.Errorf("'%s' does not contain an audio stream", relPath)
	}
	cmd = exec.Command(r.probeCmd, "-print_format", "json", "-show_format", absPath)
	if stdout, err = cmd.StdoutPipe(); err != nil {
		return
	}
	if err = cmd.Start(); err != nil {
		return
	}
	var format struct {
		Format map[string]interface{}
	}
	if err = json.NewDecoder(stdout).Decode(&format); err != nil {
		return
	}
	if err = cmd.Wait(); err != nil {
		return
	}
	s = &song{Path: relPath, Tags: make(map[string]string)}
	s.readTags(format.Format)
	for _, stream := range streams.Streams {
		s.readTags(stream)
	}
	return
}

func id() string {
	letters := []byte("abcdefghijklmnopqrstuvwxyz")
	rand.Seed(time.Now().UnixNano())
	idBytes := make([]byte, 0, 8)
	for i := 0; i < cap(idBytes); i++ {
		idBytes = append(idBytes, letters[rand.Intn(len(letters))])
	}
	return string(idBytes)
}

func (r *river) readDir(relDir string) (err error) {
	absDir := path.Join(r.library, relDir)
	fis, err := ioutil.ReadDir(absDir)
	if err != nil {
		return
	}
	for _, fi := range fis {
		relPath := path.Join(relDir, fi.Name())
		if fi.IsDir() {
			if err = r.readDir(relPath); err != nil {
				return
			}
		} else {
			s, err := r.newSong(relPath)
			if err != nil {
				continue
			}
			r.Songs[id()] = s
		}
	}
	return
}

func newRiver(library string, port uint16) (r *river, err error) {
	r = &river{library: library, port: port}
	convCmd, err := chooseCmd("ffmpeg", "avconv")
	if err != nil {
		return nil, err
	}
	r.convCmd = convCmd
	probeCmd, err := chooseCmd("ffprobe", "avprobe")
	if err != nil {
		return nil, err
	}
	r.probeCmd = probeCmd
	r.Songs = make(map[string]*song)
	if err = r.readDir(""); err != nil {
		return nil, err
	}
	return
}

type songsHandler struct {
	river
}

func (songsh songsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if err := json.NewEncoder(w).Encode(songsh); err != nil {
		http.Error(w, "unable to encode song list", 500)
		return
	}
}

type songHandler struct {
	river
}

func (songh songHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	base := path.Base(r.URL.Path)
	ext := path.Ext(base)
	song, ok := songh.Songs[strings.TrimSuffix(base, ext)]
	if !ok {
		http.Error(w, "file not found", 404)
		return
	}
	var codec string
	var qFlag string
	var quality string
	var format string
	switch ext {
	case ".opus":
		codec = "opus"
		qFlag = "-compression_level"
		quality = "10"
		format = "opus"
		break
	case ".mp3":
		codec = "libmp3lame"
		qFlag = "-q"
		quality = "0"
		format = "mp3"
		break
	default:
		http.Error(w, "unsupported file extension", 403)
		return
	}
	cmd := exec.Command(songh.convCmd, "-i", path.Join(songh.library, song.Path), "-c", codec, qFlag, quality, "-f", format, "-")
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		http.Error(w, "unable to pipe output from encoder", 500)
		return
	}
	if err = cmd.Start(); err != nil {
		http.Error(w, "unable to start encoding file", 500)
		return
	}
	if _, err = io.Copy(w, stdout); err != nil {
		http.Error(w, "unable to stream file", 500)
		return
	}
	if err = cmd.Wait(); err != nil {
		http.Error(w, "error while encoding file", 500)
		return
	}
}

func (r river) serve() {
	http.Handle("/songs", songsHandler{r})
	http.Handle("/songs/", songHandler{r})
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", r.port), nil))
}

func main() {
	var flagLibrary = flag.String("library", "", "the music library")
	var flagPort = flag.Uint("port", 8080, "the port to listen on")
	flag.Parse()
	if *flagLibrary == "" {
		log.Fatal("no library path specified")
	}
	r, err := newRiver(*flagLibrary, uint16(*flagPort))
	if err != nil {
		log.Fatal(err)
	}
	r.serve()
}
<|endoftext|>


<commit_before>package birect_test

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"testing"

	"github.com/marcuswestin/go-birect"
	"github.com/marcuswestin/go-errs"
)

type ResponseError struct {
	RetryIn int
}

func TestApplicationError(t *testing.T) {
	var err error
	server, client := setupServerClient()

	type AuthenticateParams struct{ Secret string }
	type AuthenticateResponse struct{ SessionToken string }
	server.HandleJSONReq("Authenticate", func(req *birect.JSONReq) (res interface{}, err error) {
		var params AuthenticateParams
		req.ParseParams(&params)

		if params.Secret != "foobarcat" {
			err = birect.NewError(errs.Info{"Secret": params.Secret}, "Wrong secret:", params.Secret)
			return
		}

		res = AuthenticateResponse{"asd71lgd1892d1"}
		return
	})

	var res AuthenticateResponse
	err = client.SendJSONReq("Authenticate", AuthenticateParams{"badsecret"}, &res)
	assert(t, err != nil)
	assert(t, err.Error() == "Wrong secret: badsecret")
	assert(t, res.SessionToken == "")
	err = client.SendJSONReq("Authenticate", AuthenticateParams{"foobarcat"}, &res)
	assert(t, err == nil)
	assert(t, res.SessionToken == "asd71lgd1892d1")
}

func TestDefaultErrorMessage(t *testing.T) {
	var err error
	server, client := setupServerClient()
	server.HandleJSONReq("TestDefaultErrorMessage", func(req *birect.JSONReq) (res interface{}, err error) {
		return nil, errors.New("Internal, secret error message")
	})
	err = client.SendJSONReq("TestDefaultErrorMessage", nil, nil)
	assert(t, err.Error() == birect.DefaultPublicErrorMessage)
}

// Misc utils
/////////////
var lastPort = 25001

func setupServerClient() (*birect.Server, *birect.Client) {
	lastPort += 1
	go http.ListenAndServe(fmt.Sprintf("localhost:%d", lastPort), nil)
	server := birect.UpgradeRequests(fmt.Sprintf("/birect/upgrade/%d", lastPort))
	client, err := birect.Connect(fmt.Sprintf("http://localhost:%d/birect/upgrade/%d", lastPort, lastPort))
	if err != nil {
		panic(err)
	}
	return server, client
}

func assert(t *testing.T, ok bool, msg ...interface{}) {
	if !ok {
		// t.Fatal("assert failed", msg)
		log.Panic(msg...)
	}
}
<commit_msg>Try another range of ports for tests on circleci<commit_after>package birect_test

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"testing"

	"github.com/marcuswestin/go-birect"
	"github.com/marcuswestin/go-errs"
)

type ResponseError struct {
	RetryIn int
}

func TestApplicationError(t *testing.T) {
	var err error
	server, client := setupServerClient()

	type AuthenticateParams struct{ Secret string }
	type AuthenticateResponse struct{ SessionToken string }
	server.HandleJSONReq("Authenticate", func(req *birect.JSONReq) (res interface{}, err error) {
		var params AuthenticateParams
		req.ParseParams(&params)

		if params.Secret != "foobarcat" {
			err = birect.NewError(errs.Info{"Secret": params.Secret}, "Wrong secret:", params.Secret)
			return
		}

		res = AuthenticateResponse{"asd71lgd1892d1"}
		return
	})

	var res AuthenticateResponse
	err = client.SendJSONReq("Authenticate", AuthenticateParams{"badsecret"}, &res)
	assert(t, err != nil)
	assert(t, err.Error() == "Wrong secret: badsecret")
	assert(t, res.SessionToken == "")
	err = client.SendJSONReq("Authenticate", AuthenticateParams{"foobarcat"}, &res)
	assert(t, err == nil)
	assert(t, res.SessionToken == "asd71lgd1892d1")
}

func TestDefaultErrorMessage(t *testing.T) {
	var err error
	server, client := setupServerClient()
	server.HandleJSONReq("TestDefaultErrorMessage", func(req *birect.JSONReq) (res interface{}, err error) {
		return nil, errors.New("Internal, secret error message")
	})
	err = client.SendJSONReq("TestDefaultErrorMessage", nil, nil)
	assert(t, err.Error() == birect.DefaultPublicErrorMessage)
}

// Misc utils
/////////////
var lastPort = 18000

func setupServerClient() (*birect.Server, *birect.Client) {
	lastPort += 1
	go http.ListenAndServe(fmt.Sprintf("localhost:%d", lastPort), nil)
	server := birect.UpgradeRequests(fmt.Sprintf("/birect/upgrade/%d", lastPort))
	client, err := birect.Connect(fmt.Sprintf("http://localhost:%d/birect/upgrade/%d", lastPort, lastPort))
	if err != nil {
		panic(err)
	}
	return server, client
}

func assert(t *testing.T, ok bool, msg ...interface{}) {
	if !ok {
		// t.Fatal("assert failed", msg)
		log.Panic(msg...)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2011 Google Inc.

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

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

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

package main

import (
	"errors"
	"flag"
	"fmt"
	"strings"
	"time"

	"camlistore.org/pkg/client"
	"camlistore.org/pkg/cmdmain"
	"camlistore.org/pkg/schema"
)

type permanodeCmd struct {
	title   string
	tag     string
	key     string // else random
	sigTime string
}

func init() {
	cmdmain.RegisterCommand("permanode", func(flags *flag.FlagSet) cmdmain.CommandRunner {
		cmd := new(permanodeCmd)
		flags.StringVar(&cmd.title, "title", "", "Optional 'title' attribute to set on new permanode")
		flags.StringVar(&cmd.tag, "tag", "", "Optional tag(s) to set on new permanode; comma separated.")
		flags.StringVar(&cmd.key, "key", "", "Optional key to create deterministic ('planned') permanodes. Must also use --sigtime.")
		flags.StringVar(&cmd.sigTime, "sigtime", "", "Optional time to put in the OpenPGP signature packet instead of the current time. Required when producing a deterministic permanode (with --key). In format YYYY-MM-DD HH:MM:SS")
		return cmd
	})
}

func (c *permanodeCmd) Describe() string {
	return "Create and upload a permanode."
}

func (c *permanodeCmd) Usage() {
	cmdmain.Errorf("Usage: camput [globalopts] permanode [permanodeopts]\n")
}

func (c *permanodeCmd) Examples() []string {
	return []string{
		"                               (create a new permanode)",
		`-name="Some Name" -tag=foo,bar (with attributes added)`,
	}
}

func (c *permanodeCmd) RunCommand(args []string) error {
	if len(args) > 0 {
		return errors.New("Permanode command doesn't take any additional arguments")
	}

	var (
		permaNode *client.PutResult
		err       error
		up        = getUploader()
	)
	if (c.key != "") != (c.sigTime != "") {
		return errors.New("Both --key and --sigtime must be used to produce deterministic permanodes.")
	}
	if c.key == "" {
		// Normal case, with a random permanode.
		permaNode, err = up.UploadNewPermanode()
	} else {
		const format = "2006-01-02 15:04:05"
		sigTime, err := time.Parse(format, c.sigTime)
		if err != nil {
			return fmt.Errorf("Error parsing time %q; expecting time of form %q", c.sigTime, format)
		}
		permaNode, err = up.UploadPlannedPermanode(c.key, sigTime)
	}
	if handleResult("permanode", permaNode, err) != nil {
		return err
	}

	if c.title != "" {
		put, err := up.UploadAndSignBlob(schema.NewSetAttributeClaim(permaNode.BlobRef, "title", c.title))
		handleResult("claim-permanode-title", put, err)
	}
	if c.tag != "" {
		tags := strings.Split(c.tag, ",")
		m := schema.NewSetAttributeClaim(permaNode.BlobRef, "tag", tags[0])
		for _, tag := range tags {
			m = schema.NewAddAttributeClaim(permaNode.BlobRef, "tag", tag)
			put, err := up.UploadAndSignBlob(m)
			handleResult("claim-permanode-tag", put, err)
		}
	}
	return nil
}
<commit_msg>camput: s/name/title in usage string<commit_after>/*
Copyright 2011 Google Inc.

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

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

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

package main

import (
	"errors"
	"flag"
	"fmt"
	"strings"
	"time"

	"camlistore.org/pkg/client"
	"camlistore.org/pkg/cmdmain"
	"camlistore.org/pkg/schema"
)

type permanodeCmd struct {
	title   string
	tag     string
	key     string // else random
	sigTime string
}

func init() {
	cmdmain.RegisterCommand("permanode", func(flags *flag.FlagSet) cmdmain.CommandRunner {
		cmd := new(permanodeCmd)
		flags.StringVar(&cmd.title, "title", "", "Optional 'title' attribute to set on new permanode")
		flags.StringVar(&cmd.tag, "tag", "", "Optional tag(s) to set on new permanode; comma separated.")
		flags.StringVar(&cmd.key, "key", "", "Optional key to create deterministic ('planned') permanodes. Must also use --sigtime.")
		flags.StringVar(&cmd.sigTime, "sigtime", "", "Optional time to put in the OpenPGP signature packet instead of the current time. Required when producing a deterministic permanode (with --key). In format YYYY-MM-DD HH:MM:SS")
		return cmd
	})
}

func (c *permanodeCmd) Describe() string {
	return "Create and upload a permanode."
}

func (c *permanodeCmd) Usage() {
	cmdmain.Errorf("Usage: camput [globalopts] permanode [permanodeopts]\n")
}

func (c *permanodeCmd) Examples() []string {
	return []string{
		"                               (create a new permanode)",
		`-title="Some Title" -tag=foo,bar (with attributes added)`,
	}
}

func (c *permanodeCmd) RunCommand(args []string) error {
	if len(args) > 0 {
		return errors.New("Permanode command doesn't take any additional arguments")
	}

	var (
		permaNode *client.PutResult
		err       error
		up        = getUploader()
	)
	if (c.key != "") != (c.sigTime != "") {
		return errors.New("Both --key and --sigtime must be used to produce deterministic permanodes.")
	}
	if c.key == "" {
		// Normal case, with a random permanode.
		permaNode, err = up.UploadNewPermanode()
	} else {
		const format = "2006-01-02 15:04:05"
		sigTime, err := time.Parse(format, c.sigTime)
		if err != nil {
			return fmt.Errorf("Error parsing time %q; expecting time of form %q", c.sigTime, format)
		}
		permaNode, err = up.UploadPlannedPermanode(c.key, sigTime)
	}
	if handleResult("permanode", permaNode, err) != nil {
		return err
	}

	if c.title != "" {
		put, err := up.UploadAndSignBlob(schema.NewSetAttributeClaim(permaNode.BlobRef, "title", c.title))
		handleResult("claim-permanode-title", put, err)
	}
	if c.tag != "" {
		tags := strings.Split(c.tag, ",")
		m := schema.NewSetAttributeClaim(permaNode.BlobRef, "tag", tags[0])
		for _, tag := range tags {
			m = schema.NewAddAttributeClaim(permaNode.BlobRef, "tag", tag)
			put, err := up.UploadAndSignBlob(m)
			handleResult("claim-permanode-tag", put, err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"runtime"
	"time"

	"code.cloudfoundry.org/cc-uploader/ccclient"
	"code.cloudfoundry.org/cc-uploader/config"
	"code.cloudfoundry.org/cc-uploader/handlers"
	"code.cloudfoundry.org/cfhttp"
	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/consuladapter"
	"code.cloudfoundry.org/debugserver"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagerflags"
	"code.cloudfoundry.org/locket"
	"github.com/cloudfoundry/dropsonde"
	"github.com/hashicorp/consul/api"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
)

var configPath = flag.String(
	"configPath",
	"",
	"path to config",
)

const (
	ccUploadDialTimeout         = 10 * time.Second
	ccUploadKeepAlive           = 30 * time.Second
	ccUploadTLSHandshakeTimeout = 10 * time.Second
	dropsondeOrigin             = "cc_uploader"
	communicationTimeout        = 30 * time.Second
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	uploaderConfig, err := config.NewUploaderConfig(*configPath)
	if err != nil {
		panic(err.Error())
	}

	logger, reconfigurableSink := lagerflags.NewFromConfig("cc-uploader", uploaderConfig.LagerConfig)

	cfhttp.Initialize(communicationTimeout)

	initializeDropsonde(logger, uploaderConfig)
	consulClient, err := consuladapter.NewClientFromUrl(uploaderConfig.ConsulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	registrationRunner := initializeRegistrationRunner(logger, consulClient, uploaderConfig.ListenAddress, clock.NewClock())

	members := grouper.Members{
		{"cc-uploader", initializeServer(logger, uploaderConfig)},
		{"registration-runner", registrationRunner},
	}

	if uploaderConfig.DebugServerConfig.DebugAddress != "" {
		members = append(grouper.Members{
			{"debug-server", debugserver.Runner(uploaderConfig.DebugServerConfig.DebugAddress, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))
	logger.Info("ready")

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func initializeDropsonde(logger lager.Logger, uploaderConfig config.UploaderConfig) {
	dropsondeDestination := fmt.Sprint("localhost:", uploaderConfig.DropsondePort)
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func initializeTlsConfig(uploaderConfig config.UploaderConfig) *tls.Config {
	cert, err := tls.LoadX509KeyPair(uploaderConfig.CCClientCert, uploaderConfig.CCClientKey)
	if err != nil {
		log.Fatalln("Unable to load cert", err)
	}

	clientCACert, err := ioutil.ReadFile(uploaderConfig.CCCACert)
	if err != nil {
		log.Fatal("Unable to open cert", err)
	}

	clientCertPool := x509.NewCertPool()
	clientCertPool.AppendCertsFromPEM(clientCACert)
	return &tls.Config{
		InsecureSkipVerify: false,
		Certificates:       []tls.Certificate{cert},
		RootCAs:            clientCertPool,
	}
}

func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig) ifrit.Runner {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   ccUploadDialTimeout,
			KeepAlive: ccUploadKeepAlive,
		}).Dial,
		TLSClientConfig:     initializeTlsConfig(uploaderConfig),
		TLSHandshakeTimeout: ccUploadTLSHandshakeTimeout,
	}

	pollerHttpClient := cfhttp.NewClient()
	pollerHttpClient.Transport = transport

	// NewUploader takes two http.Clients, one TLS, one not?)
	uploader := ccclient.NewUploader(logger, &http.Client{Transport: transport})
	poller := ccclient.NewPoller(logger, pollerHttpClient, time.Duration(uploaderConfig.CCJobPollingInterval))

	ccUploaderHandler, err := handlers.New(uploader, poller, logger)
	if err != nil {
		logger.Error("router-building-failed", err)
		os.Exit(1)
	}

	return http_server.New(uploaderConfig.ListenAddress, ccUploaderHandler)
}

func initializeRegistrationRunner(logger lager.Logger, consulClient consuladapter.Client, listenAddress string, clock clock.Clock) ifrit.Runner {
	_, portString, err := net.SplitHostPort(listenAddress)
	if err != nil {
		logger.Fatal("failed-invalid-listen-address", err)
	}
	portNum, err := net.LookupPort("tcp", portString)
	if err != nil {
		logger.Fatal("failed-invalid-listen-port", err)
	}

	registration := &api.AgentServiceRegistration{
		Name: "cc-uploader",
		Port: portNum,
		Check: &api.AgentServiceCheck{
			TTL: "20s",
		},
	}

	return locket.NewRegistrationRunner(logger, registration, consulClient, locket.RetryInterval, clock)
}
<commit_msg>Consolidate http clients for uploading and polling<commit_after>package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"runtime"
	"time"

	"code.cloudfoundry.org/cc-uploader/ccclient"
	"code.cloudfoundry.org/cc-uploader/config"
	"code.cloudfoundry.org/cc-uploader/handlers"
	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/consuladapter"
	"code.cloudfoundry.org/debugserver"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagerflags"
	"code.cloudfoundry.org/locket"
	"github.com/cloudfoundry/dropsonde"
	"github.com/hashicorp/consul/api"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
)

var configPath = flag.String(
	"configPath",
	"",
	"path to config",
)

const (
	ccUploadDialTimeout         = 10 * time.Second
	ccUploadKeepAlive           = 30 * time.Second
	ccUploadTLSHandshakeTimeout = 10 * time.Second
	dropsondeOrigin             = "cc_uploader"
	communicationTimeout        = 30 * time.Second
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	uploaderConfig, err := config.NewUploaderConfig(*configPath)
	if err != nil {
		panic(err.Error())
	}

	logger, reconfigurableSink := lagerflags.NewFromConfig("cc-uploader", uploaderConfig.LagerConfig)

	initializeDropsonde(logger, uploaderConfig)
	consulClient, err := consuladapter.NewClientFromUrl(uploaderConfig.ConsulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	registrationRunner := initializeRegistrationRunner(logger, consulClient, uploaderConfig.ListenAddress, clock.NewClock())

	members := grouper.Members{
		{"cc-uploader", initializeServer(logger, uploaderConfig)},
		{"registration-runner", registrationRunner},
	}

	if uploaderConfig.DebugServerConfig.DebugAddress != "" {
		members = append(grouper.Members{
			{"debug-server", debugserver.Runner(uploaderConfig.DebugServerConfig.DebugAddress, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))
	logger.Info("ready")

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func initializeDropsonde(logger lager.Logger, uploaderConfig config.UploaderConfig) {
	dropsondeDestination := fmt.Sprint("localhost:", uploaderConfig.DropsondePort)
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func initializeTlsConfig(uploaderConfig config.UploaderConfig) *tls.Config {
	cert, err := tls.LoadX509KeyPair(uploaderConfig.CCClientCert, uploaderConfig.CCClientKey)
	if err != nil {
		log.Fatalln("Unable to load cert", err)
	}

	clientCACert, err := ioutil.ReadFile(uploaderConfig.CCCACert)
	if err != nil {
		log.Fatal("Unable to open cert", err)
	}

	clientCertPool := x509.NewCertPool()
	clientCertPool.AppendCertsFromPEM(clientCACert)
	return &tls.Config{
		InsecureSkipVerify: false,
		Certificates:       []tls.Certificate{cert},
		RootCAs:            clientCertPool,
	}
}

func initializeServer(logger lager.Logger, uploaderConfig config.UploaderConfig) ifrit.Runner {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial: (&net.Dialer{
			Timeout:   ccUploadDialTimeout,
			KeepAlive: ccUploadKeepAlive,
		}).Dial,
		TLSClientConfig:     initializeTlsConfig(uploaderConfig),
		TLSHandshakeTimeout: ccUploadTLSHandshakeTimeout,
	}

	// NewUploader takes two http.Clients, one TLS, one not?)
	uploader := ccclient.NewUploader(logger, &http.Client{Transport: transport})
	poller := ccclient.NewPoller(logger, &http.Client{Transport: transport}, time.Duration(uploaderConfig.CCJobPollingInterval))

	ccUploaderHandler, err := handlers.New(uploader, poller, logger)
	if err != nil {
		logger.Error("router-building-failed", err)
		os.Exit(1)
	}

	return http_server.New(uploaderConfig.ListenAddress, ccUploaderHandler)
}

func initializeRegistrationRunner(logger lager.Logger, consulClient consuladapter.Client, listenAddress string, clock clock.Clock) ifrit.Runner {
	_, portString, err := net.SplitHostPort(listenAddress)
	if err != nil {
		logger.Fatal("failed-invalid-listen-address", err)
	}
	portNum, err := net.LookupPort("tcp", portString)
	if err != nil {
		logger.Fatal("failed-invalid-listen-port", err)
	}

	registration := &api.AgentServiceRegistration{
		Name: "cc-uploader",
		Port: portNum,
		Check: &api.AgentServiceCheck{
			TTL: "20s",
		},
	}

	return locket.NewRegistrationRunner(logger, registration, consulClient, locket.RetryInterval, clock)
}
<|endoftext|>


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

package s3

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"sort"
	"strings"
	"time"

	"github.com/alecthomas/template"
	"github.com/blang/semver"
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/keybase/release/update"
	"github.com/keybase/release/version"
)

type Section struct {
	Header   string
	Releases []Release
}

type Release struct {
	Name       string
	Key        s3.Key
	URL        string
	Version    string
	DateString string
	Date       time.Time
	Commit     string
}

type ByRelease []Release

func (s ByRelease) Len() int {
	return len(s)
}

func (s ByRelease) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s ByRelease) Less(i, j int) bool {
	// Reverse date order
	return s[j].Date.Before(s[i].Date)
}

type Client struct {
	s3 *s3.S3
}

func NewClient() (*Client, error) {
	auth, err := aws.EnvAuth()
	if err != nil {
		return nil, err
	}
	s3 := s3.New(auth, aws.USEast)
	return &Client{s3: s3}, nil
}

func convertEastern(t time.Time) time.Time {
	locationNewYork, err := time.LoadLocation("America/New_York")
	if err != nil {
		log.Printf("Couldn't load location: %s", err)
	}
	return t.In(locationNewYork)
}

func loadReleases(keys []s3.Key, bucketName string, prefix string, suffix string, truncate int) []Release {
	var releases []Release
	for _, k := range keys {
		if strings.HasSuffix(k.Key, suffix) {
			urlString, name := urlStringForKey(k, bucketName, prefix)
			version, date, commit, err := version.Parse(name)
			if err != nil {
				log.Printf("Couldn't get version from name: %s\n", name)
			}
			date = convertEastern(date)
			releases = append(releases,
				Release{
					Name:       name,
					Key:        k,
					URL:        urlString,
					Version:    version,
					Date:       date,
					DateString: date.Format("Mon Jan _2 15:04:05 MST 2006"),
					Commit:     commit,
				})
		}
	}
	// TODO: Should also sanity check that version sort is same as time sort
	// otherwise something got messed up
	sort.Sort(ByRelease(releases))
	if truncate > 0 && len(releases) > truncate {
		releases = releases[0:truncate]
	}
	return releases
}

func WriteHTML(path string, bucketName string, prefixes string, suffix string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.s3.Bucket(bucketName)
	if bucket == nil {
		return fmt.Errorf("Bucket %s not found", bucketName)
	}

	var sections []Section
	for _, prefix := range strings.Split(prefixes, ",") {
		resp, err := bucket.List(prefix, "", "", 0)
		if err != nil {
			return err
		}

		releases := loadReleases(resp.Contents, bucketName, prefix, suffix, 20)
		if len(releases) > 0 {
			log.Printf("Found %d release(s) at %s\n", len(releases), prefix)
			for _, release := range releases {
				log.Printf(" %s %s %s\n", release.Name, release.Version, release.DateString)
			}
		}
		sections = append(sections, Section{
			Header:   prefix,
			Releases: releases,
		})
	}

	return WriteHTMLForLinks(path, bucketName, sections)
}

var htmlTemplate = `
<!doctype html>

  <title>{{ .Title }}</title>

	{{ range $index, $sec := .Sections }}

{{ $sec.Header }}

		{{ range $index2, $rel := $sec.Releases }}

{{ $rel.Name }} {{ $rel.Version }} {{ $rel.Date }} {{ $rel.Commit }}

		{{ end }}

	{{ end }}

`

func WriteHTMLForLinks(path string, title string, sections []Section) error {
	vars := map[string]interface{}{
		"Title":    title,
		"Sections": sections,
	}

	t, err := template.New("t").Parse(htmlTemplate)
	if err != nil {
		return err
	}

	if path != "" {
		var data bytes.Buffer
		err = t.Execute(&data, vars)
		if err != nil {
			return err
		}
		err := makeParentDirs(path)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(path, data.Bytes(), 0644)
	}
	return nil
}

type Platform struct {
	Name          string
	Prefix        string
	PrefixSupport string
	Suffix        string
	LatestName    string
}

func CopyLatest(bucketName string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	return client.CopyLatest(bucketName)
}

func Platforms() []Platform {
	return []Platform{
		Platform{Name: "darwin", Prefix: "darwin/", PrefixSupport: "darwin-support/", LatestName: "Keybase.dmg"},
		Platform{Name: "deb", Prefix: "linux_binaries/deb/", Suffix: "_amd64.deb", LatestName: "keybase_amd64.deb"},
		Platform{Name: "rpm", Prefix: "linux_binaries/rpm/", Suffix: ".x86_64.rpm", LatestName: "keybase_amd64.rpm"},
		Platform{Name: "windows", Prefix: "windows/", Suffix: ".386.exe", LatestName: "keybase_setup_386.exe"},
	}
}

func FindPlatform(name string) *Platform {
	platforms := Platforms()
	for _, p := range platforms {
		if p.Name == name {
			return &p
		}
	}
	return nil
}

func (p *Platform) FindRelease(bucket s3.Bucket, f func(r Release) bool) (*Release, error) {
	resp, err := bucket.List(p.Prefix, "", "", 0)
	if err != nil {
		return nil, err
	}
	releases := loadReleases(resp.Contents, bucket.Name, p.Prefix, p.Suffix, 0)
	for _, release := range releases {
		k := release.Key
		if !strings.HasSuffix(k.Key, p.Suffix) {
			continue
		}
		if f(release) {
			return &release, nil
		}
	}
	return nil, nil
}

func (c *Client) CopyLatest(bucketName string) error {
	bucket := c.s3.Bucket(bucketName)

	platforms := Platforms()

	for _, platform := range platforms {
		release, err := platform.FindRelease(*bucket, func(r Release) bool { return true })
		if err != nil {
			return err
		}
		if release == nil {
			continue
		}
		k := release.Key
		url, _ := urlStringForKey(k, bucketName, platform.Prefix)
		// Instead of linking, we're making copies. S3 linking has some issues.
		// headers := map[string][]string{
		// 	"x-amz-website-redirect-location": []string{url},
		// }
		//err = bucket.PutHeader(name, []byte{}, headers, s3.PublicRead)
		//_, err = bucket.PutCopy(platform.LatestName, s3.PublicRead, s3.CopyOptions{}, url)
		_, err = putCopy(bucket, platform.LatestName, url)
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *Client) CurrentUpdate(bucketName string, channel string, platformName string, env string) (currentUpdate *keybase1.Update, err error) {
	bucket := c.s3.Bucket(bucketName)

	data, err := bucket.Get(updateJSONName(channel, platformName, env))
	if err != nil {
		return
	}
	currentUpdate, err = update.DecodeJSON(data)
	return
}

func PromoteRelease(bucketName string, delay time.Duration, hourEastern int, channel string, platform string, env string) (*Release, error) {
	client, err := NewClient()
	if err != nil {
		return nil, err
	}
	return client.PromoteRelease(bucketName, delay, hourEastern, channel, platform, env)
}

func updateJSONName(channel string, platformName string, env string) string {
	if channel == "" {
		return fmt.Sprintf("update-%s-%s.json", platformName, env)
	}
	return fmt.Sprintf("update-%s-%s-%s.json", platformName, env, channel)
}

func (c *Client) PromoteRelease(bucketName string, delay time.Duration, beforeHourEastern int, channel string, platformName string, env string) (*Release, error) {
	if channel == "" {
		log.Printf("Finding release to promote for public (%s delay, < %dam)", delay, beforeHourEastern)
	} else {
		log.Printf("Finding release to promote for %s channel (%s delay)", channel, delay)
	}
	bucket := c.s3.Bucket(bucketName)

	platform := FindPlatform(platformName)
	if platform == nil {
		return nil, fmt.Errorf("Unsupported platform")
	}
	release, err := platform.FindRelease(*bucket, func(r Release) bool {
		if delay != 0 && time.Since(r.Date) < delay {
			return false
		}
		hour, _, _ := r.Date.Clock()
		if beforeHourEastern != 0 && hour >= beforeHourEastern {
			return false
		}
		return true
	})
	if err != nil {
		return nil, err
	}

	if release == nil {
		log.Printf("No matching release found")
		return nil, nil
	}
	log.Printf("Found release %s (%s), %s", release.Name, time.Since(release.Date), release.Version)

	currentUpdate, err := c.CurrentUpdate(bucketName, channel, platformName, env)
	if err != nil {
		log.Printf("Error looking for current update: %s", err)
	}
	if currentUpdate != nil {
		log.Printf("Found update: %s", currentUpdate.Version)
		currentVer, err := semver.Make(currentUpdate.Version)
		if err != nil {
			return nil, err
		}
		releaseVer, err := semver.Make(release.Version)
		if err != nil {
			return nil, err
		}

		if releaseVer.Equals(currentVer) {
			log.Printf("Release unchanged")
			return nil, nil
		} else if releaseVer.LT(currentVer) {
			log.Printf("Release older than current update")
			return nil, nil
		}
	}

	jsonName := updateJSONName(channel, platformName, env)
	jsonURL := urlString(bucketName, platform.PrefixSupport, fmt.Sprintf("update-%s-%s-%s.json", platformName, env, release.Version))
	//_, err = bucket.PutCopy(jsonName, s3.PublicRead, s3.CopyOptions{}, jsonURL)
	_, err = putCopy(bucket, jsonName, jsonURL)
	if err != nil {
		return nil, err
	}
	return release, nil
}

// Temporary until amz/go PR is live
func putCopy(b *s3.Bucket, destPath string, sourceURL string) (res *s3.CopyObjectResult, err error) {
	for attempt := b.S3.AttemptStrategy.Start(); attempt.Next(); {
		log.Printf("PutCopying %s to %s\n", sourceURL, destPath)
		res, err = b.PutCopy(destPath, s3.PublicRead, s3.CopyOptions{}, sourceURL)
		if err == nil {
			return
		}
	}
	return
}
<commit_msg>Copy latest from latest update instead of most recent<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package s3

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"sort"
	"strings"
	"time"

	"github.com/alecthomas/template"
	"github.com/blang/semver"
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/keybase/release/update"
	"github.com/keybase/release/version"
)

type Section struct {
	Header   string
	Releases []Release
}

type Release struct {
	Name       string
	Key        s3.Key
	URL        string
	Version    string
	DateString string
	Date       time.Time
	Commit     string
}

type ByRelease []Release

func (s ByRelease) Len() int {
	return len(s)
}

func (s ByRelease) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s ByRelease) Less(i, j int) bool {
	// Reverse date order
	return s[j].Date.Before(s[i].Date)
}

type Client struct {
	s3 *s3.S3
}

func NewClient() (*Client, error) {
	auth, err := aws.EnvAuth()
	if err != nil {
		return nil, err
	}
	s3 := s3.New(auth, aws.USEast)
	return &Client{s3: s3}, nil
}

func convertEastern(t time.Time) time.Time {
	locationNewYork, err := time.LoadLocation("America/New_York")
	if err != nil {
		log.Printf("Couldn't load location: %s", err)
	}
	return t.In(locationNewYork)
}

func loadReleases(keys []s3.Key, bucketName string, prefix string, suffix string, truncate int) []Release {
	var releases []Release
	for _, k := range keys {
		if strings.HasSuffix(k.Key, suffix) {
			urlString, name := urlStringForKey(k, bucketName, prefix)
			version, date, commit, err := version.Parse(name)
			if err != nil {
				log.Printf("Couldn't get version from name: %s\n", name)
			}
			date = convertEastern(date)
			releases = append(releases,
				Release{
					Name:       name,
					Key:        k,
					URL:        urlString,
					Version:    version,
					Date:       date,
					DateString: date.Format("Mon Jan _2 15:04:05 MST 2006"),
					Commit:     commit,
				})
		}
	}
	// TODO: Should also sanity check that version sort is same as time sort
	// otherwise something got messed up
	sort.Sort(ByRelease(releases))
	if truncate > 0 && len(releases) > truncate {
		releases = releases[0:truncate]
	}
	return releases
}

func WriteHTML(path string, bucketName string, prefixes string, suffix string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	bucket := client.s3.Bucket(bucketName)
	if bucket == nil {
		return fmt.Errorf("Bucket %s not found", bucketName)
	}

	var sections []Section
	for _, prefix := range strings.Split(prefixes, ",") {
		resp, err := bucket.List(prefix, "", "", 0)
		if err != nil {
			return err
		}

		releases := loadReleases(resp.Contents, bucketName, prefix, suffix, 20)
		if len(releases) > 0 {
			log.Printf("Found %d release(s) at %s\n", len(releases), prefix)
			for _, release := range releases {
				log.Printf(" %s %s %s\n", release.Name, release.Version, release.DateString)
			}
		}
		sections = append(sections, Section{
			Header:   prefix,
			Releases: releases,
		})
	}

	return WriteHTMLForLinks(path, bucketName, sections)
}

var htmlTemplate = `
<!doctype html>

  <title>{{ .Title }}</title>

	{{ range $index, $sec := .Sections }}

{{ $sec.Header }}

		{{ range $index2, $rel := $sec.Releases }}

{{ $rel.Name }} {{ $rel.Version }} {{ $rel.Date }} {{ $rel.Commit }}

		{{ end }}

	{{ end }}

`

func WriteHTMLForLinks(path string, title string, sections []Section) error {
	vars := map[string]interface{}{
		"Title":    title,
		"Sections": sections,
	}

	t, err := template.New("t").Parse(htmlTemplate)
	if err != nil {
		return err
	}

	if path != "" {
		var data bytes.Buffer
		err = t.Execute(&data, vars)
		if err != nil {
			return err
		}
		err := makeParentDirs(path)
		if err != nil {
			return err
		}
		return ioutil.WriteFile(path, data.Bytes(), 0644)
	}
	return nil
}

type Platform struct {
	Name          string
	Prefix        string
	PrefixSupport string
	Suffix        string
	LatestName    string
}

func CopyLatest(bucketName string) error {
	client, err := NewClient()
	if err != nil {
		return err
	}
	return client.CopyLatest(bucketName)
}

func Platforms() []Platform {
	return []Platform{
		Platform{Name: "darwin", Prefix: "darwin/", PrefixSupport: "darwin-support/", LatestName: "Keybase.dmg"},
		Platform{Name: "deb", Prefix: "linux_binaries/deb/", Suffix: "_amd64.deb", LatestName: "keybase_amd64.deb"},
		Platform{Name: "rpm", Prefix: "linux_binaries/rpm/", Suffix: ".x86_64.rpm", LatestName: "keybase_amd64.rpm"},
		Platform{Name: "windows", Prefix: "windows/", Suffix: ".386.exe", LatestName: "keybase_setup_386.exe"},
	}
}

func FindPlatform(name string) *Platform {
	platforms := Platforms()
	for _, p := range platforms {
		if p.Name == name {
			return &p
		}
	}
	return nil
}

func (p *Platform) FindRelease(bucket s3.Bucket, f func(r Release) bool) (*Release, error) {
	resp, err := bucket.List(p.Prefix, "", "", 0)
	if err != nil {
		return nil, err
	}
	releases := loadReleases(resp.Contents, bucket.Name, p.Prefix, p.Suffix, 0)
	for _, release := range releases {
		k := release.Key
		if !strings.HasSuffix(k.Key, p.Suffix) {
			continue
		}
		if f(release) {
			return &release, nil
		}
	}
	return nil, nil
}

func (c *Client) CopyLatest(bucketName string) error {
	platforms := Platforms()
	for _, platform := range platforms {
		currentUpdate, path, err := c.CurrentUpdate(bucketName, "", platform.Name, "prod")
		if err != nil || currentUpdate == nil {
			log.Printf("%s No latest for %s at %s", err, platform.Name, path)
			continue
		}

		bucket := c.s3.Bucket(bucketName)
		// Instead of linking, we're making copies. S3 linking has some issues.
		//_, err = bucket.PutCopy(platform.LatestName, s3.PublicRead, s3.CopyOptions{}, url)
		_, err = putCopy(bucket, platform.LatestName, currentUpdate.Asset.Url)
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *Client) CurrentUpdate(bucketName string, channel string, platformName string, env string) (currentUpdate *keybase1.Update, path string, err error) {
	bucket := c.s3.Bucket(bucketName)
	path = updateJSONName(channel, platformName, env)
	data, err := bucket.Get(path)
	if err != nil {
		return
	}
	currentUpdate, err = update.DecodeJSON(data)
	return
}

func PromoteRelease(bucketName string, delay time.Duration, hourEastern int, channel string, platform string, env string) (*Release, error) {
	client, err := NewClient()
	if err != nil {
		return nil, err
	}
	return client.PromoteRelease(bucketName, delay, hourEastern, channel, platform, env)
}

func updateJSONName(channel string, platformName string, env string) string {
	if channel == "" {
		return fmt.Sprintf("update-%s-%s.json", platformName, env)
	}
	return fmt.Sprintf("update-%s-%s-%s.json", platformName, env, channel)
}

func (c *Client) PromoteRelease(bucketName string, delay time.Duration, beforeHourEastern int, channel string, platformName string, env string) (*Release, error) {
	if channel == "" {
		log.Printf("Finding release to promote for public (%s delay, < %dam)", delay, beforeHourEastern)
	} else {
		log.Printf("Finding release to promote for %s channel (%s delay)", channel, delay)
	}
	bucket := c.s3.Bucket(bucketName)

	platform := FindPlatform(platformName)
	if platform == nil {
		return nil, fmt.Errorf("Unsupported platform")
	}
	release, err := platform.FindRelease(*bucket, func(r Release) bool {
		if delay != 0 && time.Since(r.Date) < delay {
			return false
		}
		hour, _, _ := r.Date.Clock()
		if beforeHourEastern != 0 && hour >= beforeHourEastern {
			return false
		}
		return true
	})
	if err != nil {
		return nil, err
	}

	if release == nil {
		log.Printf("No matching release found")
		return nil, nil
	}
	log.Printf("Found release %s (%s), %s", release.Name, time.Since(release.Date), release.Version)

	currentUpdate, _, err := c.CurrentUpdate(bucketName, channel, platformName, env)
	if err != nil {
		log.Printf("Error looking for current update: %s (%s)", err, platformName)
	}
	if currentUpdate != nil {
		log.Printf("Found update: %s", currentUpdate.Version)
		currentVer, err := semver.Make(currentUpdate.Version)
		if err != nil {
			return nil, err
		}
		releaseVer, err := semver.Make(release.Version)
		if err != nil {
			return nil, err
		}

		if releaseVer.Equals(currentVer) {
			log.Printf("Release unchanged")
			return nil, nil
		} else if releaseVer.LT(currentVer) {
			log.Printf("Release older than current update")
			return nil, nil
		}
	}

	jsonName := updateJSONName(channel, platformName, env)
	jsonURL := urlString(bucketName, platform.PrefixSupport, fmt.Sprintf("update-%s-%s-%s.json", platformName, env, release.Version))
	//_, err = bucket.PutCopy(jsonName, s3.PublicRead, s3.CopyOptions{}, jsonURL)
	_, err = putCopy(bucket, jsonName, jsonURL)
	if err != nil {
		return nil, err
	}
	return release, nil
}

// Temporary until amz/go PR is live
func putCopy(b *s3.Bucket, destPath string, sourceURL string) (res *s3.CopyObjectResult, err error) {
	for attempt := b.S3.AttemptStrategy.Start(); attempt.Next(); {
		log.Printf("PutCopying %s to %s\n", sourceURL, destPath)
		res, err = b.PutCopy(destPath, s3.PublicRead, s3.CopyOptions{}, sourceURL)
		if err == nil {
			return
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"time"

	"golang.org/x/crypto/ssh"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/restart"
	"github.com/tedsuo/ifrit/sigmon"
)

type WorkerCommand struct {
	Name string   `long:"name" description:"The name to set for the worker during registration. If not specified, the hostname will be used."`
	Tags []string `long:"tag" description:"A tag to set during registration. Can be specified multiple times."`

	HTTPProxy  URLFlag  `long:"http-proxy"  env:"http_proxy"                  description:"HTTP proxy endpoint to use for containers."`
	HTTPSProxy URLFlag  `long:"https-proxy" env:"https_proxy"                 description:"HTTPS proxy endpoint to use for containers."`
	NoProxy    []string `long:"no-proxy"    env:"no_proxy"    env-delim:","   description:"Blacklist of addresses to skip the proxy when reaching."`

	WorkDir string `long:"work-dir" required:"true" description:"Directory in which to place container data."`

	BindIP   IPFlag `long:"bind-ip"   default:"0.0.0.0" description:"IP address on which to listen for the Garden server."`
	BindPort uint16 `long:"bind-port" default:"7777"    description:"Port on which to listen for the Garden server."`

	PeerIP IPFlag `long:"peer-ip" description:"IP used to reach this worker from the ATC nodes. If omitted, the worker will be forwarded through the SSH connection to the TSA."`

	Garden GardenBackend `group:"Garden Configuration" namespace:"garden"`

	Baggageclaim baggageclaimcmd.BaggageclaimCommand `group:"Baggageclaim Configuration" namespace:"baggageclaim"`

	TSA BeaconConfig `group:"TSA Configuration" namespace:"tsa"`

	Metrics struct {
		YellerAPIKey      string `long:"yeller-api-key"     description:"Yeller API key. If specified, all errors logged will be emitted."`
		YellerEnvironment string `long:"yeller-environment" description:"Environment to tag on all Yeller events emitted."`
	} `group:"Metrics & Diagnostics"`
}

func (cmd *WorkerCommand) Execute(args []string) error {
	logger := lager.NewLogger("worker")
	logger.RegisterSink(lager.NewWriterSink(os.Stdout, lager.INFO))

	worker, gardenRunner, err := cmd.gardenRunner(logger.Session("garden"), args)
	if err != nil {
		return err
	}

	baggageclaimRunner, err := cmd.baggageclaimRunner(logger.Session("baggageclaim"))
	if err != nil {
		return err
	}

	members := grouper.Members{
		{
			Name:   "garden",
			Runner: gardenRunner,
		},
		{
			Name:   "baggageclaim",
			Runner: baggageclaimRunner,
		},
	}

	if cmd.TSA.WorkerPrivateKey != "" {
		members = append(members, grouper.Member{
			Name:   "beacon",
			Runner: cmd.beaconRunner(logger.Session("beacon"), worker),
		})
	}

	runner := sigmon.New(grouper.NewParallel(os.Interrupt, members))

	return <-ifrit.Invoke(runner).Wait()
}

func (cmd *WorkerCommand) workerName() (string, error) {
	if cmd.Name != "" {
		return cmd.Name, nil
	}

	return os.Hostname()
}

func (cmd *WorkerCommand) beaconRunner(logger lager.Logger, worker atc.Worker) ifrit.Runner {
	beacon := Beacon{
		Logger: logger,
		Config: cmd.TSA,
	}

	var beaconRunner ifrit.RunFunc
	if cmd.PeerIP != nil {
		worker.GardenAddr = fmt.Sprintf("%s:%d", cmd.PeerIP.IP(), cmd.BindPort)
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", cmd.PeerIP.IP(), cmd.Baggageclaim.BindPort)
		beaconRunner = beacon.Register
	} else {
		worker.GardenAddr = fmt.Sprintf("%s:%d", cmd.BindIP.IP(), cmd.BindPort)
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", cmd.Baggageclaim.BindIP.IP(), cmd.Baggageclaim.BindPort)
		beaconRunner = beacon.Forward
	}

	beacon.Worker = worker

	return restart.Restarter{
		Runner: beaconRunner,
		Load: func(prevRunner ifrit.Runner, prevErr error) ifrit.Runner {
			if _, ok := prevErr.(*ssh.ExitError); !ok {
				logger.Error("restarting", prevErr)
				time.Sleep(5 * time.Second)
				return beaconRunner
			}

			return nil
		},
	}
}
<commit_msg>add somehow missing import<commit_after>package main

import (
	"fmt"
	"os"
	"time"

	"golang.org/x/crypto/ssh"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc"
	"github.com/concourse/baggageclaim/baggageclaimcmd"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/restart"
	"github.com/tedsuo/ifrit/sigmon"
)

type WorkerCommand struct {
	Name string   `long:"name" description:"The name to set for the worker during registration. If not specified, the hostname will be used."`
	Tags []string `long:"tag" description:"A tag to set during registration. Can be specified multiple times."`

	HTTPProxy  URLFlag  `long:"http-proxy"  env:"http_proxy"                  description:"HTTP proxy endpoint to use for containers."`
	HTTPSProxy URLFlag  `long:"https-proxy" env:"https_proxy"                 description:"HTTPS proxy endpoint to use for containers."`
	NoProxy    []string `long:"no-proxy"    env:"no_proxy"    env-delim:","   description:"Blacklist of addresses to skip the proxy when reaching."`

	WorkDir string `long:"work-dir" required:"true" description:"Directory in which to place container data."`

	BindIP   IPFlag `long:"bind-ip"   default:"0.0.0.0" description:"IP address on which to listen for the Garden server."`
	BindPort uint16 `long:"bind-port" default:"7777"    description:"Port on which to listen for the Garden server."`

	PeerIP IPFlag `long:"peer-ip" description:"IP used to reach this worker from the ATC nodes. If omitted, the worker will be forwarded through the SSH connection to the TSA."`

	Garden GardenBackend `group:"Garden Configuration" namespace:"garden"`

	Baggageclaim baggageclaimcmd.BaggageclaimCommand `group:"Baggageclaim Configuration" namespace:"baggageclaim"`

	TSA BeaconConfig `group:"TSA Configuration" namespace:"tsa"`

	Metrics struct {
		YellerAPIKey      string `long:"yeller-api-key"     description:"Yeller API key. If specified, all errors logged will be emitted."`
		YellerEnvironment string `long:"yeller-environment" description:"Environment to tag on all Yeller events emitted."`
	} `group:"Metrics & Diagnostics"`
}

func (cmd *WorkerCommand) Execute(args []string) error {
	logger := lager.NewLogger("worker")
	logger.RegisterSink(lager.NewWriterSink(os.Stdout, lager.INFO))

	worker, gardenRunner, err := cmd.gardenRunner(logger.Session("garden"), args)
	if err != nil {
		return err
	}

	baggageclaimRunner, err := cmd.baggageclaimRunner(logger.Session("baggageclaim"))
	if err != nil {
		return err
	}

	members := grouper.Members{
		{
			Name:   "garden",
			Runner: gardenRunner,
		},
		{
			Name:   "baggageclaim",
			Runner: baggageclaimRunner,
		},
	}

	if cmd.TSA.WorkerPrivateKey != "" {
		members = append(members, grouper.Member{
			Name:   "beacon",
			Runner: cmd.beaconRunner(logger.Session("beacon"), worker),
		})
	}

	runner := sigmon.New(grouper.NewParallel(os.Interrupt, members))

	return <-ifrit.Invoke(runner).Wait()
}

func (cmd *WorkerCommand) workerName() (string, error) {
	if cmd.Name != "" {
		return cmd.Name, nil
	}

	return os.Hostname()
}

func (cmd *WorkerCommand) beaconRunner(logger lager.Logger, worker atc.Worker) ifrit.Runner {
	beacon := Beacon{
		Logger: logger,
		Config: cmd.TSA,
	}

	var beaconRunner ifrit.RunFunc
	if cmd.PeerIP != nil {
		worker.GardenAddr = fmt.Sprintf("%s:%d", cmd.PeerIP.IP(), cmd.BindPort)
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", cmd.PeerIP.IP(), cmd.Baggageclaim.BindPort)
		beaconRunner = beacon.Register
	} else {
		worker.GardenAddr = fmt.Sprintf("%s:%d", cmd.BindIP.IP(), cmd.BindPort)
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", cmd.Baggageclaim.BindIP.IP(), cmd.Baggageclaim.BindPort)
		beaconRunner = beacon.Forward
	}

	beacon.Worker = worker

	return restart.Restarter{
		Runner: beaconRunner,
		Load: func(prevRunner ifrit.Runner, prevErr error) ifrit.Runner {
			if _, ok := prevErr.(*ssh.ExitError); !ok {
				logger.Error("restarting", prevErr)
				time.Sleep(5 * time.Second)
				return beaconRunner
			}

			return nil
		},
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"syscall"
	"time"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/guardian/rundmc/dadoo"
	"github.com/eapache/go-resiliency/retrier"
	"github.com/kr/pty"
	"github.com/opencontainers/runc/libcontainer/system"
)

var uid = flag.Int("uid", 0, "uid to chown console to")
var gid = flag.Int("gid", 0, "gid to chown console to")
var rows = flag.Int("rows", 0, "rows for tty")
var cols = flag.Int("cols", 0, "cols for tty")
var tty = flag.Bool("tty", false, "tty requested")

func main() {
	os.Exit(run())
}

func run() int {
	flag.Parse()

	runtime := flag.Args()[1] // e.g. runc
	dir := flag.Args()[2]     // bundlePath for run, processPath for exec
	containerId := flag.Args()[3]

	signals := make(chan os.Signal, 100)
	signal.Notify(signals, syscall.SIGCHLD)

	fd3 := os.NewFile(3, "/proc/self/fd/3")
	logFile := fmt.Sprintf("/proc/%d/fd/4", os.Getpid())
	logFD := os.NewFile(4, "/proc/self/fd/4")
	pidFilePath := filepath.Join(dir, "pidfile")

	stdin, stdout, stderr, winsz := openPipes(dir)

	var runcStartCmd *exec.Cmd
	if *tty {
		ttySlave := setupTty(stdin, stdout, pidFilePath, winsz, garden.WindowSize{Rows: *rows, Columns: *cols})
		runcStartCmd = exec.Command(runtime, "-debug", "-log", logFile, "exec", "-d", "-tty", "-console", ttySlave.Name(), "-p", fmt.Sprintf("/proc/%d/fd/0", os.Getpid()), "-pid-file", pidFilePath, containerId)
	} else {
		runcStartCmd = exec.Command(runtime, "-debug", "-log", logFile, "exec", "-p", fmt.Sprintf("/proc/%d/fd/0", os.Getpid()), "-d", "-pid-file", pidFilePath, containerId)
		runcStartCmd.Stdin = stdin
		runcStartCmd.Stdout = stdout
		runcStartCmd.Stderr = stderr
	}

	// we need to be the subreaper so we can wait on the detached container process
	system.SetSubreaper(os.Getpid())

	if err := runcStartCmd.Start(); err != nil {
		fd3.Write([]byte{2})
		return 2
	}

	var status syscall.WaitStatus
	var rusage syscall.Rusage
	_, err := syscall.Wait4(runcStartCmd.Process.Pid, &status, 0, &rusage)
	check(err)    // Start succeeded but Wait4 failed, this can only be a programmer error
	logFD.Close() // No more logs from runc so close fd

	fd3.Write([]byte{byte(status.ExitStatus())})
	if status.ExitStatus() != 0 {
		return 3 // nothing to wait for, container didn't launch
	}

	containerPid, err := parsePid(pidFilePath)
	check(err)

	return waitForContainerToExit(dir, containerPid, signals)
}

func waitForContainerToExit(dir string, containerPid int, signals chan os.Signal) (exitCode int) {
	for range signals {
		for {
			var status syscall.WaitStatus
			var rusage syscall.Rusage
			wpid, err := syscall.Wait4(-1, &status, syscall.WNOHANG, &rusage)
			if err != nil || wpid <= 0 {
				break // wait for next SIGCHLD
			}

			if wpid == containerPid {
				exitCode = status.ExitStatus()
				if status.Signaled() {
					exitCode = 128 + int(status.Signal())
				}

				check(ioutil.WriteFile(filepath.Join(dir, "exitcode"), []byte(strconv.Itoa(exitCode)), 0700))
				return exitCode
			}
		}
	}

	panic("ran out of signals") // cant happen
}

func openPipes(dir string) (io.Reader, io.Writer, io.Writer, io.Reader) {
	stdin := openFifo(filepath.Join(dir, "stdin"), os.O_RDONLY)
	stdout := openFifo(filepath.Join(dir, "stdout"), os.O_WRONLY|os.O_APPEND)
	stderr := openFifo(filepath.Join(dir, "stderr"), os.O_WRONLY|os.O_APPEND)
	winsz := openFifo(filepath.Join(dir, "winsz"), os.O_RDWR)
	openFifo(filepath.Join(dir, "exit"), os.O_RDWR) // open just so guardian can detect it being closed when we exit

	return stdin, stdout, stderr, winsz
}

func openFifo(path string, flags int) io.ReadWriter {
	r, err := os.OpenFile(path, flags, 0600)
	if os.IsNotExist(err) {
		return nil
	}

	check(err)
	return r
}

func setupTty(stdin io.Reader, stdout io.Writer, pidFilePath string, winszFifo io.Reader, defaultWinSize garden.WindowSize) *os.File {
	m, s, err := pty.Open()
	if err != nil {
		check(err)
	}

	go io.Copy(stdout, m)

	go func() {
		io.Copy(m, stdin)
		m.Close()
	}()

	dadoo.SetWinSize(m, defaultWinSize)

	go func() {
		for {
			pid, err := readPid(pidFilePath)
			if err != nil {
				println("Timed out trying to open pidfile: ", err.Error())
				return
			}

			p, err := os.FindProcess(pid)
			check(err) // cant happen on linux

			var winSize garden.WindowSize
			if err := json.NewDecoder(winszFifo).Decode(&winSize); err != nil {
				println("invalid winsz event", err)
				continue // not much we can do here..
			}

			dadoo.SetWinSize(m, winSize)
			p.Signal(syscall.SIGWINCH)
		}
	}()

	check(s.Chown(*uid, *gid))
	return s
}

func readPid(pidFilePath string) (int, error) {
	retrier := retrier.New(retrier.ConstantBackoff(20, 500*time.Millisecond), nil)
	var (
		pid int = -1
		err error
	)
	retrier.Run(func() error {
		pid, err = parsePid(pidFilePath)
		return err
	})

	return pid, err
}

func parsePid(pidFile string) (int, error) {
	b, err := ioutil.ReadFile(pidFile)
	if err != nil {
		return -1, err
	}

	var pid int
	if _, err := fmt.Sscanf(string(b), "%d", &pid); err != nil {
		return -1, err
	}

	return pid, nil
}

func check(err error) {
	if err != nil {
		panic(err)
	}
}
<commit_msg>Wait for all output to be copied before exiting<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
	"time"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/guardian/rundmc/dadoo"
	"github.com/eapache/go-resiliency/retrier"
	"github.com/kr/pty"
	"github.com/opencontainers/runc/libcontainer/system"
)

var uid = flag.Int("uid", 0, "uid to chown console to")
var gid = flag.Int("gid", 0, "gid to chown console to")
var rows = flag.Int("rows", 0, "rows for tty")
var cols = flag.Int("cols", 0, "cols for tty")
var tty = flag.Bool("tty", false, "tty requested")

var ioWg sync.WaitGroup = sync.WaitGroup{}

func main() {
	os.Exit(run())
}

func run() int {
	flag.Parse()

	runtime := flag.Args()[1] // e.g. runc
	dir := flag.Args()[2]     // bundlePath for run, processPath for exec
	containerId := flag.Args()[3]

	signals := make(chan os.Signal, 100)
	signal.Notify(signals, syscall.SIGCHLD)

	fd3 := os.NewFile(3, "/proc/self/fd/3")
	logFile := fmt.Sprintf("/proc/%d/fd/4", os.Getpid())
	logFD := os.NewFile(4, "/proc/self/fd/4")
	pidFilePath := filepath.Join(dir, "pidfile")

	stdin, stdout, stderr, winsz := openPipes(dir)

	var runcStartCmd *exec.Cmd
	if *tty {
		ttySlave := setupTty(stdin, stdout, pidFilePath, winsz, garden.WindowSize{Rows: *rows, Columns: *cols})
		runcStartCmd = exec.Command(runtime, "-debug", "-log", logFile, "exec", "-d", "-tty", "-console", ttySlave.Name(), "-p", fmt.Sprintf("/proc/%d/fd/0", os.Getpid()), "-pid-file", pidFilePath, containerId)
	} else {
		runcStartCmd = exec.Command(runtime, "-debug", "-log", logFile, "exec", "-p", fmt.Sprintf("/proc/%d/fd/0", os.Getpid()), "-d", "-pid-file", pidFilePath, containerId)
		runcStartCmd.Stdin = stdin
		runcStartCmd.Stdout = stdout
		runcStartCmd.Stderr = stderr
	}

	// we need to be the subreaper so we can wait on the detached container process
	system.SetSubreaper(os.Getpid())

	if err := runcStartCmd.Start(); err != nil {
		fd3.Write([]byte{2})
		return 2
	}

	var status syscall.WaitStatus
	var rusage syscall.Rusage
	_, err := syscall.Wait4(runcStartCmd.Process.Pid, &status, 0, &rusage)
	check(err)    // Start succeeded but Wait4 failed, this can only be a programmer error
	logFD.Close() // No more logs from runc so close fd

	fd3.Write([]byte{byte(status.ExitStatus())})
	if status.ExitStatus() != 0 {
		return 3 // nothing to wait for, container didn't launch
	}

	containerPid, err := parsePid(pidFilePath)
	check(err)

	return waitForContainerToExit(dir, containerPid, signals)
}

func waitForContainerToExit(dir string, containerPid int, signals chan os.Signal) (exitCode int) {
	for range signals {
		for {
			var status syscall.WaitStatus
			var rusage syscall.Rusage
			wpid, err := syscall.Wait4(-1, &status, syscall.WNOHANG, &rusage)
			if err != nil || wpid <= 0 {
				break // wait for next SIGCHLD
			}

			if wpid == containerPid {
				exitCode = status.ExitStatus()
				if status.Signaled() {
					exitCode = 128 + int(status.Signal())
				}

				ioWg.Wait() // wait for full output to be collected

				check(ioutil.WriteFile(filepath.Join(dir, "exitcode"), []byte(strconv.Itoa(exitCode)), 0700))
				return exitCode
			}
		}
	}

	panic("ran out of signals") // cant happen
}

func openPipes(dir string) (io.Reader, io.Writer, io.Writer, io.Reader) {
	stdin := openFifo(filepath.Join(dir, "stdin"), os.O_RDONLY)
	stdout := openFifo(filepath.Join(dir, "stdout"), os.O_WRONLY|os.O_APPEND)
	stderr := openFifo(filepath.Join(dir, "stderr"), os.O_WRONLY|os.O_APPEND)
	winsz := openFifo(filepath.Join(dir, "winsz"), os.O_RDWR)
	openFifo(filepath.Join(dir, "exit"), os.O_RDWR) // open just so guardian can detect it being closed when we exit

	return stdin, stdout, stderr, winsz
}

func openFifo(path string, flags int) io.ReadWriter {
	r, err := os.OpenFile(path, flags, 0600)
	if os.IsNotExist(err) {
		return nil
	}

	check(err)
	return r
}

func setupTty(stdin io.Reader, stdout io.Writer, pidFilePath string, winszFifo io.Reader, defaultWinSize garden.WindowSize) *os.File {
	m, s, err := pty.Open()
	if err != nil {
		check(err)
	}

	ioWg.Add(1)
	go func() {
		defer ioWg.Done()
		io.Copy(stdout, m)
	}()

	go io.Copy(m, stdin)

	dadoo.SetWinSize(m, defaultWinSize)

	go func() {
		for {
			pid, err := readPid(pidFilePath)
			if err != nil {
				println("Timed out trying to open pidfile: ", err.Error())
				return
			}

			// free up slave fd as soon as container process is running to avoid hanging
			s.Close()

			p, err := os.FindProcess(pid)
			check(err) // cant happen on linux

			var winSize garden.WindowSize
			if err := json.NewDecoder(winszFifo).Decode(&winSize); err != nil {
				println("invalid winsz event", err)
				continue // not much we can do here..
			}

			dadoo.SetWinSize(m, winSize)
			p.Signal(syscall.SIGWINCH)
		}
	}()

	check(s.Chown(*uid, *gid))
	return s
}

func readPid(pidFilePath string) (int, error) {
	retrier := retrier.New(retrier.ConstantBackoff(20, 500*time.Millisecond), nil)
	var (
		pid int = -1
		err error
	)
	retrier.Run(func() error {
		pid, err = parsePid(pidFilePath)
		return err
	})

	return pid, err
}

func parsePid(pidFile string) (int, error) {
	b, err := ioutil.ReadFile(pidFile)
	if err != nil {
		return -1, err
	}

	var pid int
	if _, err := fmt.Sscanf(string(b), "%d", &pid); err != nil {
		return -1, err
	}

	return pid, nil
}

func check(err error) {
	if err != nil {
		panic(err)
	}
}
<|endoftext|>