<commit_before>package api

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/grafana/metrictank/api/middleware"
	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/api/response"
	"github.com/grafana/metrictank/cluster"
	"github.com/grafana/metrictank/stats"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/tinylib/msgp/msgp"
)

var NotFoundErr = errors.New("not found")

var (

	// metric api.cluster.speculative.attempts is how many peer queries resulted in speculation
	speculativeAttempts = stats.NewCounter32("api.cluster.speculative.attempts")

	// metric api.cluster.speculative.wins is how many peer queries were improved due to speculation
	speculativeWins = stats.NewCounter32("api.cluster.speculative.wins")

	// metric api.cluster.speculative.requests is how many speculative http requests made to peers
	speculativeRequests = stats.NewCounter32("api.cluster.speculative.requests")
)

func (s *Server) explainPriority(ctx *middleware.Context) {
	var data []interface{}
	for _, p := range s.prioritySetters {
		data = append(data, p.ExplainPriority())
	}
	response.Write(ctx, response.NewJson(200, data, ""))
}

func (s *Server) getNodeStatus(ctx *middleware.Context) {
	response.Write(ctx, response.NewJson(200, cluster.Manager.ThisNode(), ""))
}

func (s *Server) setNodeStatus(ctx *middleware.Context, status models.NodeStatus) {
	primary, err := strconv.ParseBool(status.Primary)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, fmt.Sprintf(
			"could not parse status to bool. %s",
			err.Error())),
		)
		return
	}
	cluster.Manager.SetPrimary(primary)
	ctx.PlainText(200, []byte("OK"))
}

func (s *Server) appStatus(ctx *middleware.Context) {
	if cluster.Manager.IsReady() {
		ctx.PlainText(200, []byte("OK"))
		return
	}

	response.Write(ctx, response.NewError(http.StatusServiceUnavailable, "node not ready"))
}

func (s *Server) getClusterStatus(ctx *middleware.Context) {
	status := models.ClusterStatus{
		ClusterName: cluster.ClusterName,
		NodeName:    cluster.Manager.ThisNode().GetName(),
		Members:     cluster.Manager.MemberList(),
	}
	response.Write(ctx, response.NewJson(200, status, ""))
}

func (s *Server) postClusterMembers(ctx *middleware.Context, req models.ClusterMembers) {
	memberNames := make(map[string]struct{})
	var toJoin []string

	for _, memberNode := range cluster.Manager.MemberList() {
		memberNames[memberNode.GetName()] = struct{}{}
	}

	for _, peerName := range req.Members {
		if _, ok := memberNames[peerName]; !ok {
			toJoin = append(toJoin, peerName)
		}
	}

	resp := models.ClusterMembersResp{
		Status:       "ok",
		MembersAdded: 0,
	}

	if len(toJoin) == 0 {
		response.Write(ctx, response.NewJson(200, resp, ""))
		return
	}

	n, err := cluster.Manager.Join(toJoin)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, fmt.Sprintf(
			"error when joining cluster members: %s", err.Error())),
		)
		return
	}
	resp.MembersAdded = n
	response.Write(ctx, response.NewJson(200, resp, ""))
}

// IndexFind returns a sequence of msgp encoded idx.Node's
func (s *Server) indexFind(ctx *middleware.Context, req models.IndexFind) {
	resp := models.NewIndexFindResp()

	for _, pattern := range req.Patterns {
		nodes, err := s.MetricIndex.Find(req.OrgId, pattern, req.From)
		if err != nil {
			response.Write(ctx, response.WrapError(err))
			return
		}
		resp.Nodes[pattern] = nodes
	}
	response.Write(ctx, response.NewMsgp(200, resp))
}

func (s *Server) indexTagDetails(ctx *middleware.Context, req models.IndexTagDetails) {
	values, err := s.MetricIndex.TagDetails(req.OrgId, req.Tag, req.Filter, req.From)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, &models.IndexTagDetailsResp{Values: values}))
}

func (s *Server) indexTags(ctx *middleware.Context, req models.IndexTags) {
	tags, err := s.MetricIndex.Tags(req.OrgId, req.Filter, req.From)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, &models.IndexTagsResp{Tags: tags}))
}

func (s *Server) indexAutoCompleteTags(ctx *middleware.Context, req models.IndexAutoCompleteTags) {
	tags, err := s.MetricIndex.FindTags(req.OrgId, req.Prefix, req.Expr, req.From, req.Limit)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, models.StringList(tags)))
}

func (s *Server) indexAutoCompleteTagValues(ctx *middleware.Context, req models.IndexAutoCompleteTagValues) {
	tags, err := s.MetricIndex.FindTagValues(req.OrgId, req.Tag, req.Prefix, req.Expr, req.From, req.Limit)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, models.StringList(tags)))
}

func (s *Server) indexTagDelSeries(ctx *middleware.Context, request models.IndexTagDelSeries) {
	deleted, err := s.MetricIndex.DeleteTagged(request.OrgId, request.Paths)
	if err != nil {
		response.Write(ctx, response.WrapErrorForTagDB(err))
		return
	}

	res := models.IndexTagDelSeriesResp{}
	res.Count = len(deleted)

	response.Write(ctx, response.NewMsgp(200, res))
}

func (s *Server) indexFindByTag(ctx *middleware.Context, req models.IndexFindByTag) {
	metrics, err := s.MetricIndex.FindByTag(req.OrgId, req.Expr, req.From)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, &models.IndexFindByTagResp{Metrics: metrics}))
}

// IndexGet returns a msgp encoded schema.MetricDefinition
func (s *Server) indexGet(ctx *middleware.Context, req models.IndexGet) {
	def, ok := s.MetricIndex.Get(req.MKey)
	if !ok {
		response.Write(ctx, response.NewError(http.StatusNotFound, "Not Found"))
		return
	}

	response.Write(ctx, response.NewMsgp(200, &def))
}

// IndexList returns msgp encoded schema.MetricDefinition's
func (s *Server) indexList(ctx *middleware.Context, req models.IndexList) {
	defs := s.MetricIndex.List(req.OrgId)
	resp := make([]msgp.Marshaler, len(defs))
	for i := range defs {
		d := defs[i]
		resp[i] = &d
	}
	response.Write(ctx, response.NewMsgpArray(200, resp))
}

func (s *Server) getData(ctx *middleware.Context, request models.GetData) {
	series, err := s.getTargetsLocal(ctx.Req.Context(), request.Requests)
	if err != nil {
		// the only errors returned are from us catching panics, so we should treat them
		// all as internalServerErrors
		log.Error(3, "HTTP getData() %s", err.Error())
		response.Write(ctx, response.WrapError(err))
		return
	}
	response.Write(ctx, response.NewMsgp(200, &models.GetDataResp{Series: series}))
}

func (s *Server) indexDelete(ctx *middleware.Context, req models.IndexDelete) {
	defs, err := s.MetricIndex.Delete(req.OrgId, req.Query)
	if err != nil {
		// errors can only be caused by bad request.
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}

	resp := models.MetricsDeleteResp{
		DeletedDefs: len(defs),
	}
	response.Write(ctx, response.NewMsgp(200, &resp))
}

type PeerResponse struct {
	peer cluster.Node
	buf  []byte
}

// peerQuery takes a request and the path to request it on, then fans it out
// across the cluster, except to the local peer. If any peer fails requests to
// other peers are aborted.
// ctx:          request context
// data:         request to be submitted
// name:         name to be used in logging & tracing
// path:         path to request on
func (s *Server) peerQuery(ctx context.Context, data cluster.Traceable, name, path string, allPeers bool) (map[string]PeerResponse, error) {
	var peers []cluster.Node
	var err error

	if allPeers {
		peers = cluster.Manager.MemberList()
	} else {
		peers, err = cluster.MembersForQuery()
		if err != nil {
			log.Error(3, "HTTP peerQuery unable to get peers, %s", err)
			return nil, err
		}
	}
	log.Debug("HTTP %s across %d instances", name, len(peers)-1)

	reqCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	responses := make(chan struct {
		data PeerResponse
		err  error
	}, 1)
	var wg sync.WaitGroup
	for _, peer := range peers {
		if peer.IsLocal() {
			continue
		}
		wg.Add(1)
		go func(peer cluster.Node) {
			defer wg.Done()
			log.Debug("HTTP Render querying %s%s", peer.GetName(), path)
			buf, err := peer.Post(reqCtx, name, path, data)
			if err != nil {
				cancel()
				log.Error(4, "HTTP Render error querying %s%s: %q", peer.GetName(), path, err)
			}
			responses <- struct {
				data PeerResponse
				err  error
			}{PeerResponse{peer, buf}, err}
		}(peer)
	}
	// wait for all list goroutines to end, then close our responses channel
	go func() {
		wg.Wait()
		close(responses)
	}()

	result := make(map[string]PeerResponse)
	for resp := range responses {
		if resp.err != nil {
			return nil, err
		}
		result[resp.data.peer.GetName()] = resp.data
	}

	return result, nil
}

// peerQuerySpeculative takes a request and the path to request it on, then fans it out
// across the cluster, except to the local peer. If any peer fails requests to
// other peers are aborted. If enough peers have been heard from (based on
// speculation-threshold configuration), and we are missing the others, try to
// speculatively query other members of the shard group.
// ctx:          request context
// data:         request to be submitted
// name:         name to be used in logging & tracing
// path:         path to request on
func (s *Server) peerQuerySpeculative(ctx context.Context, data cluster.Traceable, name, path string) (map[string]PeerResponse, error) {
	peerGroups, err := cluster.MembersForSpeculativeQuery()
	if err != nil {
		log.Error(3, "HTTP peerQuery unable to get peers, %s", err)
		return nil, err
	}
	log.Debug("HTTP %s across %d instances", name, len(peerGroups)-1)

	reqCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	originalPeers := make(map[string]struct{}, len(peerGroups))
	pendingResponses := make(map[int32]struct{}, len(peerGroups))
	receivedResponses := make(map[int32]struct{}, len(peerGroups))

	responses := make(chan struct {
		shardGroup int32
		data       PeerResponse
		err        error
	}, 1)

	askPeer := func(shardGroup int32, peer cluster.Node) {
		log.Debug("HTTP Render querying %s%s", peer.GetName(), path)
		buf, err := peer.Post(reqCtx, name, path, data)

		select {
		case <-ctx.Done():
			return
		default:
			// Not canceled, continue
		}

		if err != nil {
			cancel()
			log.Error(4, "HTTP Render error querying %s%s: %q", peer.GetName(), path, err)
		}
		responses <- struct {
			shardGroup int32
			data       PeerResponse
			err        error
		}{shardGroup, PeerResponse{peer, buf}, err}
	}

	for group, peers := range peerGroups {
		peer := peers[0]
		originalPeers[peer.GetName()] = struct{}{}
		pendingResponses[group] = struct{}{}
		go askPeer(group, peer)
	}

	result := make(map[string]PeerResponse)

	specCheckTicker := time.NewTicker(5 * time.Millisecond)

	for len(pendingResponses) > 0 {
		select {
		case resp := <-responses:
			if _, ok := receivedResponses[resp.shardGroup]; ok {
				// already received this response (possibly speculatively)
				continue
			}

			if resp.err != nil {
				return nil, err
			}

			result[resp.data.peer.GetName()] = resp.data
			receivedResponses[resp.shardGroup] = struct{}{}
			delete(pendingResponses, resp.shardGroup)
			delete(originalPeers, resp.data.peer.GetName())

		case <-specCheckTicker.C:
			// Check if it's time to speculate!
			percentReceived := 1 - (float64(len(pendingResponses)) / float64(len(peerGroups)))
			if percentReceived > speculationThreshold {
				// kick off speculative queries to other members now
				specCheckTicker.Stop()
				speculativeAttempts.Inc()
				for shardGroup := range pendingResponses {
					eligiblePeers := peerGroups[shardGroup][1:]
					for _, peer := range eligiblePeers {
						speculativeRequests.Inc()
						go askPeer(shardGroup, peer)
					}
				}
			}
		}
	}

	if len(originalPeers) > 0 {
		speculativeWins.Inc()
	}

	return result, nil
}
<commit_msg>clarify<commit_after>package api

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/grafana/metrictank/api/middleware"
	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/api/response"
	"github.com/grafana/metrictank/cluster"
	"github.com/grafana/metrictank/stats"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/tinylib/msgp/msgp"
)

var NotFoundErr = errors.New("not found")

var (

	// metric api.cluster.speculative.attempts is how many peer queries resulted in speculation
	speculativeAttempts = stats.NewCounter32("api.cluster.speculative.attempts")

	// metric api.cluster.speculative.wins is how many peer queries were improved due to speculation
	speculativeWins = stats.NewCounter32("api.cluster.speculative.wins")

	// metric api.cluster.speculative.requests is how many speculative http requests made to peers
	speculativeRequests = stats.NewCounter32("api.cluster.speculative.requests")
)

