<commit_before>package scheduler

import (
	"runtime"
	"sort"
	"sync"
	"sync/atomic"
	// "time"
	// "fmt"

	"github.com/henrylee2cn/pholcus/app/aid/history"
	"github.com/henrylee2cn/pholcus/app/aid/proxy"
	"github.com/henrylee2cn/pholcus/app/downloader/context"
	"github.com/henrylee2cn/pholcus/common/util"
	"github.com/henrylee2cn/pholcus/logs"
	"github.com/henrylee2cn/pholcus/runtime/cache"
	"github.com/henrylee2cn/pholcus/runtime/status"
)

type scheduler struct {
	// Spider实例的请求矩阵列表
	matrices map[int]*Matrix
	// 总并发量计数
	count chan bool
	// 运行状态
	status int
	// 全局代理IP
	proxy *proxy.Proxy
	// 标记是否使用代理IP
	useProxy bool
	// 全局历史记录
	history history.Historier
	// 全局读写锁
	sync.RWMutex
}

// 定义全局调度
var sdl = &scheduler{
	history: history.New(),
	status:  status.RUN,
	count:   make(chan bool, cache.Task.ThreadNum),
	proxy:   proxy.New(),
}

func Init() {
	for sdl.proxy == nil {
		runtime.Gosched()
	}
	sdl.matrices = make(map[int]*Matrix)
	sdl.count = make(chan bool, cache.Task.ThreadNum)
	if cache.Task.Mode == status.OFFLINE {
		sdl.history.ReadSuccess(cache.Task.OutType, cache.Task.SuccessInherit)
		sdl.history.ReadFailure(cache.Task.OutType, cache.Task.FailureInherit)
	}
	if cache.Task.ProxyMinute > 0 {
		if sdl.proxy.Count() > 0 {
			sdl.useProxy = true
			sdl.proxy.UpdateTicker(cache.Task.ProxyMinute)
			logs.Log.Informational(" *     使用代理IP，代理IP更换频率为 %v 分钟\n", cache.Task.ProxyMinute)
		} else {
			sdl.useProxy = false
			logs.Log.Informational(" *     在线代理IP列表为空，无法使用代理IP\n")
		}
	} else {
		sdl.useProxy = false
		logs.Log.Informational(" *     不使用代理IP\n")
	}
	sdl.status = status.RUN
}

// 暂停\恢复所有爬行任务
func PauseRecover() {
	sdl.Lock()
	defer sdl.Unlock()
	switch sdl.status {
	case status.PAUSE:
		sdl.status = status.RUN
	case status.RUN:
		sdl.status = status.PAUSE
	}
}

// 终止任务
func Stop() {
	sdl.Lock()
	defer func() {
		recover()
		sdl.Unlock()
	}()

	sdl.status = status.STOP
	for _, v := range sdl.matrices {
		for _, vv := range v.reqs {
			for _, req := range vv {
				// 删除队列中未执行请求的去重记录
				DeleteSuccess(req)
			}
		}
		v.reqs = make(map[int][]*context.Request)
		v.priorities = []int{}
	}

	// 清空
	close(sdl.count)
	sdl.matrices = make(map[int]*Matrix)
}

func UpsertSuccess(record history.Record) bool {
	return sdl.history.UpsertSuccess(record)
}

func DeleteSuccess(record history.Record) {
	sdl.history.DeleteSuccess(record)
}

func UpsertFailure(req *context.Request) bool {
	return sdl.history.UpsertFailure(req)
}

func DeleteFailure(req *context.Request) {
	sdl.history.DeleteFailure(req)
}

// 获取指定蜘蛛在上一次运行时失败的请求
func PullFailure(spiderName string) []*context.Request {
	return sdl.history.PullFailure(spiderName)
}

func TryFlushHistory() {
	if cache.Task.SuccessInherit {
		sdl.history.FlushSuccess(cache.Task.OutType)
	}
	if cache.Task.FailureInherit {
		sdl.history.FlushFailure(cache.Task.OutType)
	}
}

// 一个Spider实例的请求矩阵
type Matrix struct {
	// [优先级]队列，优先级默认为0
	reqs map[int][]*context.Request
	// 优先级顺序，从低到高
	priorities []int
	// 资源使用情况计数
	resCount int32
	// 最大采集页数，以负数形式表示
	maxPage int64
	// 历史及本次失败请求
	failures map[string]*context.Request
	sync.Mutex
}

// 注册资源队列
func NewMatrix(spiderId int, maxPage int64) *Matrix {
	sdl.Lock()
	defer sdl.Unlock()
	matrix := &Matrix{
		reqs:       make(map[int][]*context.Request),
		priorities: []int{},
		maxPage:    maxPage,
		failures:   make(map[string]*context.Request),
	}
	sdl.matrices[spiderId] = matrix
	return matrix
}

// 添加请求到队列
func (self *Matrix) Push(req *context.Request) {
	// 暂停添加操作，避免大量请求堆积
	for sdl.status == status.PAUSE {
		runtime.Gosched()
	}

	sdl.RLock()
	defer sdl.RUnlock()

	// 终止添加操作
	if sdl.status == status.STOP ||
		self.maxPage >= 0 ||
		// 当req不可重复下载时，已存在成功记录则返回
		!req.IsReloadable() && !UpsertSuccess(req) {
		return
	}

	// 根据资源使用情况，一定程度避免大量请求堆积
	avg := int32(cap(sdl.count) / len(sdl.matrices))
	if avg == 0 {
		avg = 1
	}
	for self.resCount > avg {
		runtime.Gosched()
	}

	// 大致限制加入队列的请求量，并发情况下应该会比maxPage多
	atomic.AddInt64(&self.maxPage, 1)

	priority := req.GetPriority()

	// 初始化该蜘蛛下该优先级队列
	if _, found := self.reqs[priority]; !found {
		self.priorities = append(self.priorities, priority)
		sort.Ints(self.priorities) // 从小到大排序
		self.reqs[priority] = []*context.Request{}
	}

	// 添加请求到队列
	self.reqs[priority] = append(self.reqs[priority], req)
}

// 从队列取出请求，不存在时返回nil
func (self *Matrix) Pull() (req *context.Request) {
	sdl.RLock()
	defer sdl.RUnlock()
	if sdl.status != status.RUN {
		return
	}

	// 按优先级从高到低取出请求
	for i := len(self.reqs) - 1; i >= 0; i-- {
		idx := self.priorities[i]
		if len(self.reqs[idx]) > 0 {
			req = self.reqs[idx][0]
			self.reqs[idx] = self.reqs[idx][1:]
			if sdl.useProxy {
				req.SetProxy(sdl.proxy.GetOne(req.GetUrl()))
			} else {
				req.SetProxy("")
			}
			return
		}
	}
	return
}

func (self *Matrix) Len() int {
	var l int
	for _, reqs := range self.reqs {
		l += len(reqs)
	}
	return l
}

func (self *Matrix) Use() {
	defer func() {
		recover()
	}()
	sdl.count <- true
	atomic.AddInt32(&self.resCount, 1)
}

func (self *Matrix) Free() {
	<-sdl.count
	atomic.AddInt32(&self.resCount, -1)
}

func (self *Matrix) CanStop() bool {
	sdl.RLock()
	if sdl.status == status.STOP {
		sdl.RUnlock()
		return true
	}
	sdl.RUnlock()

	if self.maxPage >= 0 {
		return true
	}

	if self.resCount != 0 {
		return false
	}

	for i, count := 0, len(self.reqs); i < count; i++ {
		if len(self.reqs[i]) > 0 {
			return false
		}
	}

	if len(self.failures) > 0 {
		// 重新下载历史记录中失败的请求
		var goon bool
		for unique, req := range self.failures {
			if req == nil {
				continue
			}
			self.failures[unique] = nil
			goon = true
			logs.Log.Informational(" *     - 失败请求: [%v]\n", req.GetUrl())
			self.Push(req)
		}
		if goon {
			return false
		}
	}
	return true
}

// 等待处理中的请求完成
func (self *Matrix) Flush() {
	for self.resCount != 0 {
		runtime.Gosched()
		// time.Sleep(5e8)
	}
}

func (self *Matrix) SetFailures(reqs []*context.Request) {
	for _, req := range reqs {
		self.failures[makeUnique(req)] = req
		logs.Log.Informational(" *     + 失败请求: [%v]\n", req.GetUrl())
	}
}

func (self *Matrix) SetFailure(req *context.Request) bool {
	self.Lock()
	defer self.Unlock()
	unique := makeUnique(req)
	if _, ok := self.failures[unique]; !ok {
		// 首次失败时，在任务队列末尾重新执行一次
		self.failures[unique] = req
		logs.Log.Informational(" *     + 失败请求: [%v]\n", req.GetUrl())
		return true
	}
	// 失败两次后，加入历史失败记录
	UpsertFailure(req)
	return false
}

func makeUnique(req *context.Request) string {
	return util.MakeUnique(req.GetUrl() + req.GetMethod())
}
<commit_msg>优化sdl.Push(),降低内存占用<commit_after>package scheduler

import (
	"runtime"
	"sort"
	"sync"
	"sync/atomic"
	// "time"
	// "fmt"

	"github.com/henrylee2cn/pholcus/app/aid/history"
	"github.com/henrylee2cn/pholcus/app/aid/proxy"
	"github.com/henrylee2cn/pholcus/app/downloader/context"
	"github.com/henrylee2cn/pholcus/common/util"
	"github.com/henrylee2cn/pholcus/logs"
	"github.com/henrylee2cn/pholcus/runtime/cache"
	"github.com/henrylee2cn/pholcus/runtime/status"
)

type scheduler struct {
	// Spider实例的请求矩阵列表
	matrices map[int]*Matrix
	// 总并发量计数
	count chan bool
	// 运行状态
	status int
	// 全局代理IP
	proxy *proxy.Proxy
	// 标记是否使用代理IP
	useProxy bool
	// 全局历史记录
	history history.Historier
	// 全局读写锁
	sync.RWMutex
}

// 定义全局调度
var sdl = &scheduler{
	history: history.New(),
	status:  status.RUN,
	count:   make(chan bool, cache.Task.ThreadNum),
	proxy:   proxy.New(),
}

func Init() {
	for sdl.proxy == nil {
		runtime.Gosched()
	}
	sdl.matrices = make(map[int]*Matrix)
	sdl.count = make(chan bool, cache.Task.ThreadNum)
	if cache.Task.Mode == status.OFFLINE {
		sdl.history.ReadSuccess(cache.Task.OutType, cache.Task.SuccessInherit)
		sdl.history.ReadFailure(cache.Task.OutType, cache.Task.FailureInherit)
	}
	if cache.Task.ProxyMinute > 0 {
		if sdl.proxy.Count() > 0 {
			sdl.useProxy = true
			sdl.proxy.UpdateTicker(cache.Task.ProxyMinute)
			logs.Log.Informational(" *     使用代理IP，代理IP更换频率为 %v 分钟\n", cache.Task.ProxyMinute)
		} else {
			sdl.useProxy = false
			logs.Log.Informational(" *     在线代理IP列表为空，无法使用代理IP\n")
		}
	} else {
		sdl.useProxy = false
		logs.Log.Informational(" *     不使用代理IP\n")
	}
	sdl.status = status.RUN
}

// 暂停\恢复所有爬行任务
func PauseRecover() {
	sdl.Lock()
	defer sdl.Unlock()
	switch sdl.status {
	case status.PAUSE:
		sdl.status = status.RUN
	case status.RUN:
		sdl.status = status.PAUSE
	}
}

// 终止任务
func Stop() {
	sdl.Lock()
	defer func() {
		recover()
		sdl.Unlock()
	}()

	sdl.status = status.STOP
	for _, v := range sdl.matrices {
		for _, vv := range v.reqs {
			for _, req := range vv {
				// 删除队列中未执行请求的去重记录
				DeleteSuccess(req)
			}
		}
		v.reqs = make(map[int][]*context.Request)
		v.priorities = []int{}
	}

	// 清空
	close(sdl.count)
	sdl.matrices = make(map[int]*Matrix)
}

func UpsertSuccess(record history.Record) bool {
	return sdl.history.UpsertSuccess(record)
}

func DeleteSuccess(record history.Record) {
	sdl.history.DeleteSuccess(record)
}

func UpsertFailure(req *context.Request) bool {
	return sdl.history.UpsertFailure(req)
}

func DeleteFailure(req *context.Request) {
	sdl.history.DeleteFailure(req)
}