func (s *Server) explainPriority(ctx *middleware.Context) {
	var data []interface{}
	for _, p := range s.prioritySetters {
		data = append(data, p.ExplainPriority())
	}
	response.Write(ctx, response.NewJson(200, data, ""))
}

func (s *Server) getNodeStatus(ctx *middleware.Context) {
	response.Write(ctx, response.NewJson(200, cluster.Manager.ThisNode(), ""))
}

func (s *Server) setNodeStatus(ctx *middleware.Context, status models.NodeStatus) {
	primary, err := strconv.ParseBool(status.Primary)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, fmt.Sprintf(
			"could not parse status to bool. %s",
			err.Error())),
		)
		return
	}
	cluster.Manager.SetPrimary(primary)
	ctx.PlainText(200, []byte("OK"))
}

func (s *Server) appStatus(ctx *middleware.Context) {
	if cluster.Manager.IsReady() {
		ctx.PlainText(200, []byte("OK"))
		return
	}

	response.Write(ctx, response.NewError(http.StatusServiceUnavailable, "node not ready"))
}

func (s *Server) getClusterStatus(ctx *middleware.Context) {
	status := models.ClusterStatus{
		ClusterName: cluster.ClusterName,
		NodeName:    cluster.Manager.ThisNode().GetName(),
		Members:     cluster.Manager.MemberList(),
	}
	response.Write(ctx, response.NewJson(200, status, ""))
}

func (s *Server) postClusterMembers(ctx *middleware.Context, req models.ClusterMembers) {
	memberNames := make(map[string]struct{})
	var toJoin []string

	for _, memberNode := range cluster.Manager.MemberList() {
		memberNames[memberNode.GetName()] = struct{}{}
	}

	for _, peerName := range req.Members {
		if _, ok := memberNames[peerName]; !ok {
			toJoin = append(toJoin, peerName)
		}
	}

	resp := models.ClusterMembersResp{
		Status:       "ok",
		MembersAdded: 0,
	}

	if len(toJoin) == 0 {
		response.Write(ctx, response.NewJson(200, resp, ""))
		return
	}

	n, err := cluster.Manager.Join(toJoin)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, fmt.Sprintf(
			"error when joining cluster members: %s", err.Error())),
		)
		return
	}
	resp.MembersAdded = n
	response.Write(ctx, response.NewJson(200, resp, ""))
}

// IndexFind returns a sequence of msgp encoded idx.Node's
func (s *Server) indexFind(ctx *middleware.Context, req models.IndexFind) {
	resp := models.NewIndexFindResp()

	for _, pattern := range req.Patterns {
		nodes, err := s.MetricIndex.Find(req.OrgId, pattern, req.From)
		if err != nil {
			response.Write(ctx, response.WrapError(err))
			return
		}
		resp.Nodes[pattern] = nodes
	}
	response.Write(ctx, response.NewMsgp(200, resp))
}

func (s *Server) indexTagDetails(ctx *middleware.Context, req models.IndexTagDetails) {
	values, err := s.MetricIndex.TagDetails(req.OrgId, req.Tag, req.Filter, req.From)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, &models.IndexTagDetailsResp{Values: values}))
}

func (s *Server) indexTags(ctx *middleware.Context, req models.IndexTags) {
	tags, err := s.MetricIndex.Tags(req.OrgId, req.Filter, req.From)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, &models.IndexTagsResp{Tags: tags}))
}

func (s *Server) indexAutoCompleteTags(ctx *middleware.Context, req models.IndexAutoCompleteTags) {
	tags, err := s.MetricIndex.FindTags(req.OrgId, req.Prefix, req.Expr, req.From, req.Limit)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, models.StringList(tags)))
}

func (s *Server) indexAutoCompleteTagValues(ctx *middleware.Context, req models.IndexAutoCompleteTagValues) {
	tags, err := s.MetricIndex.FindTagValues(req.OrgId, req.Tag, req.Prefix, req.Expr, req.From, req.Limit)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, models.StringList(tags)))
}

func (s *Server) indexTagDelSeries(ctx *middleware.Context, request models.IndexTagDelSeries) {
	deleted, err := s.MetricIndex.DeleteTagged(request.OrgId, request.Paths)
	if err != nil {
		response.Write(ctx, response.WrapErrorForTagDB(err))
		return
	}

	res := models.IndexTagDelSeriesResp{}
	res.Count = len(deleted)

	response.Write(ctx, response.NewMsgp(200, res))
}

func (s *Server) indexFindByTag(ctx *middleware.Context, req models.IndexFindByTag) {
	metrics, err := s.MetricIndex.FindByTag(req.OrgId, req.Expr, req.From)
	if err != nil {
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}
	response.Write(ctx, response.NewMsgp(200, &models.IndexFindByTagResp{Metrics: metrics}))
}

// IndexGet returns a msgp encoded schema.MetricDefinition
func (s *Server) indexGet(ctx *middleware.Context, req models.IndexGet) {
	def, ok := s.MetricIndex.Get(req.MKey)
	if !ok {
		response.Write(ctx, response.NewError(http.StatusNotFound, "Not Found"))
		return
	}

	response.Write(ctx, response.NewMsgp(200, &def))
}

// IndexList returns msgp encoded schema.MetricDefinition's
func (s *Server) indexList(ctx *middleware.Context, req models.IndexList) {
	defs := s.MetricIndex.List(req.OrgId)
	resp := make([]msgp.Marshaler, len(defs))
	for i := range defs {
		d := defs[i]
		resp[i] = &d
	}
	response.Write(ctx, response.NewMsgpArray(200, resp))
}

func (s *Server) getData(ctx *middleware.Context, request models.GetData) {
	series, err := s.getTargetsLocal(ctx.Req.Context(), request.Requests)
	if err != nil {
		// the only errors returned are from us catching panics, so we should treat them
		// all as internalServerErrors
		log.Error(3, "HTTP getData() %s", err.Error())
		response.Write(ctx, response.WrapError(err))
		return
	}
	response.Write(ctx, response.NewMsgp(200, &models.GetDataResp{Series: series}))
}

func (s *Server) indexDelete(ctx *middleware.Context, req models.IndexDelete) {
	defs, err := s.MetricIndex.Delete(req.OrgId, req.Query)
	if err != nil {
		// errors can only be caused by bad request.
		response.Write(ctx, response.NewError(http.StatusBadRequest, err.Error()))
		return
	}

	resp := models.MetricsDeleteResp{
		DeletedDefs: len(defs),
	}
	response.Write(ctx, response.NewMsgp(200, &resp))
}

type PeerResponse struct {
	peer cluster.Node
	buf  []byte
}

// peerQuery takes a request and the path to request it on, then fans it out
// across the cluster, except to the local peer. If any peer fails requests to
// other peers are aborted.
// ctx:          request context
// data:         request to be submitted
// name:         name to be used in logging & tracing
// path:         path to request on
func (s *Server) peerQuery(ctx context.Context, data cluster.Traceable, name, path string, allPeers bool) (map[string]PeerResponse, error) {
	var peers []cluster.Node
	var err error

	if allPeers {
		peers = cluster.Manager.MemberList()
	} else {
		peers, err = cluster.MembersForQuery()
		if err != nil {
			log.Error(3, "HTTP peerQuery unable to get peers, %s", err)
			return nil, err
		}
	}
	log.Debug("HTTP %s across %d instances", name, len(peers)-1)

	reqCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	responses := make(chan struct {
		data PeerResponse
		err  error
	}, 1)
	var wg sync.WaitGroup
	for _, peer := range peers {
		if peer.IsLocal() {
			continue
		}
		wg.Add(1)
		go func(peer cluster.Node) {
			defer wg.Done()
			log.Debug("HTTP Render querying %s%s", peer.GetName(), path)
			buf, err := peer.Post(reqCtx, name, path, data)
			if err != nil {
				cancel()
				log.Error(4, "HTTP Render error querying %s%s: %q", peer.GetName(), path, err)
			}
			responses <- struct {
				data PeerResponse
				err  error
			}{PeerResponse{peer, buf}, err}
		}(peer)
	}
	// wait for all list goroutines to end, then close our responses channel
	go func() {
		wg.Wait()
		close(responses)
	}()

	result := make(map[string]PeerResponse)
	for resp := range responses {
		if resp.err != nil {
			return nil, err
		}
		result[resp.data.peer.GetName()] = resp.data
	}

	return result, nil
}

// peerQuerySpeculative takes a request and the path to request it on, then fans it out
// across the cluster, except to the local peer. If any peer fails requests to
// other peers are aborted. If enough peers have been heard from (based on
// speculation-threshold configuration), and we are missing the others, try to
// speculatively query each other member of each shard group.
// ctx:          request context
// data:         request to be submitted
// name:         name to be used in logging & tracing
// path:         path to request on
func (s *Server) peerQuerySpeculative(ctx context.Context, data cluster.Traceable, name, path string) (map[string]PeerResponse, error) {
	peerGroups, err := cluster.MembersForSpeculativeQuery()
	if err != nil {
		log.Error(3, "HTTP peerQuery unable to get peers, %s", err)
		return nil, err
	}
	log.Debug("HTTP %s across %d instances", name, len(peerGroups)-1)

	reqCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	originalPeers := make(map[string]struct{}, len(peerGroups))
	pendingResponses := make(map[int32]struct{}, len(peerGroups))
	receivedResponses := make(map[int32]struct{}, len(peerGroups))

	responses := make(chan struct {
		shardGroup int32
		data       PeerResponse
		err        error
	}, 1)

	askPeer := func(shardGroup int32, peer cluster.Node) {
		log.Debug("HTTP Render querying %s%s", peer.GetName(), path)
		buf, err := peer.Post(reqCtx, name, path, data)

		select {
		case <-ctx.Done():
			return
		default:
			// Not canceled, continue
		}

		if err != nil {
			cancel()
			log.Error(4, "HTTP Render error querying %s%s: %q", peer.GetName(), path, err)
		}
		responses <- struct {
			shardGroup int32
			data       PeerResponse
			err        error
		}{shardGroup, PeerResponse{peer, buf}, err}
	}

	for group, peers := range peerGroups {
		peer := peers[0]
		originalPeers[peer.GetName()] = struct{}{}
		pendingResponses[group] = struct{}{}
		go askPeer(group, peer)
	}

	result := make(map[string]PeerResponse)

	specCheckTicker := time.NewTicker(5 * time.Millisecond)

	for len(pendingResponses) > 0 {
		select {
		case resp := <-responses:
			if _, ok := receivedResponses[resp.shardGroup]; ok {
				// already received this response (possibly speculatively)
				continue
			}

			if resp.err != nil {
				return nil, err
			}

			result[resp.data.peer.GetName()] = resp.data
			receivedResponses[resp.shardGroup] = struct{}{}
			delete(pendingResponses, resp.shardGroup)
			delete(originalPeers, resp.data.peer.GetName())

		case <-specCheckTicker.C:
			// Check if it's time to speculate!
			percentReceived := 1 - (float64(len(pendingResponses)) / float64(len(peerGroups)))
			if percentReceived > speculationThreshold {
				// kick off speculative queries to other members now
				specCheckTicker.Stop()
				speculativeAttempts.Inc()
				for shardGroup := range pendingResponses {
					eligiblePeers := peerGroups[shardGroup][1:]
					for _, peer := range eligiblePeers {
						speculativeRequests.Inc()
						go askPeer(shardGroup, peer)
					}
				}
			}
		}
	}

	if len(originalPeers) > 0 {
		speculativeWins.Inc()
	}

	return result, nil
}
<|endoftext|>


<commit_before>package nsqd

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"math/rand"
	"os"
	"path"
	"sync"
	"sync/atomic"
	"time"
)

// diskQueueWriter implements the BackendQueue interface
// providing a filesystem backed FIFO queue
type diskQueueWriter struct {
	// 64bit atomic vars need to be first for proper alignment on 32bit platforms
	// run-time state (also persisted to disk)
	writePos           int64
	writeFileNum       int64
	readablePos        int64
	totalMsgCnt        int64
	virtualEnd         BackendOffset
	virtualReadableEnd BackendOffset
	sync.RWMutex

	// instantiation time metadata
	name            string
	dataPath        string
	maxBytesPerFile int64 // currently this cannot change once created
	minMsgSize      int32
	maxMsgSize      int32
	exitFlag        int32
	needSync        bool

	writeFile    *os.File
	bufferWriter *bufio.Writer
}

// newDiskQueue instantiates a new instance of diskQueueWriter, retrieving metadata
// from the filesystem and starting the read ahead goroutine
func newDiskQueueWriter(name string, dataPath string, maxBytesPerFile int64,
	minMsgSize int32, maxMsgSize int32,
	syncEvery int64, syncTimeout time.Duration) BackendQueueWriter {

	d := diskQueueWriter{
		name:            name,
		dataPath:        dataPath,
		maxBytesPerFile: maxBytesPerFile,
		minMsgSize:      minMsgSize,
		maxMsgSize:      maxMsgSize,
	}

	// no need to lock here, nothing else could possibly be touching this instance
	err := d.retrieveMetaData()
	if err != nil && !os.IsNotExist(err) {
		nsqLog.LogErrorf("diskqueue(%s) failed to retrieveMetaData - %s", d.name, err)
	}

	return &d
}

// Put writes a []byte to the queue
func (d *diskQueueWriter) Put(data []byte) (BackendQueueEnd, error) {
	d.Lock()
	defer d.Unlock()

	if d.exitFlag == 1 {
		return nil, errors.New("exiting")
	}
	e, werr := d.writeOne(data)
	d.needSync = true
	return e, werr
}

// Close cleans up the queue and persists metadata
func (d *diskQueueWriter) Close() error {
	return d.exit(false)
}

func (d *diskQueueWriter) Delete() error {
	return d.exit(true)
}

func (d *diskQueueWriter) exit(deleted bool) error {
	d.Lock()
	defer d.Unlock()

	d.exitFlag = 1

	if deleted {
		nsqLog.Logf("DISKQUEUE(%s): deleting", d.name)
	} else {
		nsqLog.Logf("DISKQUEUE(%s): closing", d.name)
	}

	d.sync()
	if deleted {
		return d.deleteAllFiles(deleted)
	}
	return nil
}

// Empty destructively clears out any pending data in the queue
// by fast forwarding read positions and removing intermediate files
func (d *diskQueueWriter) Empty() error {
	d.Lock()
	defer d.Unlock()

	if d.exitFlag == 1 {
		return errors.New("exiting")
	}

	nsqLog.Logf("DISKQUEUE(%s): emptying", d.name)
	return d.deleteAllFiles(false)
}

func (d *diskQueueWriter) deleteAllFiles(deleted bool) error {
	d.skipToNextRWFile()

	if deleted {
		innerErr := os.Remove(d.metaDataFileName())
		if innerErr != nil && !os.IsNotExist(innerErr) {
			nsqLog.LogErrorf("diskqueue(%s) failed to remove metadata file - %s", d.name, innerErr)
			return innerErr
		}
	}
	return nil
}

func (d *diskQueueWriter) skipToNextRWFile() error {
	if d.bufferWriter != nil {
		d.bufferWriter.Flush()
	}
	d.virtualReadableEnd = d.virtualEnd
	d.readablePos = 0
	if d.writeFile != nil {
		d.writeFile.Close()
		d.writeFile = nil
	}

	for i := int64(0); i <= d.writeFileNum; i++ {
		fn := d.fileName(i)
		innerErr := os.Remove(fn)
		if innerErr != nil && !os.IsNotExist(innerErr) {
			nsqLog.LogErrorf("diskqueue(%s) failed to remove data file - %s", d.name, innerErr)
		}
	}

	d.writeFileNum++
	d.writePos = 0
	return nil
}

func (d *diskQueueWriter) GetQueueReadEnd() BackendQueueEnd {
	d.RLock()
	defer d.RUnlock()

	if d.exitFlag == 1 {
		return nil
	}
	return d.internalGetQueueReadEnd()
}

func (d *diskQueueWriter) internalGetQueueReadEnd() BackendQueueEnd {
	e := &diskQueueEndInfo{}
	e.EndFileNum = d.writeFileNum
	e.EndPos = d.readablePos
	e.TotalMsgCnt = d.totalMsgCnt
	e.VirtualEnd = d.virtualReadableEnd
	return e
}

// writeOne performs a low level filesystem write for a single []byte
// while advancing write positions and rolling files, if necessary
func (d *diskQueueWriter) writeOne(data []byte) (*diskQueueEndInfo, error) {
	var err error

	if d.writeFile == nil {
		curFileName := d.fileName(d.writeFileNum)
		d.writeFile, err = os.OpenFile(curFileName, os.O_RDWR|os.O_CREATE, 0600)
		if err != nil {
			return nil, err
		}

		nsqLog.Logf("DISKQUEUE(%s): writeOne() opened %s", d.name, curFileName)

		if d.writePos > 0 {
			_, err = d.writeFile.Seek(d.writePos, 0)
			if err != nil {
				d.writeFile.Close()
				d.writeFile = nil
				return nil, err
			}
		}
		if d.bufferWriter == nil {
			d.bufferWriter = bufio.NewWriter(d.writeFile)
		} else {
			d.bufferWriter.Reset(d.writeFile)
		}
	}

	dataLen := int32(len(data))

	if dataLen < d.minMsgSize || dataLen > d.maxMsgSize {
		return nil, fmt.Errorf("invalid message write size (%d) maxMsgSize=%d", dataLen, d.maxMsgSize)
	}

	err = binary.Write(d.bufferWriter, binary.BigEndian, dataLen)
	if err != nil {
		return nil, err
	}

	_, err = d.bufferWriter.Write(data)
	if err != nil {
		d.writeFile.Close()
		d.writeFile = nil
		return nil, err
	}

	totalBytes := int64(4 + dataLen)
	d.writePos += totalBytes
	d.virtualEnd += BackendOffset(totalBytes)
	totalCnt := atomic.AddInt64(&d.totalMsgCnt, 1)

	endInfo := &diskQueueEndInfo{
		EndFileNum:  d.writeFileNum,
		EndPos:      d.writePos,
		TotalMsgCnt: totalCnt,
		VirtualEnd:  d.virtualEnd,
	}

	if d.writePos > d.maxBytesPerFile {
		d.writeFileNum++
		d.writePos = 0

		// sync every time we start writing to a new file
		err = d.sync()
		if err != nil {
			nsqLog.LogErrorf("diskqueue(%s) failed to sync - %s", d.name, err)
		}

		if d.bufferWriter != nil {
			d.bufferWriter.Flush()
		}
		d.virtualReadableEnd = d.virtualEnd
		d.readablePos = 0
		if d.writeFile != nil {
			d.writeFile.Close()
			d.writeFile = nil
		}
	}

	return endInfo, err
}

func (d *diskQueueWriter) Flush() error {
	d.Lock()
	defer d.Unlock()

	if d.exitFlag == 1 {
		return errors.New("exiting")
	}
	if d.needSync {
		return d.sync()
	}
	return nil
}

// sync fsyncs the current writeFile and persists metadata
func (d *diskQueueWriter) sync() error {
	if d.bufferWriter != nil {
		d.bufferWriter.Flush()
	}
	d.virtualReadableEnd = d.virtualEnd
	d.readablePos = d.writePos
	if d.writeFile != nil {
		err := d.writeFile.Sync()
		if err != nil {
			d.writeFile.Close()
			d.writeFile = nil
			return err
		}
	}

	err := d.persistMetaData()
	if err != nil {
		return err
	}

	d.needSync = false
	return nil
}

// retrieveMetaData initializes state from the filesystem
func (d *diskQueueWriter) retrieveMetaData() error {
	var f *os.File
	var err error

	fileName := d.metaDataFileName()
	f, err = os.OpenFile(fileName, os.O_RDONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()

	var totalCnt int64
	_, err = fmt.Fscanf(f, "%d\n%d,%d,%d\n",
		&totalCnt,
		&d.writeFileNum, &d.writePos, &d.virtualEnd)
	if err != nil {
		return err
	}
	d.readablePos = d.writePos
	d.virtualReadableEnd = d.virtualEnd
	atomic.StoreInt64(&d.totalMsgCnt, totalCnt)

	return nil
}

// persistMetaData atomically writes state to the filesystem
func (d *diskQueueWriter) persistMetaData() error {
	var f *os.File
	var err error

	fileName := d.metaDataFileName()
	tmpFileName := fmt.Sprintf("%s.%d.tmp", fileName, rand.Int())

	// write to tmp file
	f, err = os.OpenFile(tmpFileName, os.O_RDWR|os.O_CREATE, 0600)
	if err != nil {
		return err
	}

	_, err = fmt.Fprintf(f, "%d\n%d,%d,%d\n",
		atomic.LoadInt64(&d.totalMsgCnt),
		d.writeFileNum, d.writePos, d.virtualEnd)
	if err != nil {
		f.Close()
		return err
	}
	f.Sync()
	f.Close()

	// atomically rename
	return atomicRename(tmpFileName, fileName)
}

func (d *diskQueueWriter) metaDataFileName() string {
	return fmt.Sprintf(path.Join(d.dataPath, "%s.diskqueue.meta.writer.dat"), d.name)
}

func (d *diskQueueWriter) fileName(fileNum int64) string {
	return fmt.Sprintf(path.Join(d.dataPath, "%s.diskqueue.%06d.dat"), d.name, fileNum)
}
<commit_msg>remove unuseful code<commit_after>package nsqd

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"math/rand"
	"os"
	"path"
	"sync"
	"sync/atomic"
	"time"
)

// diskQueueWriter implements the BackendQueue interface
// providing a filesystem backed FIFO queue
type diskQueueWriter struct {
	// 64bit atomic vars need to be first for proper alignment on 32bit platforms
	// run-time state (also persisted to disk)
	writePos           int64
	writeFileNum       int64
	readablePos        int64
	totalMsgCnt        int64
	virtualEnd         BackendOffset
	virtualReadableEnd BackendOffset
	sync.RWMutex

	// instantiation time metadata
	name            string
	dataPath        string
	maxBytesPerFile int64 // currently this cannot change once created
	minMsgSize      int32
	maxMsgSize      int32
	exitFlag        int32
	needSync        bool

	writeFile    *os.File
	bufferWriter *bufio.Writer
}

// newDiskQueue instantiates a new instance of diskQueueWriter, retrieving metadata
// from the filesystem and starting the read ahead goroutine
func newDiskQueueWriter(name string, dataPath string, maxBytesPerFile int64,
	minMsgSize int32, maxMsgSize int32,
	syncEvery int64, syncTimeout time.Duration) BackendQueueWriter {

	d := diskQueueWriter{
		name:            name,
		dataPath:        dataPath,
		maxBytesPerFile: maxBytesPerFile,
		minMsgSize:      minMsgSize,
		maxMsgSize:      maxMsgSize,
	}

	// no need to lock here, nothing else could possibly be touching this instance
	err := d.retrieveMetaData()
	if err != nil && !os.IsNotExist(err) {
		nsqLog.LogErrorf("diskqueue(%s) failed to retrieveMetaData - %s", d.name, err)
	}

	return &d
}

// Put writes a []byte to the queue
func (d *diskQueueWriter) Put(data []byte) (BackendQueueEnd, error) {
	d.Lock()
	defer d.Unlock()

	if d.exitFlag == 1 {
		return nil, errors.New("exiting")
	}
	e, werr := d.writeOne(data)
	d.needSync = true
	return e, werr
}

// Close cleans up the queue and persists metadata
func (d *diskQueueWriter) Close() error {
	return d.exit(false)
}

func (d *diskQueueWriter) Delete() error {
	return d.exit(true)
}

func (d *diskQueueWriter) exit(deleted bool) error {
	d.Lock()
	defer d.Unlock()

	d.exitFlag = 1

	if deleted {
		nsqLog.Logf("DISKQUEUE(%s): deleting", d.name)
	} else {
		nsqLog.Logf("DISKQUEUE(%s): closing", d.name)
	}

	d.sync()
	if deleted {
		return d.deleteAllFiles(deleted)
	}
	return nil
}

// Empty destructively clears out any pending data in the queue
// by fast forwarding read positions and removing intermediate files
func (d *diskQueueWriter) Empty() error {
	d.Lock()
	defer d.Unlock()

	if d.exitFlag == 1 {
		return errors.New("exiting")
	}

	nsqLog.Logf("DISKQUEUE(%s): emptying", d.name)
	return d.deleteAllFiles(false)
}

func (d *diskQueueWriter) deleteAllFiles(deleted bool) error {
	d.skipToNextRWFile()

	if deleted {
		innerErr := os.Remove(d.metaDataFileName())
		if innerErr != nil && !os.IsNotExist(innerErr) {
			nsqLog.LogErrorf("diskqueue(%s) failed to remove metadata file - %s", d.name, innerErr)
			return innerErr
		}
	}
	return nil
}

func (d *diskQueueWriter) skipToNextRWFile() error {
	if d.bufferWriter != nil {
		d.bufferWriter.Flush()
	}
	d.virtualReadableEnd = d.virtualEnd
	d.readablePos = 0
	if d.writeFile != nil {
		d.writeFile.Close()
		d.writeFile = nil
	}

	for i := int64(0); i <= d.writeFileNum; i++ {
		fn := d.fileName(i)
		innerErr := os.Remove(fn)
		if innerErr != nil && !os.IsNotExist(innerErr) {
			nsqLog.LogErrorf("diskqueue(%s) failed to remove data file - %s", d.name, innerErr)
		}
	}

	d.writeFileNum++
	d.writePos = 0
	return nil
}