// 获取指定蜘蛛在上一次运行时失败的请求
func PullFailure(spiderName string) []*context.Request {
	return sdl.history.PullFailure(spiderName)
}

func TryFlushHistory() {
	if cache.Task.SuccessInherit {
		sdl.history.FlushSuccess(cache.Task.OutType)
	}
	if cache.Task.FailureInherit {
		sdl.history.FlushFailure(cache.Task.OutType)
	}
}

// 每个spider实例分配到的平均资源量
func (self *scheduler) avgRes() int32 {
	avg := int32(cap(sdl.count) / len(sdl.matrices))
	if avg == 0 {
		avg = 1
	}
	return avg
}

// 一个Spider实例的请求矩阵
type Matrix struct {
	// [优先级]队列，优先级默认为0
	reqs map[int][]*context.Request
	// 优先级顺序，从低到高
	priorities []int
	// 资源使用情况计数
	resCount int32
	// 最大采集页数，以负数形式表示
	maxPage int64
	// 历史及本次失败请求
	failures map[string]*context.Request
	sync.Mutex
}

// 注册资源队列
func NewMatrix(spiderId int, maxPage int64) *Matrix {
	sdl.Lock()
	defer sdl.Unlock()
	matrix := &Matrix{
		reqs:       make(map[int][]*context.Request),
		priorities: []int{},
		maxPage:    maxPage,
		failures:   make(map[string]*context.Request),
	}
	sdl.matrices[spiderId] = matrix
	return matrix
}

// 添加请求到队列
func (self *Matrix) Push(req *context.Request) {
	// 禁止并发，降低请求积存量
	self.Lock()
	defer self.Unlock()

	// 根据运行状态及资源使用情况，降低请求积存量
	for sdl.status == status.PAUSE || self.resCount > sdl.avgRes() {
		runtime.Gosched()
	}

	sdl.RLock()
	defer sdl.RUnlock()

	// 终止添加操作
	if sdl.status == status.STOP ||
		self.maxPage >= 0 ||
		// 当req不可重复下载时，已存在成功记录则返回
		!req.IsReloadable() && !UpsertSuccess(req) {
		return
	}

	// 大致限制加入队列的请求量，并发情况下应该会比maxPage多
	atomic.AddInt64(&self.maxPage, 1)

	priority := req.GetPriority()

	// 初始化该蜘蛛下该优先级队列
	if _, found := self.reqs[priority]; !found {
		self.priorities = append(self.priorities, priority)
		sort.Ints(self.priorities) // 从小到大排序
		self.reqs[priority] = []*context.Request{}
	}

	// 添加请求到队列
	self.reqs[priority] = append(self.reqs[priority], req)
}

// 从队列取出请求，不存在时返回nil
func (self *Matrix) Pull() (req *context.Request) {
	sdl.RLock()
	defer sdl.RUnlock()
	if sdl.status != status.RUN {
		return
	}

	// 按优先级从高到低取出请求
	for i := len(self.reqs) - 1; i >= 0; i-- {
		idx := self.priorities[i]
		if len(self.reqs[idx]) > 0 {
			req = self.reqs[idx][0]
			self.reqs[idx] = self.reqs[idx][1:]
			if sdl.useProxy {
				req.SetProxy(sdl.proxy.GetOne(req.GetUrl()))
			} else {
				req.SetProxy("")
			}
			return
		}
	}
	return
}

func (self *Matrix) Len() int {
	var l int
	for _, reqs := range self.reqs {
		l += len(reqs)
	}
	return l
}

func (self *Matrix) Use() {
	defer func() {
		recover()
	}()
	sdl.count <- true
	atomic.AddInt32(&self.resCount, 1)
}

func (self *Matrix) Free() {
	<-sdl.count
	atomic.AddInt32(&self.resCount, -1)
}

func (self *Matrix) CanStop() bool {
	sdl.RLock()
	if sdl.status == status.STOP {
		sdl.RUnlock()
		return true
	}
	sdl.RUnlock()

	if self.maxPage >= 0 {
		return true
	}

	if self.resCount != 0 {
		return false
	}

	for i, count := 0, len(self.reqs); i < count; i++ {
		if len(self.reqs[i]) > 0 {
			return false
		}
	}

	if len(self.failures) > 0 {
		// 重新下载历史记录中失败的请求
		var goon bool
		for unique, req := range self.failures {
			if req == nil {
				continue
			}
			self.failures[unique] = nil
			goon = true
			logs.Log.Informational(" *     - 失败请求: [%v]\n", req.GetUrl())
			self.Push(req)
		}
		if goon {
			return false
		}
	}
	return true
}

// 等待处理中的请求完成
func (self *Matrix) Flush() {
	for self.resCount != 0 {
		runtime.Gosched()
		// time.Sleep(5e8)
	}
}

func (self *Matrix) SetFailures(reqs []*context.Request) {
	for _, req := range reqs {
		self.failures[makeUnique(req)] = req
		logs.Log.Informational(" *     + 失败请求: [%v]\n", req.GetUrl())
	}
}

func (self *Matrix) SetFailure(req *context.Request) bool {
	self.Lock()
	defer self.Unlock()
	unique := makeUnique(req)
	if _, ok := self.failures[unique]; !ok {
		// 首次失败时，在任务队列末尾重新执行一次
		self.failures[unique] = req
		logs.Log.Informational(" *     + 失败请求: [%v]\n", req.GetUrl())
		return true
	}
	// 失败两次后，加入历史失败记录
	UpsertFailure(req)
	return false
}

func makeUnique(req *context.Request) string {
	return util.MakeUnique(req.GetUrl() + req.GetMethod())
}
<|endoftext|>


<commit_before>package distributor

import (
	"context"
	"flag"
	"hash/fnv"
	"sync/atomic"

	cortex_client "github.com/cortexproject/cortex/pkg/ingester/client"
	"github.com/cortexproject/cortex/pkg/ring"
	cortex_util "github.com/cortexproject/cortex/pkg/util"
	"github.com/cortexproject/cortex/pkg/util/validation"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/weaveworks/common/user"
	"google.golang.org/grpc/health/grpc_health_v1"

	"github.com/grafana/loki/pkg/ingester/client"
	"github.com/grafana/loki/pkg/logproto"
	"github.com/grafana/loki/pkg/util"
)

var (
	ingesterAppends = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_ingester_appends_total",
		Help:      "The total number of batch appends sent to ingesters.",
	}, []string{"ingester"})
	ingesterAppendFailures = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_ingester_append_failures_total",
		Help:      "The total number of failed batch appends sent to ingesters.",
	}, []string{"ingester"})

	bytesIngested = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_bytes_received_total",
		Help:      "The total number of uncompressed bytes received per tenant",
	}, []string{"tenant"})
	linesIngested = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_lines_received_total",
		Help:      "The total number of lines received per tenant",
	}, []string{"tenant"})
)

// Config for a Distributor.
type Config struct {
}

// RegisterFlags registers the flags.
func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
}

// Distributor coordinates replicates and distribution of log streams.
type Distributor struct {
	cfg       Config
	clientCfg client.Config
	ring      ring.ReadRing
	overrides *validation.Overrides
	pool      *cortex_client.Pool
}

// New a distributor creates.
func New(cfg Config, clientCfg client.Config, ring ring.ReadRing, overrides *validation.Overrides) (*Distributor, error) {
	factory := func(addr string) (grpc_health_v1.HealthClient, error) {
		return client.New(clientCfg, addr)
	}

	return &Distributor{
		cfg:       cfg,
		clientCfg: clientCfg,
		ring:      ring,
		overrides: overrides,
		pool:      cortex_client.NewPool(clientCfg.PoolConfig, ring, factory, cortex_util.Logger),
	}, nil
}

// TODO taken from Cortex, see if we can refactor out an usable interface.
type streamTracker struct {
	stream      *logproto.Stream
	minSuccess  int
	maxFailures int
	succeeded   int32
	failed      int32
}

// TODO taken from Cortex, see if we can refactor out an usable interface.
type pushTracker struct {
	samplesPending int32
	samplesFailed  int32
	done           chan struct{}
	err            chan error
}

// Push a set of streams.
func (d *Distributor) Push(ctx context.Context, req *logproto.PushRequest) (*logproto.PushResponse, error) {
	userID, err := user.ExtractOrgID(ctx)
	if err != nil {
		return nil, err
	}

	// Track metrics.
	bytesCount := 0
	lineCount := 0
	for _, stream := range req.Streams {
		for _, entry := range stream.Entries {
			bytesCount += len(entry.Line)
			lineCount++
		}
	}
	bytesIngested.WithLabelValues(userID).Add(float64(bytesCount))
	linesIngested.WithLabelValues(userID).Add(float64(lineCount))

	// First we flatten out the request into a list of samples.
	// We use the heuristic of 1 sample per TS to size the array.
	// We also work out the hash value at the same time.
	streams := make([]streamTracker, len(req.Streams))
	keys := make([]uint32, 0, len(req.Streams))
	var validationErr error
	for i, stream := range req.Streams {
		if err := d.validateLabels(userID, stream.Labels); err != nil {
			validationErr = err
			continue
		}

		keys = append(keys, tokenFor(userID, stream.Labels))
		streams[i].stream = stream
	}

	if len(streams) == 0 {
		return &logproto.PushResponse{}, nil
	}

	replicationSets, err := d.ring.BatchGet(keys, ring.Write)
	if err != nil {
		return nil, err
	}

	samplesByIngester := map[string][]*streamTracker{}
	ingesterDescs := map[string]ring.IngesterDesc{}
	for i, replicationSet := range replicationSets {
		streams[i].minSuccess = len(replicationSet.Ingesters) - replicationSet.MaxErrors
		streams[i].maxFailures = replicationSet.MaxErrors
		for _, ingester := range replicationSet.Ingesters {
			samplesByIngester[ingester.Addr] = append(samplesByIngester[ingester.Addr], &streams[i])
			ingesterDescs[ingester.Addr] = ingester
		}
	}

	tracker := pushTracker{
		samplesPending: int32(len(streams)),
		done:           make(chan struct{}),
		err:            make(chan error),
	}
	for ingester, samples := range samplesByIngester {
		go func(ingester ring.IngesterDesc, samples []*streamTracker) {
			// Use a background context to make sure all ingesters get samples even if we return early
			localCtx, cancel := context.WithTimeout(context.Background(), d.clientCfg.RemoteTimeout)
			defer cancel()
			localCtx = user.InjectOrgID(localCtx, userID)
			if sp := opentracing.SpanFromContext(ctx); sp != nil {
				localCtx = opentracing.ContextWithSpan(localCtx, sp)
			}
			d.sendSamples(localCtx, ingester, samples, &tracker)
		}(ingesterDescs[ingester], samples)
	}
	select {
	case err := <-tracker.err:
		return nil, err
	case <-tracker.done:
		return &logproto.PushResponse{}, validationErr
	}
}

func (d *Distributor) validateLabels(userID, labels string) error {
	ls, err := util.ToClientLabels(labels)
	if err != nil {
		return err
	}

	return d.overrides.ValidateLabels(userID, ls)
}

// TODO taken from Cortex, see if we can refactor out an usable interface.
func (d *Distributor) sendSamples(ctx context.Context, ingester ring.IngesterDesc, streamTrackers []*streamTracker, pushTracker *pushTracker) {
	err := d.sendSamplesErr(ctx, ingester, streamTrackers)

	// If we succeed, decrement each sample's pending count by one.  If we reach
	// the required number of successful puts on this sample, then decrement the
	// number of pending samples by one.  If we successfully push all samples to
	// min success ingesters, wake up the waiting rpc so it can return early.
	// Similarly, track the number of errors, and if it exceeds maxFailures
	// shortcut the waiting rpc.
	//
	// The use of atomic increments here guarantees only a single sendSamples
	// goroutine will write to either channel.
	for i := range streamTrackers {
		if err != nil {
			if atomic.AddInt32(&streamTrackers[i].failed, 1) <= int32(streamTrackers[i].maxFailures) {
				continue
			}
			if atomic.AddInt32(&pushTracker.samplesFailed, 1) == 1 {
				pushTracker.err <- err
			}
		} else {
			if atomic.AddInt32(&streamTrackers[i].succeeded, 1) != int32(streamTrackers[i].minSuccess) {
				continue
			}
			if atomic.AddInt32(&pushTracker.samplesPending, -1) == 0 {
				pushTracker.done <- struct{}{}
			}
		}
	}
}