func (d *diskQueueWriter) GetQueueReadEnd() BackendQueueEnd {
	d.RLock()
	defer d.RUnlock()

	if d.exitFlag == 1 {
		return nil
	}
	return d.internalGetQueueReadEnd()
}

func (d *diskQueueWriter) internalGetQueueReadEnd() BackendQueueEnd {
	e := &diskQueueEndInfo{}
	e.EndFileNum = d.writeFileNum
	e.EndPos = d.readablePos
	e.TotalMsgCnt = d.totalMsgCnt
	e.VirtualEnd = d.virtualReadableEnd
	return e
}

// writeOne performs a low level filesystem write for a single []byte
// while advancing write positions and rolling files, if necessary
func (d *diskQueueWriter) writeOne(data []byte) (*diskQueueEndInfo, error) {
	var err error

	if d.writeFile == nil {
		curFileName := d.fileName(d.writeFileNum)
		d.writeFile, err = os.OpenFile(curFileName, os.O_RDWR|os.O_CREATE, 0600)
		if err != nil {
			return nil, err
		}

		nsqLog.Logf("DISKQUEUE(%s): writeOne() opened %s", d.name, curFileName)

		if d.writePos > 0 {
			_, err = d.writeFile.Seek(d.writePos, 0)
			if err != nil {
				d.writeFile.Close()
				d.writeFile = nil
				return nil, err
			}
		}
		if d.bufferWriter == nil {
			d.bufferWriter = bufio.NewWriter(d.writeFile)
		} else {
			d.bufferWriter.Reset(d.writeFile)
		}
	}

	dataLen := int32(len(data))

	if dataLen < d.minMsgSize || dataLen > d.maxMsgSize {
		return nil, fmt.Errorf("invalid message write size (%d) maxMsgSize=%d", dataLen, d.maxMsgSize)
	}

	err = binary.Write(d.bufferWriter, binary.BigEndian, dataLen)
	if err != nil {
		return nil, err
	}

	_, err = d.bufferWriter.Write(data)
	if err != nil {
		d.writeFile.Close()
		d.writeFile = nil
		return nil, err
	}

	totalBytes := int64(4 + dataLen)
	d.writePos += totalBytes
	d.virtualEnd += BackendOffset(totalBytes)
	totalCnt := atomic.AddInt64(&d.totalMsgCnt, 1)

	endInfo := &diskQueueEndInfo{
		EndFileNum:  d.writeFileNum,
		EndPos:      d.writePos,
		TotalMsgCnt: totalCnt,
		VirtualEnd:  d.virtualEnd,
	}

	if d.writePos > d.maxBytesPerFile {
		d.writeFileNum++
		d.writePos = 0

		// sync every time we start writing to a new file
		err = d.sync()
		if err != nil {
			nsqLog.LogErrorf("diskqueue(%s) failed to sync - %s", d.name, err)
		}

		d.readablePos = 0
		if d.writeFile != nil {
			d.writeFile.Close()
			d.writeFile = nil
		}
	}

	return endInfo, err
}

func (d *diskQueueWriter) Flush() error {
	d.Lock()
	defer d.Unlock()

	if d.exitFlag == 1 {
		return errors.New("exiting")
	}
	if d.needSync {
		return d.sync()
	}
	return nil
}

// sync fsyncs the current writeFile and persists metadata
func (d *diskQueueWriter) sync() error {
	if d.bufferWriter != nil {
		d.bufferWriter.Flush()
	}
	d.virtualReadableEnd = d.virtualEnd
	d.readablePos = d.writePos
	if d.writeFile != nil {
		err := d.writeFile.Sync()
		if err != nil {
			d.writeFile.Close()
			d.writeFile = nil
			return err
		}
	}

	err := d.persistMetaData()
	if err != nil {
		return err
	}

	d.needSync = false
	return nil
}

// retrieveMetaData initializes state from the filesystem
func (d *diskQueueWriter) retrieveMetaData() error {
	var f *os.File
	var err error

	fileName := d.metaDataFileName()
	f, err = os.OpenFile(fileName, os.O_RDONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()

	var totalCnt int64
	_, err = fmt.Fscanf(f, "%d\n%d,%d,%d\n",
		&totalCnt,
		&d.writeFileNum, &d.writePos, &d.virtualEnd)
	if err != nil {
		return err
	}
	d.readablePos = d.writePos
	d.virtualReadableEnd = d.virtualEnd
	atomic.StoreInt64(&d.totalMsgCnt, totalCnt)

	return nil
}

// persistMetaData atomically writes state to the filesystem
func (d *diskQueueWriter) persistMetaData() error {
	var f *os.File
	var err error

	fileName := d.metaDataFileName()
	tmpFileName := fmt.Sprintf("%s.%d.tmp", fileName, rand.Int())

	// write to tmp file
	f, err = os.OpenFile(tmpFileName, os.O_RDWR|os.O_CREATE, 0600)
	if err != nil {
		return err
	}

	_, err = fmt.Fprintf(f, "%d\n%d,%d,%d\n",
		atomic.LoadInt64(&d.totalMsgCnt),
		d.writeFileNum, d.writePos, d.virtualEnd)
	if err != nil {
		f.Close()
		return err
	}
	f.Sync()
	f.Close()

	// atomically rename
	return atomicRename(tmpFileName, fileName)
}

func (d *diskQueueWriter) metaDataFileName() string {
	return fmt.Sprintf(path.Join(d.dataPath, "%s.diskqueue.meta.writer.dat"), d.name)
}

func (d *diskQueueWriter) fileName(fileNum int64) string {
	return fmt.Sprintf(path.Join(d.dataPath, "%s.diskqueue.%06d.dat"), d.name, fileNum)
}
<|endoftext|>


<commit_before>/***** BEGIN LICENSE BLOCK *****
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
# ***** END LICENSE BLOCK *****/

package s3splitfile

import (
	"fmt"
	. "github.com/mozilla-services/heka/pipeline"
	"regexp"
	"io/ioutil"
	"encoding/json"
)

type PublishAttempt struct {
	Name string
	AttemptsRemaining uint32
}

// Encapsulates the directory-splitting schema
type Schema struct {
	Fields []string
	FieldIndices map[string]int
	Dims map[string]DimensionChecker
}

// Determine whether a given value is acceptable for a given field, and if not
// return a default value instead.
func (s *Schema) GetValue(field string, value string) (rvalue string, err error) {
	checker, ok := s.Dims[field]
	if !ok {
		return value, fmt.Errorf("No such field: '%s'", field)
	}
	if checker.IsAllowed(value) {
		return value, nil
	} else {
		return "OTHER", nil
	}
}

// Extract all dimensions from the given pack.
func (s *Schema) GetDimensions(pack *PipelinePack) (dimensions []string) {
	dims := make([]string, len(s.Fields))
	for i, _ := range dims {
		dims[i] = "UNKNOWN"
	}

	// TODO: add support for top-level message fields (Timestamp, etc)
	remaining := len(dims)
	for _, field := range pack.Message.Fields {
		if remaining == 0 {
			break
		}

		idx, ok := s.FieldIndices[field.GetName()]
		if ok {
			remaining -= 1
			inValues := field.GetValueString()
			if len(inValues) > 0 {
				// We use the first available value, even if several have been
				// provided.
				v, err := s.GetValue(field.GetName(), inValues[0])
				if err != nil {
					fmt.Printf("How did this happen? %s", err)
				}
				dims[idx] = v
			} // Else there were no values, leave this field as unknown.
		}
	}

	return dims
}

// Interface for calculating whether a particular value is acceptable
// as-is, or if it should be replaced with a default value.
type DimensionChecker interface {
	IsAllowed(v string) (bool)
}

// Accept any value at all.
type AnyDimensionChecker struct {
}
func (adc AnyDimensionChecker) IsAllowed(v string) (bool) {
	return true
}

// Accept a specific list of values, anything not in the list
// will not be accepted
type ListDimensionChecker struct {
	// Use a map instead of a list internally for fast lookups.
	allowed map[string]struct{}
}
func (ldc ListDimensionChecker) IsAllowed(v string) (bool) {
	_, ok := ldc.allowed[v]
	return ok
}

// Factory for creating a ListDimensionChecker using a list instead of a map
func NewListDimensionChecker(allowed []string) *ListDimensionChecker {
	dimMap := map[string]struct{}{}
	for _, a := range(allowed) {
		dimMap[a] = struct{}{}
	}
	return &ListDimensionChecker{dimMap}
}

// If both are specified, accept any value between `min` and `max` (inclusive).
// If one of the bounds is missing, only enforce the other. If neither bound is
// present, accept all values.
type RangeDimensionChecker struct {
	min string
	max string
}
func (rdc RangeDimensionChecker) IsAllowed(v string) (bool) {
	// Min and max are optional, so treat them separately.
	// TODO: ensure that Go does string comparisons in the fashion expected
	//       by this code.
	if rdc.min != "" && rdc.min > v {
		return false
	}

	if rdc.max != "" && rdc.max < v {
		return false
	}

	return true
}

// Pattern to use for sanitizing path/file components.
var sanitizePattern = regexp.MustCompile("[^a-zA-Z0-9_/.]")

// Given a string, return a sanitized version that can be used safely as part
// of a filename (for example).
func SanitizeDimension(dim string) (cleaned string) {
    return sanitizePattern.ReplaceAllString(dim, "_")
}

// Load a schema from the given file name.  The file is expected to contain
// valid JSON describing a hierarchy of dimensions, each of which specifies
// what values are "allowed" for that dimension.
// Example schema:
//   {
//     "version": 1,
//     "dimensions": [
//       { "field_name": "submissionDate", "allowed_values": {
//           { "min": "20140120", "max": "20140125" }
//       },
//       { "field_name": "sourceName",     "allowed_values": "*" },
//       { "field_name": "sourceVersion",  "allowed_values": "*" },
//       { "field_name": "reason",         "allowed_values":
//           [ "idle-daily","saved-session" ]
//       },
//       { "field_name": "appName",        "allowed_values":
//           [ "Firefox", "Fennec", "Thunderbird", "FirefoxOS", "B2G" ]
//       },
//       { "field_name": "appUpdateChannel",
//         "allowed_values":
//           [ "default", "nightly", "aurora", "beta", "release", "esr" ]
//       },
//       { "field_name": "appVersion",     "allowed_values": "*" }
//     ]
//   }
func LoadSchema(schemaFileName string) (schema Schema, err error) {
	// Placeholder for parsing JSON
	type JSchemaDimension struct {
		Field_name string
		Allowed_values interface{}
	}

	// Placeholder for parsing JSON
	type JSchema struct {
		Version int32
		Dimensions []JSchemaDimension
	}

	schemaBytes, err := ioutil.ReadFile(schemaFileName)
	if err != nil {
		return
	}

	var js JSchema

	err = json.Unmarshal(schemaBytes, &js)
	if err != nil {
		return
	}

	fields := make([]string, len(js.Dimensions))
	fieldIndices := map[string]int{}
	dims := map[string]DimensionChecker{}
	schema = Schema{fields, fieldIndices, dims}

	for i, d := range js.Dimensions {
		schema.Fields[i] = d.Field_name
		schema.FieldIndices[d.Field_name] = i
		switch d.Allowed_values.(type) {
		case string:
			if d.Allowed_values.(string) == "*" {
				schema.Dims[d.Field_name] = AnyDimensionChecker{}
			} else {
				schema.Dims[d.Field_name] = NewListDimensionChecker([]string{d.Allowed_values.(string)})
			}
		case []interface{}:
			allowed := make([]string, len(d.Allowed_values.([]interface{})))
			for i, v := range d.Allowed_values.([]interface{}) {
				allowedValue, ok := v.(string)
				if !ok {
					return schema, fmt.Errorf("Entries in 'allowed_values' for field '%s' must be strings", d.Field_name)
				}
				allowed[i] = allowedValue
			}
			schema.Dims[d.Field_name] = NewListDimensionChecker(allowed)
		case map[string]interface{}:
			vrange := d.Allowed_values.(map[string]interface{})
			minStr, ok := vrange["min"].(string)
			if !ok {
				return schema, fmt.Errorf("Value of 'min' for field '%s' must be a string", d.Field_name)
			}
			maxStr, ok := vrange["max"].(string)
			if !ok {
				return schema, fmt.Errorf("Value of 'max' for field '%s' must be a string", d.Field_name)
			}
			schema.Dims[d.Field_name] = RangeDimensionChecker{minStr, maxStr}
		}
	}
	return
}
<commit_msg>Add common code to list/filter an S3 bucket.<commit_after>/***** BEGIN LICENSE BLOCK *****
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
# ***** END LICENSE BLOCK *****/

package s3splitfile

import (
	"fmt"
	. "github.com/mozilla-services/heka/pipeline"
	"regexp"
	"io/ioutil"
	"encoding/json"
	"github.com/crowdmob/goamz/s3"
)

type PublishAttempt struct {
	Name string
	AttemptsRemaining uint32
}

// Encapsulates the directory-splitting schema
type Schema struct {
	Fields []string
	FieldIndices map[string]int
	Dims map[string]DimensionChecker
}

// Determine whether a given value is acceptable for a given field, and if not
// return a default value instead.
func (s *Schema) GetValue(field string, value string) (rvalue string, err error) {
	checker, ok := s.Dims[field]
	if !ok {
		return value, fmt.Errorf("No such field: '%s'", field)
	}
	if checker.IsAllowed(value) {
		return value, nil
	} else {
		return "OTHER", nil
	}
}

// Extract all dimensions from the given pack.
func (s *Schema) GetDimensions(pack *PipelinePack) (dimensions []string) {
	dims := make([]string, len(s.Fields))
	for i, _ := range dims {
		dims[i] = "UNKNOWN"
	}

	// TODO: add support for top-level message fields (Timestamp, etc)
	remaining := len(dims)
	for _, field := range pack.Message.Fields {
		if remaining == 0 {
			break
		}

		idx, ok := s.FieldIndices[field.GetName()]
		if ok {
			remaining -= 1
			inValues := field.GetValueString()
			if len(inValues) > 0 {
				// We use the first available value, even if several have been
				// provided.
				v, err := s.GetValue(field.GetName(), inValues[0])
				if err != nil {
					fmt.Printf("How did this happen? %s", err)
				}
				dims[idx] = v
			} // Else there were no values, leave this field as unknown.
		}
	}

	return dims
}

// Interface for calculating whether a particular value is acceptable
// as-is, or if it should be replaced with a default value.
type DimensionChecker interface {
	IsAllowed(v string) (bool)
}

// Accept any value at all.
type AnyDimensionChecker struct {
}
func (adc AnyDimensionChecker) IsAllowed(v string) (bool) {
	return true
}

// Accept a specific list of values, anything not in the list
// will not be accepted
type ListDimensionChecker struct {
	// Use a map instead of a list internally for fast lookups.
	allowed map[string]struct{}
}
func (ldc ListDimensionChecker) IsAllowed(v string) (bool) {
	_, ok := ldc.allowed[v]
	return ok
}

// Factory for creating a ListDimensionChecker using a list instead of a map
func NewListDimensionChecker(allowed []string) *ListDimensionChecker {
	dimMap := map[string]struct{}{}
	for _, a := range(allowed) {
		dimMap[a] = struct{}{}
	}
	return &ListDimensionChecker{dimMap}
}

// If both are specified, accept any value between `min` and `max` (inclusive).
// If one of the bounds is missing, only enforce the other. If neither bound is
// present, accept all values.
type RangeDimensionChecker struct {
	min string
	max string
}
func (rdc RangeDimensionChecker) IsAllowed(v string) (bool) {
	// Min and max are optional, so treat them separately.
	// TODO: ensure that Go does string comparisons in the fashion expected
	//       by this code.
	if rdc.min != "" && rdc.min > v {
		return false
	}

	if rdc.max != "" && rdc.max < v {
		return false
	}

	return true
}

// Pattern to use for sanitizing path/file components.
var sanitizePattern = regexp.MustCompile("[^a-zA-Z0-9_/.]")

// Given a string, return a sanitized version that can be used safely as part
// of a filename (for example).
func SanitizeDimension(dim string) (cleaned string) {
    return sanitizePattern.ReplaceAllString(dim, "_")
}

// Load a schema from the given file name.  The file is expected to contain
// valid JSON describing a hierarchy of dimensions, each of which specifies
// what values are "allowed" for that dimension.
// Example schema:
//   {
//     "version": 1,
//     "dimensions": [
//       { "field_name": "submissionDate", "allowed_values": {
//           { "min": "20140120", "max": "20140125" }
//       },
//       { "field_name": "sourceName",     "allowed_values": "*" },
//       { "field_name": "sourceVersion",  "allowed_values": "*" },
//       { "field_name": "reason",         "allowed_values":
//           [ "idle-daily","saved-session" ]
//       },
//       { "field_name": "appName",        "allowed_values":
//           [ "Firefox", "Fennec", "Thunderbird", "FirefoxOS", "B2G" ]
//       },
//       { "field_name": "appUpdateChannel",
//         "allowed_values":
//           [ "default", "nightly", "aurora", "beta", "release", "esr" ]
//       },
//       { "field_name": "appVersion",     "allowed_values": "*" }
//     ]
//   }
func LoadSchema(schemaFileName string) (schema Schema, err error) {
	// Placeholder for parsing JSON
	type JSchemaDimension struct {
		Field_name string
		Allowed_values interface{}
	}

	// Placeholder for parsing JSON
	type JSchema struct {
		Version int32
		Dimensions []JSchemaDimension
	}

	schemaBytes, err := ioutil.ReadFile(schemaFileName)
	if err != nil {
		return
	}

	var js JSchema

	err = json.Unmarshal(schemaBytes, &js)
	if err != nil {
		return
	}

	fields := make([]string, len(js.Dimensions))
	fieldIndices := map[string]int{}
	dims := map[string]DimensionChecker{}
	schema = Schema{fields, fieldIndices, dims}

	for i, d := range js.Dimensions {
		schema.Fields[i] = d.Field_name
		schema.FieldIndices[d.Field_name] = i
		switch d.Allowed_values.(type) {
		case string:
			if d.Allowed_values.(string) == "*" {
				schema.Dims[d.Field_name] = AnyDimensionChecker{}
			} else {
				schema.Dims[d.Field_name] = NewListDimensionChecker([]string{d.Allowed_values.(string)})
			}
		case []interface{}:
			allowed := make([]string, len(d.Allowed_values.([]interface{})))
			for i, v := range d.Allowed_values.([]interface{}) {
				allowedValue, ok := v.(string)
				if !ok {
					return schema, fmt.Errorf("Entries in 'allowed_values' for field '%s' must be strings", d.Field_name)
				}
				allowed[i] = allowedValue
			}
			schema.Dims[d.Field_name] = NewListDimensionChecker(allowed)
		case map[string]interface{}:
			vrange := d.Allowed_values.(map[string]interface{})
			minStr, ok := vrange["min"].(string)
			if !ok {
				return schema, fmt.Errorf("Value of 'min' for field '%s' must be a string", d.Field_name)
			}
			maxStr, ok := vrange["max"].(string)
			if !ok {
				return schema, fmt.Errorf("Value of 'max' for field '%s' must be a string", d.Field_name)
			}
			schema.Dims[d.Field_name] = RangeDimensionChecker{minStr, maxStr}
		}
	}
	return
}

// Maximum number of S3 List results to fetch at once.
const listBatchSize = 1000

// Encapsulates the result of a List operation, allowing detection of errors
// along the way.
type S3Result struct {
	Key s3.Key
	Err error
}

// List the contents of the given bucket, sending matching filenames to a
// channel which can be read by the caller.
func S3Iterator(bucket *s3.Bucket, prefix string, schema Schema) <-chan S3Result {
	keyChannel := make(chan S3Result, listBatchSize)
	go FilterS3(bucket, prefix, 0, schema, keyChannel)
	return keyChannel
}

// Recursively descend into an S3 directory tree, filtering based on the given
// schema, and sending results on the given channel. The `level` parameter
// indicates how far down the tree we are, and is used to determine which schema
// field we use for filtering.
func FilterS3(bucket *s3.Bucket, prefix string, level int, schema Schema, kc chan S3Result) {
    // Update the marker as we encounter keys / prefixes. If a response is
    // truncated, the next `List` request will start from the next item after
    // the marker.
    marker := ""

    // Keep listing if the response is incomplete (there are more than
    // `listBatchSize` entries or prefixes)
	done := false
    for !done {
	    response, err := bucket.List(prefix, "/", marker, listBatchSize)
		if err != nil {
			fmt.Printf("Error listing: %s\n", err)
			// TODO: retry?
			kc <- S3Result{s3.Key{}, err}
		}

		if !response.IsTruncated {
			// Response is not truncated, so we're done.
			done = true
		}

		if level >= len(schema.Fields) {
			// We are past all the dimensions - encountered items are now
			// S3 key names. We ignore any further prefixes and assume that the
			// specified schema is correct/complete.
			for _, k := range response.Contents {
				marker = k.Key
				kc <- S3Result{k, nil}
			}
		} else {
			// We are still looking at prefixes. Recursively list each one that
			// matches the specified schema's allowed values.
			for _, pf := range response.CommonPrefixes {
				// Get just the last piece of the prefix to check it as a
				// dimension. If we have '/foo/bar/baz', we just want 'baz'.
				stripped := pf[len(prefix):len(pf)-1]
				allowed := schema.Dims[schema.Fields[level]].IsAllowed(stripped)
				marker = pf
				if allowed {
					FilterS3(bucket, pf, level + 1, schema, kc)
				}
			}
		}
	}

	if level == 0 {
		// We traverse the tree in depth-first order, so once we've reached the
		// end at the root (level 0), we know we're done.
		// Note that things could be made faster by parallelizing the recursive
		// listing, but we would need some other mechanism to know when to close
		// the channel?
		close(kc)
	}
	return
}
<|endoftext|>


<commit_before>// Parse out
// package and activity(MAIN) from apk
package main

import (
	"archive/zip"
	"fmt"
	"io/ioutil"
	"log"

	"github.com/nick-fedesna/android/apk"
)

func parseManifest(data []byte) {
	var manifest apk.Manifest
	err := apk.Unmarshal(data, &manifest)
	if err != nil {
		log.Fatal(err)
	}
	var mainAct *apk.AppActivity
	for _, act := range manifest.App.Activity {
		if mainAct != nil {
			break
		}
		for _, intent := range act.IntentFilter {
			if intent.Action.Name == "android.intent.action.MAIN" {
				mainAct = &act
				break
			}
		}
	}
	log.Println(manifest.Package)
	log.Println(mainAct.Name)
	fmt.Printf("adb shell am start -n %s/%s\n", manifest.Package, mainAct.Name)
	//out, _ := xml.MarshalIndent(manifest, "", "\t")
	//fmt.Printf("%s\n", string(out))
}

func ReadManifestFromApk(filename string) (data []byte, err error) {
	r, err := zip.OpenReader(filename)
	if err != nil {
		return
	}
	defer r.Close()
	for _, f := range r.File {
		if f.Name != "AndroidManifest.xml" {
			continue
		}
		rc, er := f.Open()
		if er != nil {
			return nil, er
		}
		data, err = ioutil.ReadAll(rc)
		rc.Close()
		return
	}
	return nil, fmt.Errorf("File not found: AndroidManifest.xml")
}

func main() {
	data, err := ReadManifestFromApk("ScreenTest.apk")
	if err != nil {
		log.Fatal(err)
	}
	parseManifest(data)
}
<commit_msg>fix bug<commit_after>// Parse out
// package and activity(MAIN) from apk
package main

import (
	"archive/zip"
	"flag"
	"fmt"
	"io/ioutil"
	"log"

	"github.com/nick-fedesna/android/apk"
)

func parseManifest(data []byte) {
	var manifest apk.Manifest
	err := apk.Unmarshal(data, &manifest)
	if err != nil {
		log.Fatal(err)
	}
	var launchActivity apk.AppActivity
	for _, act := range manifest.App.Activity {
		for _, intent := range act.IntentFilter {
			if intent.Action.Name == "android.intent.action.MAIN" &&
				intent.Category.Name == "android.intent.category.LAUNCHER" {
				launchActivity = act
				goto FOUND
			}
		}
	}
FOUND:
	fmt.Println(manifest.Package)
	fmt.Println(launchActivity.Name)
	fmt.Printf("adb shell am start -n %s/%s\n", manifest.Package, launchActivity.Name)
	//out, _ := xml.MarshalIndent(manifest, "", "\t")
	//fmt.Printf("%s\n", string(out))
}

func ReadManifestFromApk(filename string) (data []byte, err error) {
	r, err := zip.OpenReader(filename)
	if err != nil {
		return
	}
	defer r.Close()
	for _, f := range r.File {
		if f.Name != "AndroidManifest.xml" {
			continue
		}
		rc, er := f.Open()
		if er != nil {
			return nil, er
		}
		data, err = ioutil.ReadAll(rc)
		rc.Close()
		return
	}
	return nil, fmt.Errorf("File not found: AndroidManifest.xml")
}

func main() {
	flag.Parse()
	data, err := ReadManifestFromApk(flag.Arg(0))
	if err != nil {
		log.Fatal(err)
	}
	parseManifest(data)
}
<|endoftext|>


<commit_before>/*
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 (
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"

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

type chanChan chan chan struct{}

var (
	filename       = flag.String("filename", "", "Filename to upload. Can be a URL")
	thumbnail      = flag.String("thumbnail", "", "Thumbnail to upload. Can be a URL")
	title          = flag.String("title", "Video Title", "Video title")
	description    = flag.String("description", "uploaded by youtubeuploader", "Video description")
	categoryId     = flag.String("categoryId", "", "Video category Id")
	tags           = flag.String("tags", "", "Comma separated list of video tags")
	privacy        = flag.String("privacy", "private", "Video privacy status")
	quiet          = flag.Bool("quiet", false, "Suppress progress indicator")
	rate           = flag.Int("ratelimit", 0, "Rate limit upload in kbps. No limit by default")
	metaJSON       = flag.String("metaJSON", "", "JSON file containing title,description,tags etc (optional)")
	limitBetween   = flag.String("limitBetween", "", "Only rate limit between these times e.g. 10:00-14:00 (local time zone)")
	headlessAuth   = flag.Bool("headlessAuth", false, "set this if no browser available for the oauth authorisation step")
	showAppVersion = flag.Bool("v", false, "show version")
	chunksize      = flag.Int("chunksize", googleapi.DefaultUploadChunkSize, "size (in bytes) of each upload chunk. A zero value will cause all data to be uploaded in a single request")

	// this is set by compile-time to match git tag
	appVersion string = "unknown"
)

func main() {
	flag.Parse()

	if *showAppVersion {
		fmt.Printf("Youtubeuploader version: %s\n", appVersion)
		os.Exit(0)
	}

	if *filename == "" {
		fmt.Printf("You must provide a filename of a video file to upload\n")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var limitRange limitRange
	if *limitBetween != "" {
		var err error
		limitRange, err = parseLimitBetween(*limitBetween)
		if err != nil {
			fmt.Printf("Invalid value for -limitBetween: %v", err)
			os.Exit(1)
		}
	}

	reader, filesize := Open(*filename)
	defer reader.Close()

	var thumbReader io.ReadCloser
	if *thumbnail != "" {
		thumbReader, _ = Open(*thumbnail)
		defer thumbReader.Close()
	}

	ctx := context.Background()
	transport := &limitTransport{rt: http.DefaultTransport, lr: limitRange, filesize: filesize}
	ctx = context.WithValue(ctx, oauth2.HTTPClient, &http.Client{
		Transport: transport,
	})

	var quitChan chanChan
	if !*quiet {
		quitChan = make(chanChan)
		go func() {
			Progress(quitChan, transport, filesize)
		}()
	}
	client, err := buildOAuthHTTPClient(ctx, []string{youtube.YoutubeUploadScope, youtube.YoutubeScope})
	if err != nil {
		log.Fatalf("Error building OAuth client: %v", err)
	}

	upload := &youtube.Video{
		Snippet:          &youtube.VideoSnippet{},
		RecordingDetails: &youtube.VideoRecordingDetails{},
		Status:           &youtube.VideoStatus{},
	}

	videoMeta := LoadVideoMeta(*metaJSON, upload)

	service, err := youtube.New(client)
	if err != nil {
		log.Fatalf("Error creating Youtube client: %s", err)
	}

	if upload.Status.PrivacyStatus == "" {
		upload.Status.PrivacyStatus = *privacy
	}
	if upload.Snippet.Tags == nil && strings.Trim(*tags, "") != "" {
		upload.Snippet.Tags = strings.Split(*tags, ",")
	}
	if upload.Snippet.Title == "" {
		upload.Snippet.Title = *title
	}
	if upload.Snippet.Description == "" {
		upload.Snippet.Description = *description
	}
	if upload.Snippet.CategoryId == "" && *categoryId != "" {
		upload.Snippet.CategoryId = *categoryId
	}

	fmt.Printf("Uploading file '%s'...\n", *filename)

	var option googleapi.MediaOption
	var video *youtube.Video

	option = googleapi.ChunkSize(*chunksize)

	call := service.Videos.Insert("snippet,status,recordingDetails", upload)
	video, err = call.Media(reader, option).Do()

	if quitChan != nil {
		quit := make(chan struct{})
		quitChan <- quit
		<-quit
	}

	if err != nil {
		if video != nil {
			log.Fatalf("Error making YouTube API call: %v, %v", err, video.HTTPStatusCode)
		} else {
			log.Fatalf("Error making YouTube API call: %v", err)
		}
	}
	fmt.Printf("Upload successful! Video ID: %v\n", video.Id)

	plx := &Playlistx{}
	if upload.Status.PrivacyStatus != "" {
		plx.PrivacyStatus = upload.Status.PrivacyStatus
	}
	// PlaylistID is deprecated in favour of PlaylistIDs
	if videoMeta.PlaylistID != "" {
		plx.Id = videoMeta.PlaylistID
		err = plx.AddVideoToPlaylist(service, video.Id)
		if err != nil {
			log.Fatalf("Error adding video to playlist: %s", err)
		}
	}

	if len(videoMeta.PlaylistIDs) > 0 {
		plx.Title = ""
		for _, pid := range videoMeta.PlaylistIDs {
			plx.Id = pid
			err = plx.AddVideoToPlaylist(service, video.Id)
			if err != nil {
				log.Fatalf("Error adding video to playlist: %s", err)
			}
		}
	}

	if len(videoMeta.PlaylistTitles) > 0 {
		plx.Id = ""
		for _, title := range videoMeta.PlaylistTitles {
			plx.Title = title
			err = plx.AddVideoToPlaylist(service, video.Id)
			if err != nil {
				log.Fatalf("Error adding video to playlist: %s", err)
			}
		}
	}

	if thumbReader != nil {
		log.Printf("Uploading thumbnail '%s'...\n", *thumbnail)
		_, err = service.Thumbnails.Set(video.Id).Media(thumbReader).Do()
		if err != nil {
			log.Fatalf("Error making YouTube API call: %v", err)
		}
		fmt.Printf("Thumbnail uploaded!\n")
	}
}
<commit_msg>Upload thumbnail before playlist add<commit_after>/*
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 (
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"

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

type chanChan chan chan struct{}

var (
	filename       = flag.String("filename", "", "Filename to upload. Can be a URL")
	thumbnail      = flag.String("thumbnail", "", "Thumbnail to upload. Can be a URL")
	title          = flag.String("title", "Video Title", "Video title")
	description    = flag.String("description", "uploaded by youtubeuploader", "Video description")
	categoryId     = flag.String("categoryId", "", "Video category Id")
	tags           = flag.String("tags", "", "Comma separated list of video tags")
	privacy        = flag.String("privacy", "private", "Video privacy status")
	quiet          = flag.Bool("quiet", false, "Suppress progress indicator")
	rate           = flag.Int("ratelimit", 0, "Rate limit upload in kbps. No limit by default")
	metaJSON       = flag.String("metaJSON", "", "JSON file containing title,description,tags etc (optional)")
	limitBetween   = flag.String("limitBetween", "", "Only rate limit between these times e.g. 10:00-14:00 (local time zone)")
	headlessAuth   = flag.Bool("headlessAuth", false, "set this if no browser available for the oauth authorisation step")
	showAppVersion = flag.Bool("v", false, "show version")
	chunksize      = flag.Int("chunksize", googleapi.DefaultUploadChunkSize, "size (in bytes) of each upload chunk. A zero value will cause all data to be uploaded in a single request")

	// this is set by compile-time to match git tag
	appVersion string = "unknown"
)

func main() {
	flag.Parse()

	if *showAppVersion {
		fmt.Printf("Youtubeuploader version: %s\n", appVersion)
		os.Exit(0)
	}

	if *filename == "" {
		fmt.Printf("You must provide a filename of a video file to upload\n")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var limitRange limitRange
	if *limitBetween != "" {
		var err error
		limitRange, err = parseLimitBetween(*limitBetween)
		if err != nil {
			fmt.Printf("Invalid value for -limitBetween: %v", err)
			os.Exit(1)
		}
	}

	reader, filesize := Open(*filename)
	defer reader.Close()

	var thumbReader io.ReadCloser
	if *thumbnail != "" {
		thumbReader, _ = Open(*thumbnail)
		defer thumbReader.Close()
	}

	ctx := context.Background()
	transport := &limitTransport{rt: http.DefaultTransport, lr: limitRange, filesize: filesize}
	ctx = context.WithValue(ctx, oauth2.HTTPClient, &http.Client{
		Transport: transport,
	})

	var quitChan chanChan
	if !*quiet {
		quitChan = make(chanChan)
		go func() {
			Progress(quitChan, transport, filesize)
		}()
	}
	client, err := buildOAuthHTTPClient(ctx, []string{youtube.YoutubeUploadScope, youtube.YoutubeScope})
	if err != nil {
		log.Fatalf("Error building OAuth client: %v", err)
	}

	upload := &youtube.Video{
		Snippet:          &youtube.VideoSnippet{},
		RecordingDetails: &youtube.VideoRecordingDetails{},
		Status:           &youtube.VideoStatus{},
	}

	videoMeta := LoadVideoMeta(*metaJSON, upload)

	service, err := youtube.New(client)
	if err != nil {
		log.Fatalf("Error creating Youtube client: %s", err)
	}

	if upload.Status.PrivacyStatus == "" {
		upload.Status.PrivacyStatus = *privacy
	}
	if upload.Snippet.Tags == nil && strings.Trim(*tags, "") != "" {
		upload.Snippet.Tags = strings.Split(*tags, ",")
	}
	if upload.Snippet.Title == "" {
		upload.Snippet.Title = *title
	}
	if upload.Snippet.Description == "" {
		upload.Snippet.Description = *description
	}
	if upload.Snippet.CategoryId == "" && *categoryId != "" {
		upload.Snippet.CategoryId = *categoryId
	}

	fmt.Printf("Uploading file '%s'...\n", *filename)

	var option googleapi.MediaOption
	var video *youtube.Video

	option = googleapi.ChunkSize(*chunksize)

	call := service.Videos.Insert("snippet,status,recordingDetails", upload)
	video, err = call.Media(reader, option).Do()

	if quitChan != nil {
		quit := make(chan struct{})
		quitChan <- quit
		<-quit
	}

	if err != nil {
		if video != nil {
			log.Fatalf("Error making YouTube API call: %v, %v", err, video.HTTPStatusCode)
		} else {
			log.Fatalf("Error making YouTube API call: %v", err)
		}
	}
	fmt.Printf("Upload successful! Video ID: %v\n", video.Id)

	if thumbReader != nil {
		log.Printf("Uploading thumbnail '%s'...\n", *thumbnail)
		_, err = service.Thumbnails.Set(video.Id).Media(thumbReader).Do()
		if err != nil {
			log.Fatalf("Error making YouTube API call: %v", err)
		}
		fmt.Printf("Thumbnail uploaded!\n")
	}

	plx := &Playlistx{}
	if upload.Status.PrivacyStatus != "" {
		plx.PrivacyStatus = upload.Status.PrivacyStatus
	}
	// PlaylistID is deprecated in favour of PlaylistIDs
	if videoMeta.PlaylistID != "" {
		plx.Id = videoMeta.PlaylistID
		err = plx.AddVideoToPlaylist(service, video.Id)
		if err != nil {
			log.Fatalf("Error adding video to playlist: %s", err)
		}
	}

	if len(videoMeta.PlaylistIDs) > 0 {
		plx.Title = ""
		for _, pid := range videoMeta.PlaylistIDs {
			plx.Id = pid
			err = plx.AddVideoToPlaylist(service, video.Id)
			if err != nil {
				log.Fatalf("Error adding video to playlist: %s", err)
			}
		}
	}

	if len(videoMeta.PlaylistTitles) > 0 {
		plx.Id = ""
		for _, title := range videoMeta.PlaylistTitles {
			plx.Title = title
			err = plx.AddVideoToPlaylist(service, video.Id)
			if err != nil {
				log.Fatalf("Error adding video to playlist: %s", err)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/APTrust/bagman/bagman"
	"github.com/APTrust/bagman/dpn"
	"github.com/APTrust/bagman/workers"
	"math/rand"
	"os"
	"path/filepath"
)

var TEST_NODE_URLS = map[string]string {
	"chron": "http://127.0.0.1:8001",
	"hathi": "http://127.0.0.1:8002",
	"sdr":   "http://127.0.0.1:8003",
	"tdr":   "http://127.0.0.1:8004",
}
var testBagUuid = "00000000-0000-0000-0000-000000000001"
var goodBagPath = fmt.Sprintf("dpn/testdata/%s.tar", testBagUuid)
var testBagSize = uint64(268800)
var testBagDigest = "f9f39a1602cde405042dd8b4859c6a3e2c04092a76eaab858ae28e48403ccba4"
var adminTestToken = "0000000000000000000000000000000000000000"

// TODO: Remote admin clients with all-zero tokens

func main() {
	testUtil := NewTestUtil()
	err := testUtil.MakeTestData()
	if err != nil {
		testUtil.ProcUtil.MessageLog.Fatal(err)
	}
}

type TestUtil struct {
	ProcUtil        *bagman.ProcessUtil
	DPNConfig       *dpn.DPNConfig
	LocalRestClient *dpn.DPNRestClient
	RemoteClients   map[string]*dpn.DPNRestClient
}

func NewTestUtil() (*TestUtil) {
	procUtil := workers.CreateProcUtil()
	procUtil.MessageLog.Info("Starting data setup for local integration test")
	dpnConfig, err := dpn.LoadConfig("dpn/dpn_config.json")
	if err != nil {
		procUtil.MessageLog.Fatal(err.Error())
	}

	localClient, err := dpn.NewDPNRestClient(
		dpnConfig.RestClient.LocalServiceURL,
		dpnConfig.RestClient.LocalAPIRoot,
		dpnConfig.RestClient.LocalAuthToken,
		procUtil.MessageLog)
	if err != nil {
		procUtil.MessageLog.Fatal(err.Error())
	}
	remoteClients, err := dpn.GetRemoteClients(localClient, dpnConfig,
		procUtil.MessageLog)

	// Point the remote clients toward our own local DPN test cluster.
	// This means you have to run the run_cluster.sh script in the
	// DPN REST project to run these tests.
	for nodeNamespace := range remoteClients {
		remoteClient := remoteClients[nodeNamespace]
		remoteClient.HostUrl = TEST_NODE_URLS[nodeNamespace]
	}

	return &TestUtil{
		ProcUtil: procUtil,
		DPNConfig: dpnConfig,
		LocalRestClient: localClient,
		RemoteClients: remoteClients,
	}
}

func (testUtil *TestUtil) MakeTestData() (error) {
	allNodes, err := testUtil.LocalRestClient.DPNNodeListGet(nil)
	if err != nil {
		return err
	}
	count := 0
	for _, node := range allNodes.Results {
		if node.Namespace == testUtil.DPNConfig.LocalNode {
			continue
		}
		count += 1

		// Create a symlink from dpn_home/integration_test/<uuid>.tar
		// to our known good bag in dpn/testdata/000...1.tar
		bagUuid := fmt.Sprintf("%d0000000-0000-0000-0000-000000000001", count)
		linkPath, err := testUtil.CreateSymLink(bagUuid)
		if err != nil {
			return err
		} else {
			testUtil.ProcUtil.MessageLog.Info("Created symlink at %s", linkPath)
		}

		// Create an entry for this bag on the remote node.
		bag, err := testUtil.CreateBag(bagUuid, node.Namespace)
		if err != nil {
			return err
		} else {
			testUtil.ProcUtil.MessageLog.Info("Created bag %s on %s",
				bag.UUID, bag.AdminNode)
		}

		// Create a transfer record for this bag on the remote node.
		xfer, err := testUtil.CreateReplicationRequest(bag, linkPath)
		if err != nil {
			return err
		} else {
			testUtil.ProcUtil.MessageLog.Info(
				"Created replication request %s on %s",
				xfer.ReplicationId, bag.AdminNode)
		}
	}
	return nil
}

func (testUtil *TestUtil) CreateSymLink(bagUuid string) (string, error) {
	sourceFile, err := bagman.RelativeToAbsPath(goodBagPath)
	if err != nil {
		return "", err
	}
	linkPath := filepath.Join(testUtil.ProcUtil.Config.DPNHomeDirectory,
		"integration_test", bagUuid + ".tar")
	if bagman.FileExists(sourceFile) {
		return linkPath, nil
	}
	err = os.Symlink(sourceFile, linkPath)
	if err != nil {
		return "", err
	}
	return linkPath, err
}

func (testUtil *TestUtil) CreateBag(bagUuid, node string) (*dpn.DPNBag, error) {
	bag, err := testUtil.RemoteClients[node].DPNBagGet(bagUuid)
	if err == nil && bag != nil {
		// Bag already exists. No need to recreate it.
		return bag, err
	}
	bag = &dpn.DPNBag{
		UUID: bagUuid,
		LocalId: fmt.Sprintf("integration-test-%s-1", node),
		Size: testBagSize,
		FirstVersionUUID: bagUuid,
		Version: 1,
		IngestNode: node,
		AdminNode: node,
		BagType: "D",
		Rights: make([]string, 0),
		Interpretive: make([]string, 0),
		ReplicatingNodes: make([]string, 0),
		Fixities: []*dpn.DPNFixity {
			&dpn.DPNFixity{
				Sha256: testBagDigest,
			},
		},
	}
	return testUtil.RemoteClients[node].DPNBagCreate(bag)
}

func (testUtil *TestUtil) CreateReplicationRequest(bag *dpn.DPNBag, linkPath string) (*dpn.DPNReplicationTransfer, error) {
	replicationId := fmt.Sprintf("%s-%d", bag.AdminNode, rand.Intn(200000000))
	xfer := &dpn.DPNReplicationTransfer{
		FromNode: bag.AdminNode,
		ToNode: testUtil.DPNConfig.LocalNode,
		UUID: bag.UUID,
		ReplicationId: replicationId,
		FixityAlgorithm: "sha256",
		Status: "Requested",
		Protocol: "R",
		Link: linkPath,
	}
	return testUtil.RemoteClients[bag.AdminNode].ReplicationTransferCreate(xfer)
}
<commit_msg>Use admin REST clients to do admin tasks<commit_after>package main

import (
	"fmt"
	"github.com/APTrust/bagman/bagman"
	"github.com/APTrust/bagman/dpn"
	"github.com/APTrust/bagman/workers"
	"math/rand"
	"os"
	"path/filepath"
)

var TEST_NODE_URLS = map[string]string {
	"chron": "http://127.0.0.1:8001",
	"hathi": "http://127.0.0.1:8002",
	"sdr":   "http://127.0.0.1:8003",
	"tdr":   "http://127.0.0.1:8004",
}
var testBagUuid = "00000000-0000-0000-0000-000000000001"
var goodBagPath = fmt.Sprintf("dpn/testdata/%s.tar", testBagUuid)
var testBagSize = uint64(268800)
var testBagDigest = "f9f39a1602cde405042dd8b4859c6a3e2c04092a76eaab858ae28e48403ccba4"
var adminTestToken = "0000000000000000000000000000000000000000"

func main() {
	testUtil := NewTestUtil()
	err := testUtil.MakeTestData()
	if err != nil {
		testUtil.ProcUtil.MessageLog.Fatal(err)
	}
}

type TestUtil struct {
	ProcUtil             *bagman.ProcessUtil
	DPNConfig            *dpn.DPNConfig
	LocalRestClient      *dpn.DPNRestClient
	RemoteClients        map[string]*dpn.DPNRestClient
	RemoteAdminClients   map[string]*dpn.DPNRestClient
}

func NewTestUtil() (*TestUtil) {
	procUtil := workers.CreateProcUtil()
	procUtil.MessageLog.Info("Starting data setup for local integration test")
	dpnConfig, err := dpn.LoadConfig("dpn/dpn_config.json")
	if err != nil {
		procUtil.MessageLog.Fatal(err.Error())
	}

	// Create a local REST client to talk to our local APTrust DPN REST server
	localClient, err := dpn.NewDPNRestClient(
		dpnConfig.RestClient.LocalServiceURL,
		dpnConfig.RestClient.LocalAPIRoot,
		dpnConfig.RestClient.LocalAuthToken,
		procUtil.MessageLog)
	if err != nil {
		procUtil.MessageLog.Fatal(err.Error())
	}

	// Create clients to talk to remote DPN REST servers.
	// Actually, these are servers in the local cluster that
	// impersonate remote nodes. All of these admin clients
	// need to use the admin API token. We need these admin
	// clients to do admin-only work, such as creating bags
	// and transfer requests.
	remoteClients, err := dpn.GetRemoteClients(localClient, dpnConfig,
		procUtil.MessageLog)
	adminConfig := *dpnConfig
	adminConfig.RemoteNodeTokens["chron"] = adminTestToken
	adminConfig.RemoteNodeTokens["hathi"] = adminTestToken
	adminConfig.RemoteNodeTokens["sdr"] = adminTestToken
	adminConfig.RemoteNodeTokens["tdr"] = adminTestToken
	remoteAdminClients, err := dpn.GetRemoteClients(localClient,
		&adminConfig, procUtil.MessageLog)

	// Point the remote clients toward our own local DPN test cluster.
	// This means you have to run the run_cluster.sh script in the
	// DPN REST project to run these tests.
	for nodeNamespace := range remoteClients {
		remoteClient := remoteClients[nodeNamespace]
		remoteClient.HostUrl = TEST_NODE_URLS[nodeNamespace]
		remoteAdminClient := remoteAdminClients[nodeNamespace]
		remoteAdminClient.HostUrl = TEST_NODE_URLS[nodeNamespace]
	}

	return &TestUtil{
		ProcUtil: procUtil,
		DPNConfig: dpnConfig,
		LocalRestClient: localClient,
		RemoteClients: remoteClients,
		RemoteAdminClients: remoteAdminClients,
	}
}

func (testUtil *TestUtil) MakeTestData() (error) {
	allNodes, err := testUtil.LocalRestClient.DPNNodeListGet(nil)
	if err != nil {
		return err
	}
	count := 0
	for _, node := range allNodes.Results {
		if node.Namespace == testUtil.DPNConfig.LocalNode {
			continue
		}
		count += 1

		// Create a symlink from dpn_home/integration_test/<uuid>.tar
		// to our known good bag in dpn/testdata/000...1.tar
		bagUuid := fmt.Sprintf("%d0000000-0000-0000-0000-000000000001", count)
		linkPath, err := testUtil.CreateSymLink(bagUuid)
		if err != nil {
			return err
		} else {
			testUtil.ProcUtil.MessageLog.Info("Created symlink at %s", linkPath)
		}

		// Create an entry for this bag on the remote node.
		bag, err := testUtil.CreateBag(bagUuid, node.Namespace)
		if err != nil {
			return err
		} else {
			testUtil.ProcUtil.MessageLog.Info("Created bag %s on %s",
				bag.UUID, bag.AdminNode)
		}

		// Create a transfer record for this bag on the remote node.
		xfer, err := testUtil.CreateReplicationRequest(bag, linkPath)
		if err != nil {
			return err
		} else {
			testUtil.ProcUtil.MessageLog.Info(
				"Created replication request %s on %s",
				xfer.ReplicationId, bag.AdminNode)
		}
	}
	return nil
}

func (testUtil *TestUtil) CreateSymLink(bagUuid string) (string, error) {
	sourceFile, err := bagman.RelativeToAbsPath(goodBagPath)
	if err != nil {
		return "", err
	}
	linkPath := filepath.Join(testUtil.ProcUtil.Config.DPNHomeDirectory,
		"integration_test", bagUuid + ".tar")
	if bagman.FileExists(linkPath) {
		return linkPath, nil
	}
	err = os.Symlink(sourceFile, linkPath)
	if err != nil {
		return "", err
	}
	return linkPath, err
}

func (testUtil *TestUtil) CreateBag(bagUuid, node string) (*dpn.DPNBag, error) {
	bag, err := testUtil.RemoteClients[node].DPNBagGet(bagUuid)
	if err == nil && bag != nil {
		// Bag already exists. No need to recreate it.
		return bag, err
	}
	bag = &dpn.DPNBag{
		UUID: bagUuid,
		LocalId: fmt.Sprintf("integration-test-%s-1", node),
		Size: testBagSize,
		FirstVersionUUID: bagUuid,
		Version: 1,
		IngestNode: node,
		AdminNode: node,
		BagType: "D",
		Rights: make([]string, 0),
		Interpretive: make([]string, 0),
		ReplicatingNodes: make([]string, 0),
		Fixities: []*dpn.DPNFixity {
			&dpn.DPNFixity{
				Sha256: testBagDigest,
			},
		},
	}
	// You have to be node admin to create a bag, so use the admin client.
	return testUtil.RemoteAdminClients[node].DPNBagCreate(bag)
}

func (testUtil *TestUtil) CreateReplicationRequest(bag *dpn.DPNBag, linkPath string) (*dpn.DPNReplicationTransfer, error) {
	replicationId := fmt.Sprintf("%s-%d", bag.AdminNode, rand.Intn(200000000))
	xfer := &dpn.DPNReplicationTransfer{
		FromNode: bag.AdminNode,
		ToNode: testUtil.DPNConfig.LocalNode,
		UUID: bag.UUID,
		ReplicationId: replicationId,
		FixityAlgorithm: "sha256",
		Status: "Requested",
		Protocol: "R",
		Link: linkPath,
	}
	// You have to be node admin to create the transfer request,
	// so use the admin client.
	return testUtil.RemoteAdminClients[bag.AdminNode].ReplicationTransferCreate(xfer)
}
<|endoftext|>


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

// +build example

package main

import (
	"log"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/ebitenutil"
)

var runes = append(make([]rune, 0, 1024), []rune("Type on the keyboard:\n")...)

var buf = make([]rune, 1024)

var counter int

func update(screen *ebiten.Image) error {
	runes = append(runes, ebiten.InputChars()...)
	if ebiten.IsKeyPressed(ebiten.KeyEnter) {
		if len(runes) > 0 && runes[len(runes)-1] != '\n' {
			runes = append(runes, '\n')
		}
	}
	counter++
	if ebiten.IsRunningSlowly() {
		return nil
	}
	if counter%60 < 30 {
		return ebitenutil.DebugPrint(screen, string(append(runes, '_')))
	}
	return ebitenutil.DebugPrint(screen, string(runes))
}

func main() {
	log.Fatal(ebiten.Run(update, 320, 240, 2.0, "Runes (Ebiten Demo)")) // ebiterm?
}
<commit_msg>examples/runes: Refactoring<commit_after>// Copyright 2017 The Ebiten Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build example

package main

import (
	"log"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/ebitenutil"
)

var (
	runes   = append(make([]rune, 0, 1024), []rune("Type on the keyboard:\n")...)
	counter = 0
)

func update(screen *ebiten.Image) error {
	runes = append(runes, ebiten.InputChars()...)
	if ebiten.IsKeyPressed(ebiten.KeyEnter) {
		if len(runes) > 0 && runes[len(runes)-1] != '\n' {
			runes = append(runes, '\n')
		}
	}
	counter++

	if ebiten.IsRunningSlowly() {
		return nil
	}

	if counter%60 < 30 {
		return ebitenutil.DebugPrint(screen, string(append(runes, '_')))
	}
	return ebitenutil.DebugPrint(screen, string(runes))
}

func main() {
	if err := ebiten.Run(update, 320, 240, 2.0, "Runes (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package gce

import (
	"context"
	computealpha "google.golang.org/api/compute/v0.alpha"
	"strings"
)

const (
	NEGLoadBalancerType          = "LOAD_BALANCING"
	NEGIPPortNetworkEndpointType = "GCE_VM_IP_PORT"
)

func newNetworkEndpointGroupMetricContext(request string, zone string) *metricContext {
	return newGenericMetricContext("networkendpointgroup_", request, unusedMetricLabel, zone, computeAlphaVersion)
}

func (gce *GCECloud) GetNetworkEndpointGroup(name string, zone string) (*computealpha.NetworkEndpointGroup, error) {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return nil, err
	}
	mc := newNetworkEndpointGroupMetricContext("get", zone)
	v, err := gce.serviceAlpha.NetworkEndpointGroups.Get(gce.ProjectID(), zone, name).Do()
	return v, mc.Observe(err)
}

func (gce *GCECloud) ListNetworkEndpointGroup(zone string) ([]*computealpha.NetworkEndpointGroup, error) {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return nil, err
	}
	mc := newNetworkEndpointGroupMetricContext("list", zone)
	networkEndpointGroups := []*computealpha.NetworkEndpointGroup{}
	err := gce.serviceAlpha.NetworkEndpointGroups.List(gce.ProjectID(), zone).Pages(context.Background(), func(res *computealpha.NetworkEndpointGroupList) error {
		networkEndpointGroups = append(networkEndpointGroups, res.Items...)
		return nil
	})
	return networkEndpointGroups, mc.Observe(err)
}

func (gce *GCECloud) AggregatedListNetworkEndpointGroup() (map[string][]*computealpha.NetworkEndpointGroup, error) {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return nil, err
	}
	mc := newNetworkEndpointGroupMetricContext("aggregated_list", "")
	zoneNetworkEndpointGroupMap := map[string][]*computealpha.NetworkEndpointGroup{}
	err := gce.serviceAlpha.NetworkEndpointGroups.AggregatedList(gce.ProjectID()).Pages(context.Background(), func(res *computealpha.NetworkEndpointGroupAggregatedList) error {
		for key, negs := range res.Items {
			if len(negs.NetworkEndpointGroups) == 0 {
				continue
			}
			// key has the format of "zones/${zone_name}"
			zone := strings.Split(key, "/")[1]
			if _, ok := zoneNetworkEndpointGroupMap[zone]; !ok {
				zoneNetworkEndpointGroupMap[zone] = []*computealpha.NetworkEndpointGroup{}
			}
			zoneNetworkEndpointGroupMap[zone] = append(zoneNetworkEndpointGroupMap[zone], negs.NetworkEndpointGroups...)
		}
		return nil
	})
	return zoneNetworkEndpointGroupMap, mc.Observe(err)
}

func (gce *GCECloud) CreateNetworkEndpointGroup(neg *computealpha.NetworkEndpointGroup, zone string) error {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return err
	}
	mc := newNetworkEndpointGroupMetricContext("create", zone)
	op, err := gce.serviceAlpha.NetworkEndpointGroups.Insert(gce.ProjectID(), zone, neg).Do()
	if err != nil {
		return mc.Observe(err)
	}
	return gce.waitForZoneOp(op, zone, mc)
}

func (gce *GCECloud) DeleteNetworkEndpointGroup(name string, zone string) error {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return err
	}
	mc := newNetworkEndpointGroupMetricContext("delete", zone)
	op, err := gce.serviceAlpha.NetworkEndpointGroups.Delete(gce.ProjectID(), zone, name).Do()
	if err != nil {
		return mc.Observe(err)
	}
	return gce.waitForZoneOp(op, zone, mc)
}

func (gce *GCECloud) AttachNetworkEndpoints(name, zone string, endpoints []*computealpha.NetworkEndpoint) error {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return err
	}
	mc := newNetworkEndpointGroupMetricContext("attach", zone)
	op, err := gce.serviceAlpha.NetworkEndpointGroups.AttachNetworkEndpoints(gce.ProjectID(), zone, name, &computealpha.NetworkEndpointGroupsAttachEndpointsRequest{
		NetworkEndpoints: endpoints,
	}).Do()
	if err != nil {
		return mc.Observe(err)
	}
	return gce.waitForZoneOp(op, zone, mc)
}

func (gce *GCECloud) DetachNetworkEndpoints(name, zone string, endpoints []*computealpha.NetworkEndpoint) error {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return err
	}
	mc := newNetworkEndpointGroupMetricContext("detach", zone)
	op, err := gce.serviceAlpha.NetworkEndpointGroups.DetachNetworkEndpoints(gce.ProjectID(), zone, name, &computealpha.NetworkEndpointGroupsDetachEndpointsRequest{
		NetworkEndpoints: endpoints,
	}).Do()
	if err != nil {
		return mc.Observe(err)
	}
	return gce.waitForZoneOp(op, zone, mc)
}

func (gce *GCECloud) ListNetworkEndpoints(name, zone string, showHealthStatus bool) ([]*computealpha.NetworkEndpointWithHealthStatus, error) {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return nil, err
	}
	healthStatus := "SKIP"
	if showHealthStatus {
		healthStatus = "SHOW"
	}
	mc := newNetworkEndpointGroupMetricContext("list_networkendpoints", zone)
	networkEndpoints := []*computealpha.NetworkEndpointWithHealthStatus{}
	err := gce.serviceAlpha.NetworkEndpointGroups.ListNetworkEndpoints(gce.ProjectID(), zone, name, &computealpha.NetworkEndpointGroupsListEndpointsRequest{
		HealthStatus: healthStatus,
	}).Pages(context.Background(), func(res *computealpha.NetworkEndpointGroupsListNetworkEndpoints) error {
		networkEndpoints = append(networkEndpoints, res.Items...)
		return nil
	})
	return networkEndpoints, mc.Observe(err)
}
<commit_msg>Update NEG to use generated code<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package gce

import (
	"context"
	"fmt"
	"strings"

	computealpha "google.golang.org/api/compute/v0.alpha"

	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud/filter"
	"k8s.io/kubernetes/pkg/cloudprovider/providers/gce/cloud/meta"
)

const (
	NEGLoadBalancerType          = "LOAD_BALANCING"
	NEGIPPortNetworkEndpointType = "GCE_VM_IP_PORT"
)

func newNetworkEndpointGroupMetricContext(request string, zone string) *metricContext {
	return newGenericMetricContext("networkendpointgroup_", request, unusedMetricLabel, zone, computeAlphaVersion)
}

func (gce *GCECloud) GetNetworkEndpointGroup(name string, zone string) (*computealpha.NetworkEndpointGroup, error) {
	mc := newNetworkEndpointGroupMetricContext("get", zone)
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return nil, mc.Observe(err)
	}
	v, err := gce.c.AlphaNetworkEndpointGroups().Get(context.Background(), meta.ZonalKey(name, zone))
	return v, mc.Observe(err)
}

func (gce *GCECloud) ListNetworkEndpointGroup(zone string) ([]*computealpha.NetworkEndpointGroup, error) {
	mc := newNetworkEndpointGroupMetricContext("list", zone)
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return nil, mc.Observe(err)
	}
	negs, err := gce.c.AlphaNetworkEndpointGroups().List(context.Background(), zone, filter.None)
	return negs, mc.Observe(err)
}

// AggregatedListNetworkEndpointGroup returns a map of zone -> endpoint group.
func (gce *GCECloud) AggregatedListNetworkEndpointGroup() (map[string][]*computealpha.NetworkEndpointGroup, error) {
	mc := newNetworkEndpointGroupMetricContext("aggregated_list", "")
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return nil, mc.Observe(err)
	}
	// TODO: filter for the region the cluster is in.
	all, err := gce.c.AlphaNetworkEndpointGroups().AggregatedList(context.Background(), filter.None)
	if err != nil {
		return nil, mc.Observe(err)
	}
	ret := map[string][]*computealpha.NetworkEndpointGroup{}
	for key, byZone := range all {
		// key is "zones/<zone name>"
		parts := strings.Split(key, "/")
		if len(parts) != 2 {
			return nil, mc.Observe(fmt.Errorf("invalid key for AggregatedListNetworkEndpointGroup: %q", key))
		}
		zone := parts[1]
		ret[zone] = append(ret[zone], byZone...)
	}
	return ret, mc.Observe(nil)
}

func (gce *GCECloud) CreateNetworkEndpointGroup(neg *computealpha.NetworkEndpointGroup, zone string) error {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return err
	}
	mc := newNetworkEndpointGroupMetricContext("create", zone)
	return mc.Observe(gce.c.AlphaNetworkEndpointGroups().Insert(context.Background(), meta.ZonalKey(neg.Name, zone), neg))
}

func (gce *GCECloud) DeleteNetworkEndpointGroup(name string, zone string) error {
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return err
	}
	mc := newNetworkEndpointGroupMetricContext("delete", zone)
	return mc.Observe(gce.c.AlphaNetworkEndpointGroups().Delete(context.Background(), meta.ZonalKey(name, zone)))
}

func (gce *GCECloud) AttachNetworkEndpoints(name, zone string, endpoints []*computealpha.NetworkEndpoint) error {
	mc := newNetworkEndpointGroupMetricContext("attach", zone)
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return mc.Observe(err)
	}
	req := &computealpha.NetworkEndpointGroupsAttachEndpointsRequest{
		NetworkEndpoints: endpoints,
	}
	return mc.Observe(gce.c.AlphaNetworkEndpointGroups().AttachNetworkEndpoints(context.Background(), meta.ZonalKey(name, zone), req))
}

func (gce *GCECloud) DetachNetworkEndpoints(name, zone string, endpoints []*computealpha.NetworkEndpoint) error {
	mc := newNetworkEndpointGroupMetricContext("detach", zone)
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return mc.Observe(err)
	}
	req := &computealpha.NetworkEndpointGroupsDetachEndpointsRequest{
		NetworkEndpoints: endpoints,
	}
	return mc.Observe(gce.c.AlphaNetworkEndpointGroups().DetachNetworkEndpoints(context.Background(), meta.ZonalKey(name, zone), req))
}

func (gce *GCECloud) ListNetworkEndpoints(name, zone string, showHealthStatus bool) ([]*computealpha.NetworkEndpointWithHealthStatus, error) {
	mc := newNetworkEndpointGroupMetricContext("list_networkendpoints", zone)
	if err := gce.alphaFeatureEnabled(AlphaFeatureNetworkEndpointGroup); err != nil {
		return nil, mc.Observe(err)
	}
	healthStatus := "SKIP"
	if showHealthStatus {
		healthStatus = "SHOW"
	}
	req := &computealpha.NetworkEndpointGroupsListEndpointsRequest{
		HealthStatus: healthStatus,
	}
	l, err := gce.c.AlphaNetworkEndpointGroups().ListNetworkEndpoints(context.Background(), meta.ZonalKey(name, zone), req, filter.None)
	return l, mc.Observe(err)
}
<|endoftext|>