// TODO taken from Cortex, see if we can refactor out an usable interface.
func (d *Distributor) sendSamplesErr(ctx context.Context, ingester ring.IngesterDesc, streams []*streamTracker) error {
	c, err := d.pool.GetClientFor(ingester.Addr)
	if err != nil {
		return err
	}

	req := &logproto.PushRequest{
		Streams: make([]*logproto.Stream, len(streams)),
	}
	for i, s := range streams {
		req.Streams[i] = s.stream
	}

	_, err = c.(logproto.PusherClient).Push(ctx, req)
	ingesterAppends.WithLabelValues(ingester.Addr).Inc()
	if err != nil {
		ingesterAppendFailures.WithLabelValues(ingester.Addr).Inc()
	}
	return err
}

func tokenFor(userID, labels string) uint32 {
	h := fnv.New32()
	h.Write([]byte(userID))
	h.Write([]byte(labels))
	return h.Sum32()
}

// Check implements the grpc healthcheck
func (*Distributor) Check(_ context.Context, _ *grpc_health_v1.HealthCheckRequest) (*grpc_health_v1.HealthCheckResponse, error) {
	return &grpc_health_v1.HealthCheckResponse{Status: grpc_health_v1.HealthCheckResponse_SERVING}, nil
}
<commit_msg>Wait for the right number of streams.<commit_after>package distributor

import (
	"context"
	"flag"
	"hash/fnv"
	"sync/atomic"

	cortex_client "github.com/cortexproject/cortex/pkg/ingester/client"
	"github.com/cortexproject/cortex/pkg/ring"
	cortex_util "github.com/cortexproject/cortex/pkg/util"
	"github.com/cortexproject/cortex/pkg/util/validation"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/weaveworks/common/user"
	"google.golang.org/grpc/health/grpc_health_v1"

	"github.com/grafana/loki/pkg/ingester/client"
	"github.com/grafana/loki/pkg/logproto"
	"github.com/grafana/loki/pkg/util"
)

var (
	ingesterAppends = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_ingester_appends_total",
		Help:      "The total number of batch appends sent to ingesters.",
	}, []string{"ingester"})
	ingesterAppendFailures = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_ingester_append_failures_total",
		Help:      "The total number of failed batch appends sent to ingesters.",
	}, []string{"ingester"})

	bytesIngested = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_bytes_received_total",
		Help:      "The total number of uncompressed bytes received per tenant",
	}, []string{"tenant"})
	linesIngested = promauto.NewCounterVec(prometheus.CounterOpts{
		Namespace: "loki",
		Name:      "distributor_lines_received_total",
		Help:      "The total number of lines received per tenant",
	}, []string{"tenant"})
)

// Config for a Distributor.
type Config struct {
}

// RegisterFlags registers the flags.
func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
}

// Distributor coordinates replicates and distribution of log streams.
type Distributor struct {
	cfg       Config
	clientCfg client.Config
	ring      ring.ReadRing
	overrides *validation.Overrides
	pool      *cortex_client.Pool
}

// New a distributor creates.
func New(cfg Config, clientCfg client.Config, ring ring.ReadRing, overrides *validation.Overrides) (*Distributor, error) {
	factory := func(addr string) (grpc_health_v1.HealthClient, error) {
		return client.New(clientCfg, addr)
	}

	return &Distributor{
		cfg:       cfg,
		clientCfg: clientCfg,
		ring:      ring,
		overrides: overrides,
		pool:      cortex_client.NewPool(clientCfg.PoolConfig, ring, factory, cortex_util.Logger),
	}, nil
}

// TODO taken from Cortex, see if we can refactor out an usable interface.
type streamTracker struct {
	stream      *logproto.Stream
	minSuccess  int
	maxFailures int
	succeeded   int32
	failed      int32
}

// TODO taken from Cortex, see if we can refactor out an usable interface.
type pushTracker struct {
	samplesPending int32
	samplesFailed  int32
	done           chan struct{}
	err            chan error
}

// Push a set of streams.
func (d *Distributor) Push(ctx context.Context, req *logproto.PushRequest) (*logproto.PushResponse, error) {
	userID, err := user.ExtractOrgID(ctx)
	if err != nil {
		return nil, err
	}

	// Track metrics.
	bytesCount := 0
	lineCount := 0
	for _, stream := range req.Streams {
		for _, entry := range stream.Entries {
			bytesCount += len(entry.Line)
			lineCount++
		}
	}
	bytesIngested.WithLabelValues(userID).Add(float64(bytesCount))
	linesIngested.WithLabelValues(userID).Add(float64(lineCount))

	// First we flatten out the request into a list of samples.
	// We use the heuristic of 1 sample per TS to size the array.
	// We also work out the hash value at the same time.
	streams := make([]streamTracker, 0, len(req.Streams))
	keys := make([]uint32, 0, len(req.Streams))
	var validationErr error
	for _, stream := range req.Streams {
		if err := d.validateLabels(userID, stream.Labels); err != nil {
			validationErr = err
			continue
		}

		keys = append(keys, tokenFor(userID, stream.Labels))
		streams = append(streams, streamTracker{
			stream: stream,
		})
	}

	if len(streams) == 0 {
		return &logproto.PushResponse{}, nil
	}

	replicationSets, err := d.ring.BatchGet(keys, ring.Write)
	if err != nil {
		return nil, err
	}

	samplesByIngester := map[string][]*streamTracker{}
	ingesterDescs := map[string]ring.IngesterDesc{}
	for i, replicationSet := range replicationSets {
		streams[i].minSuccess = len(replicationSet.Ingesters) - replicationSet.MaxErrors
		streams[i].maxFailures = replicationSet.MaxErrors
		for _, ingester := range replicationSet.Ingesters {
			samplesByIngester[ingester.Addr] = append(samplesByIngester[ingester.Addr], &streams[i])
			ingesterDescs[ingester.Addr] = ingester
		}
	}

	tracker := pushTracker{
		samplesPending: int32(len(samplesByIngester)),
		done:           make(chan struct{}),
		err:            make(chan error),
	}
	for ingester, samples := range samplesByIngester {
		go func(ingester ring.IngesterDesc, samples []*streamTracker) {
			// Use a background context to make sure all ingesters get samples even if we return early
			localCtx, cancel := context.WithTimeout(context.Background(), d.clientCfg.RemoteTimeout)
			defer cancel()
			localCtx = user.InjectOrgID(localCtx, userID)
			if sp := opentracing.SpanFromContext(ctx); sp != nil {
				localCtx = opentracing.ContextWithSpan(localCtx, sp)
			}
			d.sendSamples(localCtx, ingester, samples, &tracker)
		}(ingesterDescs[ingester], samples)
	}
	select {
	case err := <-tracker.err:
		return nil, err
	case <-tracker.done:
		return &logproto.PushResponse{}, validationErr
	}
}

func (d *Distributor) validateLabels(userID, labels string) error {
	ls, err := util.ToClientLabels(labels)
	if err != nil {
		return err
	}

	return d.overrides.ValidateLabels(userID, ls)
}

// TODO taken from Cortex, see if we can refactor out an usable interface.
func (d *Distributor) sendSamples(ctx context.Context, ingester ring.IngesterDesc, streamTrackers []*streamTracker, pushTracker *pushTracker) {
	err := d.sendSamplesErr(ctx, ingester, streamTrackers)

	// If we succeed, decrement each sample's pending count by one.  If we reach
	// the required number of successful puts on this sample, then decrement the
	// number of pending samples by one.  If we successfully push all samples to
	// min success ingesters, wake up the waiting rpc so it can return early.
	// Similarly, track the number of errors, and if it exceeds maxFailures
	// shortcut the waiting rpc.
	//
	// The use of atomic increments here guarantees only a single sendSamples
	// goroutine will write to either channel.
	for i := range streamTrackers {
		if err != nil {
			if atomic.AddInt32(&streamTrackers[i].failed, 1) <= int32(streamTrackers[i].maxFailures) {
				continue
			}
			if atomic.AddInt32(&pushTracker.samplesFailed, 1) == 1 {
				pushTracker.err <- err
			}
		} else {
			if atomic.AddInt32(&streamTrackers[i].succeeded, 1) != int32(streamTrackers[i].minSuccess) {
				continue
			}
			if atomic.AddInt32(&pushTracker.samplesPending, -1) == 0 {
				pushTracker.done <- struct{}{}
			}
		}
	}
}

// TODO taken from Cortex, see if we can refactor out an usable interface.
func (d *Distributor) sendSamplesErr(ctx context.Context, ingester ring.IngesterDesc, streams []*streamTracker) error {
	c, err := d.pool.GetClientFor(ingester.Addr)
	if err != nil {
		return err
	}

	req := &logproto.PushRequest{
		Streams: make([]*logproto.Stream, len(streams)),
	}
	for i, s := range streams {
		req.Streams[i] = s.stream
	}

	_, err = c.(logproto.PusherClient).Push(ctx, req)
	ingesterAppends.WithLabelValues(ingester.Addr).Inc()
	if err != nil {
		ingesterAppendFailures.WithLabelValues(ingester.Addr).Inc()
	}
	return err
}

func tokenFor(userID, labels string) uint32 {
	h := fnv.New32()
	h.Write([]byte(userID))
	h.Write([]byte(labels))
	return h.Sum32()
}

// Check implements the grpc healthcheck
func (*Distributor) Check(_ context.Context, _ *grpc_health_v1.HealthCheckRequest) (*grpc_health_v1.HealthCheckResponse, error) {
	return &grpc_health_v1.HealthCheckResponse{Status: grpc_health_v1.HealthCheckResponse_SERVING}, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 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 config

import (
	"errors"
	"fmt"
	"io"

	"github.com/renstrom/dedent"
	"github.com/spf13/cobra"

	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	clientcmdapi "k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api/latest"
	"k8s.io/kubernetes/pkg/kubectl"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/util/flag"
)

type ViewOptions struct {
	ConfigAccess clientcmd.ConfigAccess
	Merge        flag.Tristate
	Flatten      bool
	Minify       bool
	RawByteData  bool
}

var (
	view_long = dedent.Dedent(`
		Display merged kubeconfig settings or a specified kubeconfig file.

		You can use --output jsonpath={...} to extract specific values using a jsonpath expression.`)
	view_example = dedent.Dedent(`
		# Show Merged kubeconfig settings.
		kubectl config view

		# Get the password for the e2e user
		kubectl config view -o jsonpath='{.users[?(@.name == "e2e")].user.password}'`)
)

func NewCmdConfigView(out io.Writer, ConfigAccess clientcmd.ConfigAccess) *cobra.Command {
	options := &ViewOptions{ConfigAccess: ConfigAccess}
	// Default to yaml
	defaultOutputFormat := "yaml"

	cmd := &cobra.Command{
		Use:     "view",
		Short:   "Display merged kubeconfig settings or a specified kubeconfig file",
		Long:    view_long,
		Example: view_example,
		Run: func(cmd *cobra.Command, args []string) {
			options.Complete()
			outputFormat := cmdutil.GetFlagString(cmd, "output")
			if outputFormat == "wide" {
				fmt.Printf("--output wide is not available in kubectl config view; reset to default output format (%s)\n\n", defaultOutputFormat)
				cmd.Flags().Set("output", defaultOutputFormat)
			}
			if outputFormat == "" {
				fmt.Printf("reset to default output format (%s) as --output is empty", defaultOutputFormat)
				cmd.Flags().Set("output", defaultOutputFormat)
			}

			printer, _, err := cmdutil.PrinterForCommand(cmd)
			cmdutil.CheckErr(err)
			version, err := cmdutil.OutputVersion(cmd, &latest.ExternalVersion)
			cmdutil.CheckErr(err)
			printer = kubectl.NewVersionedPrinter(printer, latest.Scheme, version)

			cmdutil.CheckErr(options.Run(out, printer))
		},
	}

	cmdutil.AddPrinterFlags(cmd)
	cmd.Flags().Set("output", defaultOutputFormat)

	options.Merge.Default(true)
	f := cmd.Flags().VarPF(&options.Merge, "merge", "", "merge the full hierarchy of kubeconfig files")
	f.NoOptDefVal = "true"
	cmd.Flags().BoolVar(&options.RawByteData, "raw", false, "display raw byte data")
	cmd.Flags().BoolVar(&options.Flatten, "flatten", false, "flatten the resulting kubeconfig file into self-contained output (useful for creating portable kubeconfig files)")
	cmd.Flags().BoolVar(&options.Minify, "minify", false, "remove all information not used by current-context from the output")
	return cmd
}

func (o ViewOptions) Run(out io.Writer, printer kubectl.ResourcePrinter) error {
	config, err := o.loadConfig()
	if err != nil {
		return err
	}

	if o.Minify {
		if err := clientcmdapi.MinifyConfig(config); err != nil {
			return err
		}
	}

	if o.Flatten {
		if err := clientcmdapi.FlattenConfig(config); err != nil {
			return err
		}
	} else if !o.RawByteData {
		clientcmdapi.ShortenConfig(config)
	}

	err = printer.PrintObj(config, out)
	if err != nil {
		return err
	}

	return nil
}

func (o *ViewOptions) Complete() bool {
	if o.ConfigAccess.IsExplicitFile() {
		if !o.Merge.Provided() {
			o.Merge.Set("false")
		}
	}

	return true
}

func (o ViewOptions) loadConfig() (*clientcmdapi.Config, error) {
	err := o.Validate()
	if err != nil {
		return nil, err
	}

	config, err := o.getStartingConfig()
	return config, err
}

func (o ViewOptions) Validate() error {
	if !o.Merge.Value() && !o.ConfigAccess.IsExplicitFile() {
		return errors.New("if merge==false a precise file must to specified")
	}

	return nil
}

// getStartingConfig returns the Config object built from the sources specified by the options, the filename read (only if it was a single file), and an error if something goes wrong
func (o *ViewOptions) getStartingConfig() (*clientcmdapi.Config, error) {
	switch {
	case !o.Merge.Value():
		return clientcmd.LoadFromFile(o.ConfigAccess.GetExplicitFile())

	default:
		return o.ConfigAccess.GetStartingConfig()
	}
}
<commit_msg>It is better to add "\n" in printf<commit_after>/*
Copyright 2014 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 config

import (
	"errors"
	"fmt"
	"io"

	"github.com/renstrom/dedent"
	"github.com/spf13/cobra"

	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	clientcmdapi "k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api/latest"
	"k8s.io/kubernetes/pkg/kubectl"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/util/flag"
)

type ViewOptions struct {
	ConfigAccess clientcmd.ConfigAccess
	Merge        flag.Tristate
	Flatten      bool
	Minify       bool
	RawByteData  bool
}

var (
	view_long = dedent.Dedent(`
		Display merged kubeconfig settings or a specified kubeconfig file.

		You can use --output jsonpath={...} to extract specific values using a jsonpath expression.`)
	view_example = dedent.Dedent(`
		# Show Merged kubeconfig settings.
		kubectl config view

		# Get the password for the e2e user
		kubectl config view -o jsonpath='{.users[?(@.name == "e2e")].user.password}'`)
)

func NewCmdConfigView(out io.Writer, ConfigAccess clientcmd.ConfigAccess) *cobra.Command {
	options := &ViewOptions{ConfigAccess: ConfigAccess}
	// Default to yaml
	defaultOutputFormat := "yaml"

	cmd := &cobra.Command{
		Use:     "view",
		Short:   "Display merged kubeconfig settings or a specified kubeconfig file",
		Long:    view_long,
		Example: view_example,
		Run: func(cmd *cobra.Command, args []string) {
			options.Complete()
			outputFormat := cmdutil.GetFlagString(cmd, "output")
			if outputFormat == "wide" {
				fmt.Printf("--output wide is not available in kubectl config view; reset to default output format (%s)\n\n", defaultOutputFormat)
				cmd.Flags().Set("output", defaultOutputFormat)
			}
			if outputFormat == "" {
				fmt.Printf("Reset to default output format (%s) as --output is empty\n", defaultOutputFormat)
				cmd.Flags().Set("output", defaultOutputFormat)
			}

			printer, _, err := cmdutil.PrinterForCommand(cmd)
			cmdutil.CheckErr(err)
			version, err := cmdutil.OutputVersion(cmd, &latest.ExternalVersion)
			cmdutil.CheckErr(err)
			printer = kubectl.NewVersionedPrinter(printer, latest.Scheme, version)

			cmdutil.CheckErr(options.Run(out, printer))
		},
	}

	cmdutil.AddPrinterFlags(cmd)
	cmd.Flags().Set("output", defaultOutputFormat)

	options.Merge.Default(true)
	f := cmd.Flags().VarPF(&options.Merge, "merge", "", "merge the full hierarchy of kubeconfig files")
	f.NoOptDefVal = "true"
	cmd.Flags().BoolVar(&options.RawByteData, "raw", false, "display raw byte data")
	cmd.Flags().BoolVar(&options.Flatten, "flatten", false, "flatten the resulting kubeconfig file into self-contained output (useful for creating portable kubeconfig files)")
	cmd.Flags().BoolVar(&options.Minify, "minify", false, "remove all information not used by current-context from the output")
	return cmd
}

func (o ViewOptions) Run(out io.Writer, printer kubectl.ResourcePrinter) error {
	config, err := o.loadConfig()
	if err != nil {
		return err
	}

	if o.Minify {
		if err := clientcmdapi.MinifyConfig(config); err != nil {
			return err
		}
	}

	if o.Flatten {
		if err := clientcmdapi.FlattenConfig(config); err != nil {
			return err
		}
	} else if !o.RawByteData {
		clientcmdapi.ShortenConfig(config)
	}

	err = printer.PrintObj(config, out)
	if err != nil {
		return err
	}

	return nil
}

func (o *ViewOptions) Complete() bool {
	if o.ConfigAccess.IsExplicitFile() {
		if !o.Merge.Provided() {
			o.Merge.Set("false")
		}
	}

	return true
}

func (o ViewOptions) loadConfig() (*clientcmdapi.Config, error) {
	err := o.Validate()
	if err != nil {
		return nil, err
	}

	config, err := o.getStartingConfig()
	return config, err
}

func (o ViewOptions) Validate() error {
	if !o.Merge.Value() && !o.ConfigAccess.IsExplicitFile() {
		return errors.New("if merge==false a precise file must to specified")
	}

	return nil
}

// getStartingConfig returns the Config object built from the sources specified by the options, the filename read (only if it was a single file), and an error if something goes wrong
func (o *ViewOptions) getStartingConfig() (*clientcmdapi.Config, error) {
	switch {
	case !o.Merge.Value():
		return clientcmd.LoadFromFile(o.ConfigAccess.GetExplicitFile())

	default:
		return o.ConfigAccess.GetStartingConfig()
	}
}
<|endoftext|>


<commit_before>/*
 * This file is part of the Dicot project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package v3

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"

	"github.com/dicot-project/dicot-api/pkg/api/identity"
	"github.com/dicot-project/dicot-api/pkg/api/identity/v1"
	"github.com/dicot-project/dicot-api/pkg/auth"
	"github.com/dicot-project/dicot-api/pkg/crypto"
)

type AuthReq struct {
	Auth AuthInfo `json:"auth"`
}

type AuthInfo struct {
	Scope    AuthInfoScope    `json:"scope"`
	Identity AuthInfoIdentity `json:"identity"`
}

type AuthInfoScope struct {
	Project ProjectInfoRef `json:"project"`
}

type ProjectInfoRef struct {
	ID     string        `json:"id"`
	Name   string        `json:"name"`
	Domain DomainInfoRef `json:"domain"`
}

type DomainInfoRef struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

type AuthInfoIdentity struct {
	Methods  []string         `json:"methods"`
	Password AuthInfoPassword `json:"password"`
	Token    AuthInfoToken    `json:"token"`
}

type AuthInfoToken struct {
	ID string `json:"id"`
}

type AuthInfoPassword struct {
	User UserInfoRef `json:"user"`
}

type TokenRes struct {
	Token TokenInfo `json:"token"`
}

type TokenInfo struct {
	Methods   []string           `json:"methods"`
	Roles     []RoleInfo         `json:"roles"`
	ExpiresAt string             `json:"expires_at"`
	IssuedAt  string             `json:"issued_at"`
	Project   ProjectInfoRef     `json:"project"`
	IsDomain  bool               `json:"is_domain"`
	Catalogs  []TokenInfoCatalog `json:"catalog"`
	User      UserInfoRef        `json:"user"`
	AuditIDs  []string           `json:"audit_ids"`
	Extras    map[string]string  `json:"extras"`
}

type TokenInfoCatalog struct {
	ID        string              `json:"id"`
	Endpoints []TokenInfoEndpoint `json:"endpoints"`
	Type      string              `json:"type"`
	Name      string              `json:"name"`
}

type TokenInfoEndpoint struct {
	ID        string `json:"id"`
	Region    string `json:"region"`
	RegionID  string `json:"region_id"`
	URL       string `json:"url"`
	Interface string `json:"interface"`
}

type RoleInfo struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

type UserInfoRef struct {
	ID                string        `json:"id"`
	Name              string        `json:"name"`
	Domain            DomainInfoRef `json:"domain"`
	Password          string        `json:"password"`
	PasswordExpiresAt string        `json:"password_expires_at"`
}

func (svc *service) TokensPost(c *gin.Context) {
	var req AuthReq
	err := c.BindJSON(&req)
	if err != nil {
		c.AbortWithStatus(http.StatusBadRequest)
		return
	}

	pwAuth := false
	for _, val := range req.Auth.Identity.Methods {
		if val == "password" {
			pwAuth = true
			break
		}
	}
	if !pwAuth {
		c.AbortWithStatus(http.StatusUnauthorized)
		return
	}

	domClnt := identity.NewProjectClient(svc.IdentityClient, v1.NamespaceSystem)
	var domain *v1.Project
	if req.Auth.Identity.Password.User.Domain.Name != "" {
		domain, err = domClnt.Get(req.Auth.Identity.Password.User.Domain.Name)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
	} else if req.Auth.Identity.Password.User.Domain.ID != "" {
		domain, err = domClnt.GetByUID(req.Auth.Identity.Password.User.Domain.ID)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
	} else {
		c.AbortWithStatus(http.StatusUnauthorized)
		return
	}
	namespace := identity.FormatDomainNamespace(domain.ObjectMeta.Name)

	userClnt := identity.NewUserClient(svc.IdentityClient, namespace)

	var user *v1.User
	if req.Auth.Identity.Password.User.Name != "" {
		user, err = userClnt.Get(req.Auth.Identity.Password.User.Name)
	} else {
		user, err = userClnt.GetByUID(req.Auth.Identity.Password.User.ID)
	}
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}

	secret, err := svc.K8SClient.CoreV1().Secrets(namespace).Get(user.Spec.Password.SecretRef, metav1.GetOptions{})
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}

	allowed, err := crypto.CheckPassword(
		req.Auth.Identity.Password.User.Password,
		string(secret.Data["password"]))
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}
	if !allowed {
		c.AbortWithStatus(http.StatusUnauthorized)
		return
	}

	token := svc.TokenManager.NewToken()
	token.Subject = auth.TokenSubject{
		DomainName: domain.ObjectMeta.Name,
		UserName:   user.ObjectMeta.Name,
	}
	token.Scope = auth.TokenScope{
		DomainName:  "",
		ProjectName: "",
	}

	tokensig, err := svc.TokenManager.SignToken(token)
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}

	catalog := []TokenInfoCatalog{}

	interfaces := []string{
		"internal", "admin", "public",
	}

	for _, service := range svc.Services.Services {
		endpoints := []TokenInfoEndpoint{}

		for _, iface := range interfaces {
			endpoints = append(endpoints, TokenInfoEndpoint{
				ID:        "4e7639cf-f78f-4cd2-aa2a-131196e25974",
				URL:       "http://" + c.Request.Host + service.GetPrefix() + "/",
				Region:    "RegionOne",
				RegionID:  "d3fd5ef9-7eff-422a-8df1-f2bc523d3381",
				Interface: iface,
			})
		}

		catalog = append(catalog, TokenInfoCatalog{
			ID:        service.GetUID(),
			Type:      service.GetType(),
			Name:      service.GetName(),
			Endpoints: endpoints,
		})
	}

	res := &TokenRes{
		Token: TokenInfo{
			Methods: []string{"password"},
			Roles: []RoleInfo{
				RoleInfo{
					ID:   "f56be11a-94a7-11e7-9f6d-e4b318e0afce",
					Name: "admin",
				},
			},
			IssuedAt:  token.Issued.Format(time.RFC3339),
			ExpiresAt: token.Expiry.Format(time.RFC3339),
			IsDomain:  false,
			AuditIDs: []string{
				string(uuid.NewUUID()),
			},
			Project: ProjectInfoRef{
				Domain: DomainInfoRef{
					ID:   "f4ae7bf2-94a7-11e7-b158-e4b318e0afce",
					Name: "default",
				},
				ID:   "324cd174-94a9-11e7-a705-e4b318e0afce",
				Name: "demo",
			},
			User: UserInfoRef{
				Domain: DomainInfoRef{
					ID:   string(domain.ObjectMeta.UID),
					Name: domain.ObjectMeta.Name,
				},
				ID:                string(user.ObjectMeta.UID),
				Name:              user.Spec.Name,
				PasswordExpiresAt: time.Now().Add(10 * time.Minute).Format(time.RFC3339),
			},
			Extras: map[string]string{
				"fish": "food",
			},
			Catalogs: catalog,
		},
	}
	c.Header("X-Subject-Token", tokensig)
	c.JSON(http.StatusOK, res)
}
<commit_msg>Resolve project scope when authenticating<commit_after>/*
 * This file is part of the Dicot project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package v3

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"

	"github.com/dicot-project/dicot-api/pkg/api/identity"
	"github.com/dicot-project/dicot-api/pkg/api/identity/v1"
	"github.com/dicot-project/dicot-api/pkg/auth"
	"github.com/dicot-project/dicot-api/pkg/crypto"
)

type AuthReq struct {
	Auth AuthInfo `json:"auth"`
}

type AuthInfo struct {
	Scope    AuthInfoScope    `json:"scope"`
	Identity AuthInfoIdentity `json:"identity"`
}

type AuthInfoScope struct {
	Project ProjectInfoRef `json:"project"`
}

type ProjectInfoRef struct {
	ID     string        `json:"id"`
	Name   string        `json:"name"`
	Domain DomainInfoRef `json:"domain"`
}

type DomainInfoRef struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

type AuthInfoIdentity struct {
	Methods  []string         `json:"methods"`
	Password AuthInfoPassword `json:"password"`
	Token    AuthInfoToken    `json:"token"`
}

type AuthInfoToken struct {
	ID string `json:"id"`
}

type AuthInfoPassword struct {
	User UserInfoRef `json:"user"`
}

type TokenRes struct {
	Token TokenInfo `json:"token"`
}

type TokenInfo struct {
	Methods   []string           `json:"methods"`
	Roles     []RoleInfo         `json:"roles"`
	ExpiresAt string             `json:"expires_at"`
	IssuedAt  string             `json:"issued_at"`
	Project   ProjectInfoRef     `json:"project"`
	IsDomain  bool               `json:"is_domain"`
	Catalogs  []TokenInfoCatalog `json:"catalog"`
	User      UserInfoRef        `json:"user"`
	AuditIDs  []string           `json:"audit_ids"`
	Extras    map[string]string  `json:"extras"`
}

type TokenInfoCatalog struct {
	ID        string              `json:"id"`
	Endpoints []TokenInfoEndpoint `json:"endpoints"`
	Type      string              `json:"type"`
	Name      string              `json:"name"`
}

type TokenInfoEndpoint struct {
	ID        string `json:"id"`
	Region    string `json:"region"`
	RegionID  string `json:"region_id"`
	URL       string `json:"url"`
	Interface string `json:"interface"`
}

type RoleInfo struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

type UserInfoRef struct {
	ID                string        `json:"id"`
	Name              string        `json:"name"`
	Domain            DomainInfoRef `json:"domain"`
	Password          string        `json:"password"`
	PasswordExpiresAt string        `json:"password_expires_at"`
}

func (svc *service) TokensPost(c *gin.Context) {
	var req AuthReq
	err := c.BindJSON(&req)
	if err != nil {
		c.AbortWithStatus(http.StatusBadRequest)
		return
	}

	pwAuth := false
	for _, val := range req.Auth.Identity.Methods {
		if val == "password" {
			pwAuth = true
			break
		}
	}
	if !pwAuth {
		c.AbortWithStatus(http.StatusUnauthorized)
		return
	}

	domClnt := identity.NewProjectClient(svc.IdentityClient, v1.NamespaceSystem)
	var userDomain *v1.Project
	if req.Auth.Identity.Password.User.Domain.Name != "" {
		userDomain, err = domClnt.Get(req.Auth.Identity.Password.User.Domain.Name)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
	} else if req.Auth.Identity.Password.User.Domain.ID != "" {
		userDomain, err = domClnt.GetByUID(req.Auth.Identity.Password.User.Domain.ID)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
	} else {
		c.AbortWithStatus(http.StatusUnauthorized)
		return
	}
	userNamespace := identity.FormatDomainNamespace(userDomain.ObjectMeta.Name)

	userClnt := identity.NewUserClient(svc.IdentityClient, userNamespace)

	var user *v1.User
	if req.Auth.Identity.Password.User.Name != "" {
		user, err = userClnt.Get(req.Auth.Identity.Password.User.Name)
	} else {
		user, err = userClnt.GetByUID(req.Auth.Identity.Password.User.ID)
	}
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}

	secret, err := svc.K8SClient.CoreV1().Secrets(userNamespace).Get(user.Spec.Password.SecretRef, metav1.GetOptions{})
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}

	allowed, err := crypto.CheckPassword(
		req.Auth.Identity.Password.User.Password,
		string(secret.Data["password"]))
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}
	if !allowed {
		c.AbortWithStatus(http.StatusUnauthorized)
		return
	}

	var projectDomain *v1.Project
	if req.Auth.Scope.Project.Domain.Name != "" {
		projectDomain, err = domClnt.Get(req.Auth.Scope.Project.Domain.Name)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
	} else if req.Auth.Scope.Project.Domain.ID != "" {
		projectDomain, err = domClnt.GetByUID(req.Auth.Scope.Project.Domain.ID)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
	} else {
		c.AbortWithStatus(http.StatusUnauthorized)
		return
	}
	projectNamespace := identity.FormatDomainNamespace(projectDomain.ObjectMeta.Name)

	projectClnt := identity.NewProjectClient(svc.IdentityClient, projectNamespace)

	var project *v1.Project
	if req.Auth.Scope.Project.Name != "" {
		project, err = projectClnt.Get(req.Auth.Scope.Project.Name)
	} else {
		project, err = projectClnt.GetByUID(req.Auth.Scope.Project.ID)
	}
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}

	// XXX validate user's access to the requested project

	token := svc.TokenManager.NewToken()
	token.Subject = auth.TokenSubject{
		DomainName: userDomain.ObjectMeta.Name,
		UserName:   user.ObjectMeta.Name,
	}
	token.Scope = auth.TokenScope{
		DomainName:  projectDomain.ObjectMeta.Name,
		ProjectName: project.ObjectMeta.Name,
	}

	tokensig, err := svc.TokenManager.SignToken(token)
	if err != nil {
		c.AbortWithError(http.StatusUnauthorized, err)
		return
	}

	catalog := []TokenInfoCatalog{}

	interfaces := []string{
		"internal", "admin", "public",
	}

	for _, service := range svc.Services.Services {
		endpoints := []TokenInfoEndpoint{}

		for _, iface := range interfaces {
			endpoints = append(endpoints, TokenInfoEndpoint{
				ID:        "4e7639cf-f78f-4cd2-aa2a-131196e25974",
				URL:       "http://" + c.Request.Host + service.GetPrefix() + "/",
				Region:    "RegionOne",
				RegionID:  "d3fd5ef9-7eff-422a-8df1-f2bc523d3381",
				Interface: iface,
			})
		}

		catalog = append(catalog, TokenInfoCatalog{
			ID:        service.GetUID(),
			Type:      service.GetType(),
			Name:      service.GetName(),
			Endpoints: endpoints,
		})
	}

	res := &TokenRes{
		Token: TokenInfo{
			Methods: []string{"password"},
			Roles: []RoleInfo{
				RoleInfo{
					ID:   "f56be11a-94a7-11e7-9f6d-e4b318e0afce",
					Name: "admin",
				},
			},
			IssuedAt:  token.Issued.Format(time.RFC3339),
			ExpiresAt: token.Expiry.Format(time.RFC3339),
			IsDomain:  false,
			AuditIDs: []string{
				string(uuid.NewUUID()),
			},
			Project: ProjectInfoRef{
				Domain: DomainInfoRef{
					ID:   string(projectDomain.ObjectMeta.UID),
					Name: projectDomain.ObjectMeta.Name,
				},
				ID:   string(project.ObjectMeta.UID),
				Name: project.ObjectMeta.Name,
			},
			User: UserInfoRef{
				Domain: DomainInfoRef{
					ID:   string(userDomain.ObjectMeta.UID),
					Name: userDomain.ObjectMeta.Name,
				},
				ID:                string(user.ObjectMeta.UID),
				Name:              user.Spec.Name,
				PasswordExpiresAt: time.Now().Add(10 * time.Minute).Format(time.RFC3339),
			},
			Extras: map[string]string{
				"fish": "food",
			},
			Catalogs: catalog,
		},
	}
	c.Header("X-Subject-Token", tokensig)
	c.JSON(http.StatusOK, res)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 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 mux

import (
	"fmt"
	"net/http"
	"runtime/debug"
	"sort"
	"strings"
	"sync"
	"sync/atomic"

	"k8s.io/klog/v2"

	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
)

// PathRecorderMux wraps a mux object and records the registered exposedPaths.
type PathRecorderMux struct {
	// name is used for logging so you can trace requests through
	name string

	lock            sync.Mutex
	notFoundHandler http.Handler
	pathToHandler   map[string]http.Handler
	prefixToHandler map[string]http.Handler

	// mux stores a pathHandler and is used to handle the actual serving.
	// Turns out, we want to accept trailing slashes, BUT we don't care about handling
	// everything under them.  This does exactly matches only unless its explicitly requested to
	// do something different
	mux atomic.Value

	// exposedPaths is the list of paths that should be shown at /
	exposedPaths []string

	// pathStacks holds the stacks of all registered paths.  This allows us to show a more helpful message
	// before the "http: multiple registrations for %s" panic.
	pathStacks map[string]string
}

// pathHandler is an http.Handler that will satisfy requests first by exact match, then by prefix,
// then by notFoundHandler
type pathHandler struct {
	// muxName is used for logging so you can trace requests through
	muxName string

	// pathToHandler is a map of exactly matching request to its handler
	pathToHandler map[string]http.Handler

	// this has to be sorted by most slashes then by length
	prefixHandlers []prefixHandler

	// notFoundHandler is the handler to use for satisfying requests with no other match
	notFoundHandler http.Handler
}

// prefixHandler holds the prefix it should match and the handler to use
type prefixHandler struct {
	// prefix is the prefix to test for a request match
	prefix string
	// handler is used to satisfy matching requests
	handler http.Handler
}

// NewPathRecorderMux creates a new PathRecorderMux
func NewPathRecorderMux(name string) *PathRecorderMux {
	ret := &PathRecorderMux{
		name:            name,
		pathToHandler:   map[string]http.Handler{},
		prefixToHandler: map[string]http.Handler{},
		mux:             atomic.Value{},
		exposedPaths:    []string{},
		pathStacks:      map[string]string{},
	}

	ret.mux.Store(&pathHandler{notFoundHandler: http.NotFoundHandler()})
	return ret
}

// ListedPaths returns the registered handler exposedPaths.
func (m *PathRecorderMux) ListedPaths() []string {
	handledPaths := append([]string{}, m.exposedPaths...)
	sort.Strings(handledPaths)

	return handledPaths
}

func (m *PathRecorderMux) trackCallers(path string) {
	if existingStack, ok := m.pathStacks[path]; ok {
		utilruntime.HandleError(fmt.Errorf("registered %q from %v", path, existingStack))
	}
	m.pathStacks[path] = string(debug.Stack())
}

// refreshMuxLocked creates a new mux and must be called while locked.  Otherwise the view of handlers may
// not be consistent
func (m *PathRecorderMux) refreshMuxLocked() {
	newMux := &pathHandler{
		muxName:         m.name,
		pathToHandler:   map[string]http.Handler{},
		prefixHandlers:  []prefixHandler{},
		notFoundHandler: http.NotFoundHandler(),
	}
	if m.notFoundHandler != nil {
		newMux.notFoundHandler = m.notFoundHandler
	}
	for path, handler := range m.pathToHandler {
		newMux.pathToHandler[path] = handler
	}

	keys := sets.StringKeySet(m.prefixToHandler).List()
	sort.Sort(sort.Reverse(byPrefixPriority(keys)))
	for _, prefix := range keys {
		newMux.prefixHandlers = append(newMux.prefixHandlers, prefixHandler{
			prefix:  prefix,
			handler: m.prefixToHandler[prefix],
		})
	}

	m.mux.Store(newMux)
}

// NotFoundHandler sets the handler to use if there's no match for a give path
func (m *PathRecorderMux) NotFoundHandler(notFoundHandler http.Handler) {
	m.lock.Lock()
	defer m.lock.Unlock()

	m.notFoundHandler = notFoundHandler

	m.refreshMuxLocked()
}

// Unregister removes a path from the mux.
func (m *PathRecorderMux) Unregister(path string) {
	m.lock.Lock()
	defer m.lock.Unlock()

	delete(m.pathToHandler, path)
	delete(m.prefixToHandler, path)
	delete(m.pathStacks, path)
	for i := range m.exposedPaths {
		if m.exposedPaths[i] == path {
			m.exposedPaths = append(m.exposedPaths[:i], m.exposedPaths[i+1:]...)
			break
		}
	}

	m.refreshMuxLocked()
}

// Handle registers the handler for the given pattern.
// If a handler already exists for pattern, Handle panics.
func (m *PathRecorderMux) Handle(path string, handler http.Handler) {
	m.lock.Lock()
	defer m.lock.Unlock()
	m.trackCallers(path)

	m.exposedPaths = append(m.exposedPaths, path)
	m.pathToHandler[path] = handler
	m.refreshMuxLocked()
}

// HandleFunc registers the handler function for the given pattern.
// If a handler already exists for pattern, Handle panics.
func (m *PathRecorderMux) HandleFunc(path string, handler func(http.ResponseWriter, *http.Request)) {
	m.Handle(path, http.HandlerFunc(handler))
}

// UnlistedHandle registers the handler for the given pattern, but doesn't list it.
// If a handler already exists for pattern, Handle panics.
func (m *PathRecorderMux) UnlistedHandle(path string, handler http.Handler) {
	m.lock.Lock()
	defer m.lock.Unlock()
	m.trackCallers(path)

	m.pathToHandler[path] = handler
	m.refreshMuxLocked()
}

// UnlistedHandleFunc registers the handler function for the given pattern, but doesn't list it.
// If a handler already exists for pattern, Handle panics.
func (m *PathRecorderMux) UnlistedHandleFunc(path string, handler func(http.ResponseWriter, *http.Request)) {
	m.UnlistedHandle(path, http.HandlerFunc(handler))
}

// HandlePrefix is like Handle, but matches for anything under the path.  Like a standard golang trailing slash.
func (m *PathRecorderMux) HandlePrefix(path string, handler http.Handler) {
	if !strings.HasSuffix(path, "/") {
		panic(fmt.Sprintf("%q must end in a trailing slash", path))
	}

	m.lock.Lock()
	defer m.lock.Unlock()
	m.trackCallers(path)

	m.exposedPaths = append(m.exposedPaths, path)
	m.prefixToHandler[path] = handler
	m.refreshMuxLocked()
}

// UnlistedHandlePrefix is like UnlistedHandle, but matches for anything under the path.  Like a standard golang trailing slash.
func (m *PathRecorderMux) UnlistedHandlePrefix(path string, handler http.Handler) {
	if !strings.HasSuffix(path, "/") {
		panic(fmt.Sprintf("%q must end in a trailing slash", path))
	}

	m.lock.Lock()
	defer m.lock.Unlock()
	m.trackCallers(path)

	m.prefixToHandler[path] = handler
	m.refreshMuxLocked()
}

// ServeHTTP makes it an http.Handler
func (m *PathRecorderMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	m.mux.Load().(*pathHandler).ServeHTTP(w, r)
}

// ServeHTTP makes it an http.Handler
func (h *pathHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if exactHandler, ok := h.pathToHandler[r.URL.Path]; ok {
		klog.V(5).Infof("%v: %q satisfied by exact match", h.muxName, r.URL.Path)
		exactHandler.ServeHTTP(w, r)
		return
	}

	for _, prefixHandler := range h.prefixHandlers {
		if strings.HasPrefix(r.URL.Path, prefixHandler.prefix) {
			klog.V(5).Infof("%v: %q satisfied by prefix %v", h.muxName, r.URL.Path, prefixHandler.prefix)
			prefixHandler.handler.ServeHTTP(w, r)
			return
		}
	}

	klog.V(5).Infof("%v: %q satisfied by NotFoundHandler", h.muxName, r.URL.Path)
	h.notFoundHandler.ServeHTTP(w, r)
}

// byPrefixPriority sorts url prefixes by the order in which they should be tested by the mux
// this has to be sorted by most slashes then by length so that we can iterate straight
// through to match the "best" one first.
type byPrefixPriority []string

func (s byPrefixPriority) Len() int      { return len(s) }
func (s byPrefixPriority) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s byPrefixPriority) Less(i, j int) bool {
	lhsNumParts := strings.Count(s[i], "/")
	rhsNumParts := strings.Count(s[j], "/")
	if lhsNumParts != rhsNumParts {
		return lhsNumParts < rhsNumParts
	}

	lhsLen := len(s[i])
	rhsLen := len(s[j])
	if lhsLen != rhsLen {
		return lhsLen < rhsLen
	}

	return strings.Compare(s[i], s[j]) < 0
}
<commit_msg>Improve pathrecorder duplicate registration info<commit_after>/*
Copyright 2016 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 mux

import (
	"fmt"
	"net/http"
	"runtime/debug"
	"sort"
	"strings"
	"sync"
	"sync/atomic"

	"k8s.io/klog/v2"

	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
)

// PathRecorderMux wraps a mux object and records the registered exposedPaths.
type PathRecorderMux struct {
	// name is used for logging so you can trace requests through
	name string

	lock            sync.Mutex
	notFoundHandler http.Handler
	pathToHandler   map[string]http.Handler
	prefixToHandler map[string]http.Handler

	// mux stores a pathHandler and is used to handle the actual serving.
	// Turns out, we want to accept trailing slashes, BUT we don't care about handling
	// everything under them.  This does exactly matches only unless its explicitly requested to
	// do something different
	mux atomic.Value

	// exposedPaths is the list of paths that should be shown at /
	exposedPaths []string

	// pathStacks holds the stacks of all registered paths.  This allows us to show a more helpful message
	// before the "http: multiple registrations for %s" panic.
	pathStacks map[string]string
}

// pathHandler is an http.Handler that will satisfy requests first by exact match, then by prefix,
// then by notFoundHandler
type pathHandler struct {
	// muxName is used for logging so you can trace requests through
	muxName string

	// pathToHandler is a map of exactly matching request to its handler
	pathToHandler map[string]http.Handler

	// this has to be sorted by most slashes then by length
	prefixHandlers []prefixHandler

	// notFoundHandler is the handler to use for satisfying requests with no other match
	notFoundHandler http.Handler
}

// prefixHandler holds the prefix it should match and the handler to use
type prefixHandler struct {
	// prefix is the prefix to test for a request match
	prefix string
	// handler is used to satisfy matching requests
	handler http.Handler
}

// NewPathRecorderMux creates a new PathRecorderMux
func NewPathRecorderMux(name string) *PathRecorderMux {
	ret := &PathRecorderMux{
		name:            name,
		pathToHandler:   map[string]http.Handler{},
		prefixToHandler: map[string]http.Handler{},
		mux:             atomic.Value{},
		exposedPaths:    []string{},
		pathStacks:      map[string]string{},
	}

	ret.mux.Store(&pathHandler{notFoundHandler: http.NotFoundHandler()})
	return ret
}

// ListedPaths returns the registered handler exposedPaths.
func (m *PathRecorderMux) ListedPaths() []string {
	handledPaths := append([]string{}, m.exposedPaths...)
	sort.Strings(handledPaths)

	return handledPaths
}

func (m *PathRecorderMux) trackCallers(path string) {
	stack := string(debug.Stack())
	if existingStack, ok := m.pathStacks[path]; ok {
		utilruntime.HandleError(fmt.Errorf("duplicate path registration of %q: original registration from %v\n\nnew registration from %v", path, existingStack, stack))
	}
	m.pathStacks[path] = stack
}

// refreshMuxLocked creates a new mux and must be called while locked.  Otherwise the view of handlers may
// not be consistent
func (m *PathRecorderMux) refreshMuxLocked() {
	newMux := &pathHandler{
		muxName:         m.name,
		pathToHandler:   map[string]http.Handler{},
		prefixHandlers:  []prefixHandler{},
		notFoundHandler: http.NotFoundHandler(),
	}
	if m.notFoundHandler != nil {
		newMux.notFoundHandler = m.notFoundHandler
	}
	for path, handler := range m.pathToHandler {
		newMux.pathToHandler[path] = handler
	}

	keys := sets.StringKeySet(m.prefixToHandler).List()
	sort.Sort(sort.Reverse(byPrefixPriority(keys)))
	for _, prefix := range keys {
		newMux.prefixHandlers = append(newMux.prefixHandlers, prefixHandler{
			prefix:  prefix,
			handler: m.prefixToHandler[prefix],
		})
	}

	m.mux.Store(newMux)
}

// NotFoundHandler sets the handler to use if there's no match for a give path
func (m *PathRecorderMux) NotFoundHandler(notFoundHandler http.Handler) {
	m.lock.Lock()
	defer m.lock.Unlock()

	m.notFoundHandler = notFoundHandler

	m.refreshMuxLocked()
}

// Unregister removes a path from the mux.
func (m *PathRecorderMux) Unregister(path string) {
	m.lock.Lock()
	defer m.lock.Unlock()

	delete(m.pathToHandler, path)
	delete(m.prefixToHandler, path)
	delete(m.pathStacks, path)
	for i := range m.exposedPaths {
		if m.exposedPaths[i] == path {
			m.exposedPaths = append(m.exposedPaths[:i], m.exposedPaths[i+1:]...)
			break
		}
	}

	m.refreshMuxLocked()
}

// Handle registers the handler for the given pattern.
// If a handler already exists for pattern, Handle panics.
func (m *PathRecorderMux) Handle(path string, handler http.Handler) {
	m.lock.Lock()
	defer m.lock.Unlock()
	m.trackCallers(path)

	m.exposedPaths = append(m.exposedPaths, path)
	m.pathToHandler[path] = handler
	m.refreshMuxLocked()
}

// HandleFunc registers the handler function for the given pattern.
// If a handler already exists for pattern, Handle panics.
func (m *PathRecorderMux) HandleFunc(path string, handler func(http.ResponseWriter, *http.Request)) {
	m.Handle(path, http.HandlerFunc(handler))
}

// UnlistedHandle registers the handler for the given pattern, but doesn't list it.
// If a handler already exists for pattern, Handle panics.
func (m *PathRecorderMux) UnlistedHandle(path string, handler http.Handler) {
	m.lock.Lock()
	defer m.lock.Unlock()
	m.trackCallers(path)

	m.pathToHandler[path] = handler
	m.refreshMuxLocked()
}

// UnlistedHandleFunc registers the handler function for the given pattern, but doesn't list it.
// If a handler already exists for pattern, Handle panics.
func (m *PathRecorderMux) UnlistedHandleFunc(path string, handler func(http.ResponseWriter, *http.Request)) {
	m.UnlistedHandle(path, http.HandlerFunc(handler))
}

// HandlePrefix is like Handle, but matches for anything under the path.  Like a standard golang trailing slash.
func (m *PathRecorderMux) HandlePrefix(path string, handler http.Handler) {
	if !strings.HasSuffix(path, "/") {
		panic(fmt.Sprintf("%q must end in a trailing slash", path))
	}

	m.lock.Lock()
	defer m.lock.Unlock()
	m.trackCallers(path)

	m.exposedPaths = append(m.exposedPaths, path)
	m.prefixToHandler[path] = handler
	m.refreshMuxLocked()
}

// UnlistedHandlePrefix is like UnlistedHandle, but matches for anything under the path.  Like a standard golang trailing slash.
func (m *PathRecorderMux) UnlistedHandlePrefix(path string, handler http.Handler) {
	if !strings.HasSuffix(path, "/") {
		panic(fmt.Sprintf("%q must end in a trailing slash", path))
	}

	m.lock.Lock()
	defer m.lock.Unlock()
	m.trackCallers(path)

	m.prefixToHandler[path] = handler
	m.refreshMuxLocked()
}

// ServeHTTP makes it an http.Handler
func (m *PathRecorderMux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	m.mux.Load().(*pathHandler).ServeHTTP(w, r)
}

// ServeHTTP makes it an http.Handler
func (h *pathHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if exactHandler, ok := h.pathToHandler[r.URL.Path]; ok {
		klog.V(5).Infof("%v: %q satisfied by exact match", h.muxName, r.URL.Path)
		exactHandler.ServeHTTP(w, r)
		return
	}

	for _, prefixHandler := range h.prefixHandlers {
		if strings.HasPrefix(r.URL.Path, prefixHandler.prefix) {
			klog.V(5).Infof("%v: %q satisfied by prefix %v", h.muxName, r.URL.Path, prefixHandler.prefix)
			prefixHandler.handler.ServeHTTP(w, r)
			return
		}
	}

	klog.V(5).Infof("%v: %q satisfied by NotFoundHandler", h.muxName, r.URL.Path)
	h.notFoundHandler.ServeHTTP(w, r)
}

// byPrefixPriority sorts url prefixes by the order in which they should be tested by the mux
// this has to be sorted by most slashes then by length so that we can iterate straight
// through to match the "best" one first.
type byPrefixPriority []string

func (s byPrefixPriority) Len() int      { return len(s) }
func (s byPrefixPriority) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s byPrefixPriority) Less(i, j int) bool {
	lhsNumParts := strings.Count(s[i], "/")
	rhsNumParts := strings.Count(s[j], "/")
	if lhsNumParts != rhsNumParts {
		return lhsNumParts < rhsNumParts
	}

	lhsLen := len(s[i])
	rhsLen := len(s[j])
	if lhsLen != rhsLen {
		return lhsLen < rhsLen
	}

	return strings.Compare(s[i], s[j]) < 0
}
<|endoftext|>


<commit_before>package sqlstore

import (
	"bytes"
	"fmt"
	"time"

	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/bus"
	m "github.com/grafana/grafana/pkg/models"
)

func init() {
	bus.AddHandler("sql", SaveAlerts)
	bus.AddHandler("sql", HandleAlertsQuery)
	bus.AddHandler("sql", GetAlertById)
	bus.AddHandler("sql", DeleteAlertById)
	bus.AddHandler("sql", GetAllAlertQueryHandler)
	bus.AddHandler("sql", SetAlertState)
	bus.AddHandler("sql", GetAlertStatesForDashboard)
	bus.AddHandler("sql", PauseAlertRule)
}

func GetAlertById(query *m.GetAlertByIdQuery) error {
	alert := m.Alert{}
	has, err := x.Id(query.Id).Get(&alert)
	if !has {
		return fmt.Errorf("could not find alert")
	}
	if err != nil {
		return err
	}

	query.Result = &alert
	return nil
}

func GetAllAlertQueryHandler(query *m.GetAllAlertsQuery) error {
	var alerts []*m.Alert
	err := x.Sql("select * from alert").Find(&alerts)
	if err != nil {
		return err
	}

	query.Result = alerts
	return nil
}

func DeleteAlertById(cmd *m.DeleteAlertCommand) error {
	return inTransaction(func(sess *xorm.Session) error {
		if _, err := sess.Exec("DELETE FROM alert WHERE id = ?", cmd.AlertId); err != nil {
			return err
		}

		return nil
	})
}

func HandleAlertsQuery(query *m.GetAlertsQuery) error {
	var sql bytes.Buffer
	params := make([]interface{}, 0)

	sql.WriteString(`SELECT *
						from alert
						`)

	sql.WriteString(`WHERE org_id = ?`)
	params = append(params, query.OrgId)

	if query.DashboardId != 0 {
		sql.WriteString(` AND dashboard_id = ?`)
		params = append(params, query.DashboardId)
	}

	if query.PanelId != 0 {
		sql.WriteString(` AND panel_id = ?`)
		params = append(params, query.PanelId)
	}

	if len(query.State) > 0 && query.State[0] != "ALL" {
		sql.WriteString(` AND (`)
		for i, v := range query.State {
			if i > 0 {
				sql.WriteString(" OR ")
			}
			sql.WriteString("state = ? ")
			params = append(params, v)
		}
		sql.WriteString(")")
	}

	if query.Limit != 0 {
		sql.WriteString(" LIMIT ?")
		params = append(params, query.Limit)
	}

	sql.WriteString(" ORDER BY name ASC")

	alerts := make([]*m.Alert, 0)
	if err := x.Sql(sql.String(), params...).Find(&alerts); err != nil {
		return err
	}

	query.Result = alerts
	return nil
}

func DeleteAlertDefinition(dashboardId int64, sess *xorm.Session) error {
	alerts := make([]*m.Alert, 0)
	sess.Where("dashboard_id = ?", dashboardId).Find(&alerts)

	for _, alert := range alerts {
		_, err := sess.Exec("DELETE FROM alert WHERE id = ? ", alert.Id)
		if err != nil {
			return err
		}

		sqlog.Debug("Alert deleted (due to dashboard deletion)", "name", alert.Name, "id", alert.Id)
	}

	return nil
}

func SaveAlerts(cmd *m.SaveAlertsCommand) error {
	return inTransaction(func(sess *xorm.Session) error {
		existingAlerts, err := GetAlertsByDashboardId2(cmd.DashboardId, sess)
		if err != nil {
			return err
		}

		if err := upsertAlerts(existingAlerts, cmd, sess); err != nil {
			return err
		}

		if err := deleteMissingAlerts(existingAlerts, cmd, sess); err != nil {
			return err
		}

		return nil
	})
}

func upsertAlerts(existingAlerts []*m.Alert, cmd *m.SaveAlertsCommand, sess *xorm.Session) error {
	for _, alert := range cmd.Alerts {
		update := false
		var alertToUpdate *m.Alert

		for _, k := range existingAlerts {
			if alert.PanelId == k.PanelId {
				update = true
				alert.Id = k.Id
				alertToUpdate = k
				break
			}
		}

		if update {
			if alertToUpdate.ContainsUpdates(alert) {
				alert.Updated = time.Now()
				alert.State = alertToUpdate.State
				_, err := sess.Id(alert.Id).Update(alert)
				if err != nil {
					return err
				}

				sqlog.Debug("Alert updated", "name", alert.Name, "id", alert.Id)
			}
		} else {
			alert.Updated = time.Now()
			alert.Created = time.Now()
			alert.State = m.AlertStateNoData
			alert.NewStateDate = time.Now()

			_, err := sess.Insert(alert)
			if err != nil {
				return err
			}

			sqlog.Debug("Alert inserted", "name", alert.Name, "id", alert.Id)
		}
	}

	return nil
}

func deleteMissingAlerts(alerts []*m.Alert, cmd *m.SaveAlertsCommand, sess *xorm.Session) error {
	for _, missingAlert := range alerts {
		missing := true

		for _, k := range cmd.Alerts {
			if missingAlert.PanelId == k.PanelId {
				missing = false
				break
			}
		}

		if missing {
			_, err := sess.Exec("DELETE FROM alert WHERE id = ?", missingAlert.Id)
			if err != nil {
				return err
			}

			sqlog.Debug("Alert deleted", "name", missingAlert.Name, "id", missingAlert.Id)
		}
	}

	return nil
}

func GetAlertsByDashboardId2(dashboardId int64, sess *xorm.Session) ([]*m.Alert, error) {
	alerts := make([]*m.Alert, 0)
	err := sess.Where("dashboard_id = ?", dashboardId).Find(&alerts)

	if err != nil {
		return []*m.Alert{}, err
	}

	return alerts, nil
}

func SetAlertState(cmd *m.SetAlertStateCommand) error {
	return inTransaction(func(sess *xorm.Session) error {
		alert := m.Alert{}

		if has, err := sess.Id(cmd.AlertId).Get(&alert); err != nil {
			return err
		} else if !has {
			return fmt.Errorf("Could not find alert")
		}

		alert.State = cmd.State
		alert.StateChanges += 1
		alert.NewStateDate = time.Now()
		alert.EvalData = cmd.EvalData

		if cmd.Error == "" {
			alert.ExecutionError = " " //without this space, xorm skips updating this field
		} else {
			alert.ExecutionError = cmd.Error
		}

		sess.Id(alert.Id).Update(&alert)
		return nil
	})
}

func PauseAlertRule(cmd *m.PauseAlertCommand) error {
	return inTransaction(func(sess *xorm.Session) error {
		alert := m.Alert{}

		if has, err := sess.Id(cmd.AlertId).Get(&alert); err != nil {
			return err
		} else if !has {
			return fmt.Errorf("Could not find alert")
		}

		var newState m.AlertStateType
		if cmd.Paused {
			newState = m.AlertStatePaused
		} else {
			newState = m.AlertStateNoData
		}
		alert.State = newState

		sess.Id(alert.Id).Update(&alert)
		return nil
	})
}

func GetAlertStatesForDashboard(query *m.GetAlertStatesForDashboardQuery) error {
	var rawSql = `SELECT
	                id,
	                dashboard_id,
	                panel_id,
	                state,
	                new_state_date
	                FROM alert
	                WHERE org_id = ? AND dashboard_id = ?`

	query.Result = make([]*m.AlertStateInfoDTO, 0)
	err := x.Sql(rawSql, query.OrgId, query.DashboardId).Find(&query.Result)

	return err
}
<commit_msg>feat(alerting): remove annotations when deleting alert, closes #6244<commit_after>package sqlstore

import (
	"bytes"
	"fmt"
	"time"

	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/bus"
	m "github.com/grafana/grafana/pkg/models"
)

func init() {
	bus.AddHandler("sql", SaveAlerts)
	bus.AddHandler("sql", HandleAlertsQuery)
	bus.AddHandler("sql", GetAlertById)
	bus.AddHandler("sql", DeleteAlertById)
	bus.AddHandler("sql", GetAllAlertQueryHandler)
	bus.AddHandler("sql", SetAlertState)
	bus.AddHandler("sql", GetAlertStatesForDashboard)
	bus.AddHandler("sql", PauseAlertRule)
}

func GetAlertById(query *m.GetAlertByIdQuery) error {
	alert := m.Alert{}
	has, err := x.Id(query.Id).Get(&alert)
	if !has {
		return fmt.Errorf("could not find alert")
	}
	if err != nil {
		return err
	}

	query.Result = &alert
	return nil
}

func GetAllAlertQueryHandler(query *m.GetAllAlertsQuery) error {
	var alerts []*m.Alert
	err := x.Sql("select * from alert").Find(&alerts)
	if err != nil {
		return err
	}

	query.Result = alerts
	return nil
}

func deleteAlertByIdInternal(alertId int64, reason string, sess *xorm.Session) error {
	sqlog.Debug("Deleting alert", "id", alertId, "reason", reason)

	if _, err := sess.Exec("DELETE FROM alert WHERE id = ?", alertId); err != nil {
		return err
	}

	if _, err := sess.Exec("DELETE FROM annotation WHERE alert_id = ?", alertId); err != nil {
		return err
	}

	return nil
}

func DeleteAlertById(cmd *m.DeleteAlertCommand) error {
	return inTransaction(func(sess *xorm.Session) error {
		return deleteAlertByIdInternal(cmd.AlertId, "DeleteAlertCommand", sess)
	})
}

func HandleAlertsQuery(query *m.GetAlertsQuery) error {
	var sql bytes.Buffer
	params := make([]interface{}, 0)

	sql.WriteString(`SELECT *
						from alert
						`)

	sql.WriteString(`WHERE org_id = ?`)
	params = append(params, query.OrgId)

	if query.DashboardId != 0 {
		sql.WriteString(` AND dashboard_id = ?`)
		params = append(params, query.DashboardId)
	}

	if query.PanelId != 0 {
		sql.WriteString(` AND panel_id = ?`)
		params = append(params, query.PanelId)
	}

	if len(query.State) > 0 && query.State[0] != "ALL" {
		sql.WriteString(` AND (`)
		for i, v := range query.State {
			if i > 0 {
				sql.WriteString(" OR ")
			}
			sql.WriteString("state = ? ")
			params = append(params, v)
		}
		sql.WriteString(")")
	}

	if query.Limit != 0 {
		sql.WriteString(" LIMIT ?")
		params = append(params, query.Limit)
	}

	sql.WriteString(" ORDER BY name ASC")

	alerts := make([]*m.Alert, 0)
	if err := x.Sql(sql.String(), params...).Find(&alerts); err != nil {
		return err
	}

	query.Result = alerts
	return nil
}

func DeleteAlertDefinition(dashboardId int64, sess *xorm.Session) error {
	alerts := make([]*m.Alert, 0)
	sess.Where("dashboard_id = ?", dashboardId).Find(&alerts)

	for _, alert := range alerts {
		deleteAlertByIdInternal(alert.Id, "Dashboard deleted", sess)
	}

	return nil
}

func SaveAlerts(cmd *m.SaveAlertsCommand) error {
	return inTransaction(func(sess *xorm.Session) error {
		existingAlerts, err := GetAlertsByDashboardId2(cmd.DashboardId, sess)
		if err != nil {
			return err
		}

		if err := upsertAlerts(existingAlerts, cmd, sess); err != nil {
			return err
		}

		if err := deleteMissingAlerts(existingAlerts, cmd, sess); err != nil {
			return err
		}

		return nil
	})
}

func upsertAlerts(existingAlerts []*m.Alert, cmd *m.SaveAlertsCommand, sess *xorm.Session) error {
	for _, alert := range cmd.Alerts {
		update := false
		var alertToUpdate *m.Alert

		for _, k := range existingAlerts {
			if alert.PanelId == k.PanelId {
				update = true
				alert.Id = k.Id
				alertToUpdate = k
				break
			}
		}

		if update {
			if alertToUpdate.ContainsUpdates(alert) {
				alert.Updated = time.Now()
				alert.State = alertToUpdate.State
				_, err := sess.Id(alert.Id).Update(alert)
				if err != nil {
					return err
				}

				sqlog.Debug("Alert updated", "name", alert.Name, "id", alert.Id)
			}
		} else {
			alert.Updated = time.Now()
			alert.Created = time.Now()
			alert.State = m.AlertStateNoData
			alert.NewStateDate = time.Now()

			_, err := sess.Insert(alert)
			if err != nil {
				return err
			}

			sqlog.Debug("Alert inserted", "name", alert.Name, "id", alert.Id)
		}
	}

	return nil
}

func deleteMissingAlerts(alerts []*m.Alert, cmd *m.SaveAlertsCommand, sess *xorm.Session) error {
	for _, missingAlert := range alerts {
		missing := true

		for _, k := range cmd.Alerts {
			if missingAlert.PanelId == k.PanelId {
				missing = false
				break
			}
		}

		if missing {
			deleteAlertByIdInternal(missingAlert.Id, "Removed from dashboard", sess)
		}
	}

	return nil
}

func GetAlertsByDashboardId2(dashboardId int64, sess *xorm.Session) ([]*m.Alert, error) {
	alerts := make([]*m.Alert, 0)
	err := sess.Where("dashboard_id = ?", dashboardId).Find(&alerts)

	if err != nil {
		return []*m.Alert{}, err
	}

	return alerts, nil
}

func SetAlertState(cmd *m.SetAlertStateCommand) error {
	return inTransaction(func(sess *xorm.Session) error {
		alert := m.Alert{}

		if has, err := sess.Id(cmd.AlertId).Get(&alert); err != nil {
			return err
		} else if !has {
			return fmt.Errorf("Could not find alert")
		}

		alert.State = cmd.State
		alert.StateChanges += 1
		alert.NewStateDate = time.Now()
		alert.EvalData = cmd.EvalData

		if cmd.Error == "" {
			alert.ExecutionError = " " //without this space, xorm skips updating this field
		} else {
			alert.ExecutionError = cmd.Error
		}

		sess.Id(alert.Id).Update(&alert)
		return nil
	})
}

func PauseAlertRule(cmd *m.PauseAlertCommand) error {
	return inTransaction(func(sess *xorm.Session) error {
		alert := m.Alert{}

		if has, err := sess.Id(cmd.AlertId).Get(&alert); err != nil {
			return err
		} else if !has {
			return fmt.Errorf("Could not find alert")
		}

		var newState m.AlertStateType
		if cmd.Paused {
			newState = m.AlertStatePaused
		} else {
			newState = m.AlertStateNoData
		}
		alert.State = newState

		sess.Id(alert.Id).Update(&alert)
		return nil
	})
}

func GetAlertStatesForDashboard(query *m.GetAlertStatesForDashboardQuery) error {
	var rawSql = `SELECT
	                id,
	                dashboard_id,
	                panel_id,
	                state,
	                new_state_date
	                FROM alert
	                WHERE org_id = ? AND dashboard_id = ?`

	query.Result = make([]*m.AlertStateInfoDTO, 0)
	err := x.Sql(rawSql, query.OrgId, query.DashboardId).Find(&query.Result)

	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Skaffold 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 build

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"sync"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build/tag"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/color"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/pkg/errors"
)

const bufferedLinesPerArtifact = 10000

type artifactBuilder func(ctx context.Context, out io.Writer, artifact *latest.Artifact, tag string) (string, error)

// For testing
var (
	buffSize      = bufferedLinesPerArtifact
	runInSequence = InSequence
)

// InParallel builds a list of artifacts in parallel but prints the logs in sequential order.
func InParallel(ctx context.Context, out io.Writer, tags tag.ImageTags, artifacts []*latest.Artifact, buildArtifact artifactBuilder) ([]Artifact, error) {
	if len(artifacts) == 0 {
		return nil, nil
	}

	if len(artifacts) == 1 {
		return runInSequence(ctx, out, tags, artifacts, buildArtifact)
	}

	var wg sync.WaitGroup
	defer wg.Wait()

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

	results := new(sync.Map)
	outputs := make([]chan []byte, len(artifacts))

	// Run builds in //
	wg.Add(len(artifacts))
	for i := range artifacts {
		outputs[i] = make(chan []byte, buffSize)
		r, w := io.Pipe()
		cw := setUpColorWriter(w, out)

		// Run build and write output/logs to piped writer and store build result in
		// sync.Map
		go runBuild(ctx, &wg, cw, tags, artifacts[i], results, buildArtifact)
		// Read build output/logs and write to buffered channel
		go readOutputAndWriteToChannel(r, outputs[i])
	}

	// Print logs and collect results in order.
	return collectResults(out, artifacts, results, outputs)
}

func runBuild(ctx context.Context, wg *sync.WaitGroup, cw io.WriteCloser, tags tag.ImageTags, artifact *latest.Artifact, results *sync.Map, build artifactBuilder) {
	defer wg.Done()
	event.BuildInProgress(artifact.ImageName)

	finalTag, err := getBuildResult(ctx, cw, tags, artifact, build)
	if err != nil {
		event.BuildFailed(artifact.ImageName, err)
		results.Store(artifact.ImageName, err)
	} else {
		event.BuildComplete(artifact.ImageName)
		artifact := Artifact{ImageName: artifact.ImageName, Tag: finalTag}
		results.Store(artifact.ImageName, artifact)
	}
	cw.Close()
}

func readOutputAndWriteToChannel(r io.Reader, lines chan []byte) {
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		lines <- scanner.Bytes()
	}
	close(lines)
}

func setUpColorWriter(w io.WriteCloser, out io.Writer) io.WriteCloser {
	if color.IsTerminal(out) {
		return color.ColoredWriteCloser{WriteCloser: w}
	}
	return w
}

func getBuildResult(ctx context.Context, cw io.Writer, tags tag.ImageTags, artifact *latest.Artifact, build artifactBuilder) (string, error) {
	color.Default.Fprintf(cw, "Building [%s]...\n", artifact.ImageName)
	tag, present := tags[artifact.ImageName]
	if !present {
		return "", fmt.Errorf("unable to find tag for image %s", artifact.ImageName)
	}
	return build(ctx, cw, artifact, tag)
}

func collectResults(out io.Writer, artifacts []*latest.Artifact, results *sync.Map, outputs []chan []byte) ([]Artifact, error) {
	var built []Artifact
	for i, artifact := range artifacts {
		// Wait for build to complete.
		printResult(out, outputs[i])
		v, ok := results.Load(artifact.ImageName)
		if !ok {
			return nil, fmt.Errorf("could not find build result for image %s", artifact.ImageName)
		}
		switch t := v.(type) {
		case error:
			return nil, errors.Wrapf(t, "building [%s]", artifact.ImageName)
		case Artifact:
			built = append(built, t)
		default:
			return nil, fmt.Errorf("unknown type %T for %s", t, artifact.ImageName)
		}
	}
	return built, nil
}

func printResult(out io.Writer, output chan []byte) {
	for line := range output {
		out.Write(line)
		fmt.Fprintln(out)
	}
}
<commit_msg>Keep parallel build wait group in one place<commit_after>/*
Copyright 2019 The Skaffold 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 build

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"sync"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build/tag"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/color"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/pkg/errors"
)

const bufferedLinesPerArtifact = 10000

type artifactBuilder func(ctx context.Context, out io.Writer, artifact *latest.Artifact, tag string) (string, error)

// For testing
var (
	buffSize      = bufferedLinesPerArtifact
	runInSequence = InSequence
)

// InParallel builds a list of artifacts in parallel but prints the logs in sequential order.
func InParallel(ctx context.Context, out io.Writer, tags tag.ImageTags, artifacts []*latest.Artifact, buildArtifact artifactBuilder) ([]Artifact, error) {
	if len(artifacts) == 0 {
		return nil, nil
	}

	if len(artifacts) == 1 {
		return runInSequence(ctx, out, tags, artifacts, buildArtifact)
	}

	var wg sync.WaitGroup
	defer wg.Wait()

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

	results := new(sync.Map)
	outputs := make([]chan []byte, len(artifacts))

	// Run builds in //
	wg.Add(len(artifacts))
	for i := range artifacts {
		outputs[i] = make(chan []byte, buffSize)
		r, w := io.Pipe()
		cw := setUpColorWriter(w, out)

		// Run build and write output/logs to piped writer and store build result in
		// sync.Map
		go func(i int) {
			runBuild(ctx, cw, tags, artifacts[i], results, buildArtifact)
			wg.Done()
		}(i)
		// Read build output/logs and write to buffered channel
		go readOutputAndWriteToChannel(r, outputs[i])
	}

	// Print logs and collect results in order.
	return collectResults(out, artifacts, results, outputs)
}

func runBuild(ctx context.Context, cw io.WriteCloser, tags tag.ImageTags, artifact *latest.Artifact, results *sync.Map, build artifactBuilder) {
	event.BuildInProgress(artifact.ImageName)

	finalTag, err := getBuildResult(ctx, cw, tags, artifact, build)
	if err != nil {
		event.BuildFailed(artifact.ImageName, err)
		results.Store(artifact.ImageName, err)
	} else {
		event.BuildComplete(artifact.ImageName)
		artifact := Artifact{ImageName: artifact.ImageName, Tag: finalTag}
		results.Store(artifact.ImageName, artifact)
	}
	cw.Close()
}

func readOutputAndWriteToChannel(r io.Reader, lines chan []byte) {
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		lines <- scanner.Bytes()
	}
	close(lines)
}

func setUpColorWriter(w io.WriteCloser, out io.Writer) io.WriteCloser {
	if color.IsTerminal(out) {
		return color.ColoredWriteCloser{WriteCloser: w}
	}
	return w
}

func getBuildResult(ctx context.Context, cw io.Writer, tags tag.ImageTags, artifact *latest.Artifact, build artifactBuilder) (string, error) {
	color.Default.Fprintf(cw, "Building [%s]...\n", artifact.ImageName)
	tag, present := tags[artifact.ImageName]
	if !present {
		return "", fmt.Errorf("unable to find tag for image %s", artifact.ImageName)
	}
	return build(ctx, cw, artifact, tag)
}

func collectResults(out io.Writer, artifacts []*latest.Artifact, results *sync.Map, outputs []chan []byte) ([]Artifact, error) {
	var built []Artifact
	for i, artifact := range artifacts {
		// Wait for build to complete.
		printResult(out, outputs[i])
		v, ok := results.Load(artifact.ImageName)
		if !ok {
			return nil, fmt.Errorf("could not find build result for image %s", artifact.ImageName)
		}
		switch t := v.(type) {
		case error:
			return nil, errors.Wrapf(t, "building [%s]", artifact.ImageName)
		case Artifact:
			built = append(built, t)
		default:
			return nil, fmt.Errorf("unknown type %T for %s", t, artifact.ImageName)
		}
	}
	return built, nil
}

func printResult(out io.Writer, output chan []byte) {
	for line := range output {
		out.Write(line)
		fmt.Fprintln(out)
	}
}
<|endoftext|>