<commit_before>package mpb

import (
	"container/heap"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v4/cwriter"
)

const (
	// default RefreshRate
	prr = 120 * time.Millisecond
	// default width
	pwidth = 80
)

// Progress represents the container that renders Progress bars
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	cwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	done         chan struct{}
	forceRefresh chan time.Time
	dlogger      *log.Logger
}

type pState struct {
	bHeap            *priorityQueue
	heapUpdated      bool
	pMatrix          map[int][]chan int
	aMatrix          map[int][]chan int
	barShutdownQueue []chan struct{}

	// following are provided/overrided by user
	idCount          int
	width            int
	rr               time.Duration
	uwg              *sync.WaitGroup
	manualRefreshCh  <-chan time.Time
	shutdownNotifier chan struct{}
	parkedBars       map[*Bar]*Bar
	output           io.Writer
	debugOut         io.Writer
}

// New creates new Progress container instance. It's not possible to
// reuse instance after *Progress.Wait() method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after *Progress.Wait()
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	pq := make(priorityQueue, 0)
	heap.Init(&pq)

	s := &pState{
		bHeap:      &pq,
		width:      pwidth,
		rr:         prr,
		parkedBars: make(map[*Bar]*Bar),
		output:     os.Stdout,
		debugOut:   ioutil.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		cwg:          new(sync.WaitGroup),
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		forceRefresh: make(chan time.Time),
		done:         make(chan struct{}),
		dlogger:      log.New(s.debugOut, "[mpb] ", log.Lshortfile),
	}
	p.cwg.Add(1)
	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a new progress bar and adds to the container.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.Add(total, newDefaultBarFiller(), options...)
}

// AddSpinner creates a new spinner bar and adds to the container.
func (p *Progress) AddSpinner(total int64, alignment SpinnerAlignment, options ...BarOption) *Bar {
	filler := &spinnerFiller{
		frames:    defaultSpinnerStyle,
		alignment: alignment,
	}
	return p.Add(total, filler, options...)
}

// Add creates a bar which renders itself by provided filler.
// Set total to 0, if you plan to update it later.
func (p *Progress) Add(total int64, filler Filler, options ...BarOption) *Bar {
	if filler == nil {
		filler = newDefaultBarFiller()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := &bState{
			total:    total,
			filler:   filler,
			priority: ps.idCount,
			id:       ps.idCount,
			width:    ps.width,
		}
		for _, opt := range options {
			if opt != nil {
				opt(bs)
			}
		}
		bar := newBar(p.ctx, p.bwg, bs)
		bar.forceRefresh = p.forceRefresh
		prefix := fmt.Sprintf("%sbar#%02d ", p.dlogger.Prefix(), bs.id)
		bar.dlogger = log.New(ps.debugOut, prefix, log.Lshortfile)
		if bs.runningBar != nil {
			if bar.priority == ps.idCount {
				bar.priority = bs.runningBar.priority
			}
			ps.parkedBars[bs.runningBar] = bar
		} else {
			heap.Push(ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		return <-result
	case <-p.done:
		p.bwg.Done()
		return nil
	}
}

// Abort is only effective while bar progress is running, it means
// remove bar now without waiting for its completion. If bar is already
// completed, there is nothing to abort. If you need to remove bar
// after completion, use BarRemoveOnComplete BarOption.
func (p *Progress) Abort(b *Bar, remove bool) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		if remove {
			s.heapUpdated = heap.Remove(s.bHeap, b.index) != nil
		}
		s.barShutdownQueue = append(s.barShutdownQueue, b.shutdown)
	}:
	case <-p.done:
	}
}

// UpdateBarPriority provides a way to change bar's order position.
// Zero is highest priority, i.e. bar will be on top.
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) { s.bHeap.update(b, priority) }:
	case <-p.done:
	}
}

// BarCount returns bars count
func (p *Progress) BarCount() int {
	result := make(chan int, 1)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Wait waits far all bars to complete and finally shutdowns container.
// After this method has been called, there is no way to reuse *Progress
// instance.
func (p *Progress) Wait() {
	if p.uwg != nil {
		// wait for user wg
		p.uwg.Wait()
	}

	// wait for bars to quit, if any
	p.bwg.Wait()

	close(p.done)

	// wait for container to quit
	p.cwg.Wait()
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer p.cwg.Done()

	manualOrTickCh, cleanUp := s.manualOrTick()
	defer cleanUp()

	refreshCh := fanInRefreshSrc(p.done, p.forceRefresh, manualOrTickCh)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case _, ok := <-refreshCh:
			if !ok {
				if s.shutdownNotifier != nil {
					close(s.shutdownNotifier)
				}
				return
			}
			if err := s.render(cw); err != nil {
				p.dlogger.Println(err)
			}
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	tw, err := cw.GetWidth()
	if err != nil {
		tw = s.width
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := (*s.bHeap)[i]
		go bar.render(tw)
	}

	return s.flush(cw)
}

func (s *pState) flush(cw *cwriter.Writer) error {
	var lineCount int
	for s.bHeap.Len() > 0 {
		bar := heap.Pop(s.bHeap).(*Bar)
		defer func() {
			if bar.toShutdown {
				// shutdown at next flush, in other words decrement underlying WaitGroup
				// only after the bar with completed state has been flushed. this
				// ensures no bar ends up with less than 100% rendered.
				s.barShutdownQueue = append(s.barShutdownQueue, bar.shutdown)
				if parkedBar := s.parkedBars[bar]; parkedBar != nil {
					heap.Push(s.bHeap, parkedBar)
					s.heapUpdated = true
					delete(s.parkedBars, bar)
				}
				if bar.dropOnComplete {
					s.heapUpdated = true
					return
				}
			}
			heap.Push(s.bHeap, bar)
		}()
		cw.ReadFrom(<-bar.frameCh)
		lineCount += bar.extendedLines + 1
	}

	for i := len(s.barShutdownQueue) - 1; i >= 0; i-- {
		close(s.barShutdownQueue[i])
		s.barShutdownQueue = s.barShutdownQueue[:i]
	}

	return cw.Flush(lineCount)
}

func (s *pState) manualOrTick() (<-chan time.Time, func()) {
	if s.manualRefreshCh != nil {
		return s.manualRefreshCh, func() {}
	}
	ticker := time.NewTicker(s.rr)
	return ticker.C, ticker.Stop
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := (*s.bHeap)[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		column := column
		go func() {
			var maxWidth int
			for _, ch := range column {
				w := <-ch
				if w > maxWidth {
					maxWidth = w
				}
			}
			for _, ch := range column {
				ch <- maxWidth
			}
		}()
	}
}

func fanInRefreshSrc(done <-chan struct{}, channels ...<-chan time.Time) <-chan time.Time {
	var wg sync.WaitGroup
	multiplexedStream := make(chan time.Time)

	multiplex := func(c <-chan time.Time) {
		defer wg.Done()
		// source channels are never closed (time.Ticker never closes associated
		// channel), so we cannot simply range over a c, instead we use select
		// inside infinite loop
		for {
			select {
			case v := <-c:
				select {
				case multiplexedStream <- v:
				case <-done:
					return
				}
			case <-done:
				return
			}
		}
	}

	wg.Add(len(channels))
	for _, c := range channels {
		go multiplex(c)
	}

	go func() {
		wg.Wait()
		close(multiplexedStream)
	}()

	return multiplexedStream
}
<commit_msg>avoid panic, if Wait called more than once<commit_after>package mpb

import (
	"container/heap"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v4/cwriter"
)

const (
	// default RefreshRate
	prr = 120 * time.Millisecond
	// default width
	pwidth = 80
)

// Progress represents the container that renders Progress bars
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	cwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	done         chan struct{}
	forceRefresh chan time.Time
	once         sync.Once
	dlogger      *log.Logger
}

type pState struct {
	bHeap            *priorityQueue
	heapUpdated      bool
	pMatrix          map[int][]chan int
	aMatrix          map[int][]chan int
	barShutdownQueue []chan struct{}

	// following are provided/overrided by user
	idCount          int
	width            int
	rr               time.Duration
	uwg              *sync.WaitGroup
	manualRefreshCh  <-chan time.Time
	shutdownNotifier chan struct{}
	parkedBars       map[*Bar]*Bar
	output           io.Writer
	debugOut         io.Writer
}

// New creates new Progress container instance. It's not possible to
// reuse instance after *Progress.Wait() method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after *Progress.Wait()
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	pq := make(priorityQueue, 0)
	heap.Init(&pq)

	s := &pState{
		bHeap:      &pq,
		width:      pwidth,
		rr:         prr,
		parkedBars: make(map[*Bar]*Bar),
		output:     os.Stdout,
		debugOut:   ioutil.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		cwg:          new(sync.WaitGroup),
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		forceRefresh: make(chan time.Time),
		done:         make(chan struct{}),
		dlogger:      log.New(s.debugOut, "[mpb] ", log.Lshortfile),
	}
	p.cwg.Add(1)
	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a new progress bar and adds to the container.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.Add(total, newDefaultBarFiller(), options...)
}

// AddSpinner creates a new spinner bar and adds to the container.
func (p *Progress) AddSpinner(total int64, alignment SpinnerAlignment, options ...BarOption) *Bar {
	filler := &spinnerFiller{
		frames:    defaultSpinnerStyle,
		alignment: alignment,
	}
	return p.Add(total, filler, options...)
}

// Add creates a bar which renders itself by provided filler.
// Set total to 0, if you plan to update it later.
func (p *Progress) Add(total int64, filler Filler, options ...BarOption) *Bar {
	if filler == nil {
		filler = newDefaultBarFiller()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := &bState{
			total:    total,
			filler:   filler,
			priority: ps.idCount,
			id:       ps.idCount,
			width:    ps.width,
		}
		for _, opt := range options {
			if opt != nil {
				opt(bs)
			}
		}
		bar := newBar(p.ctx, p.bwg, bs)
		bar.forceRefresh = p.forceRefresh
		prefix := fmt.Sprintf("%sbar#%02d ", p.dlogger.Prefix(), bs.id)
		bar.dlogger = log.New(ps.debugOut, prefix, log.Lshortfile)
		if bs.runningBar != nil {
			if bar.priority == ps.idCount {
				bar.priority = bs.runningBar.priority
			}
			ps.parkedBars[bs.runningBar] = bar
		} else {
			heap.Push(ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		return <-result
	case <-p.done:
		p.bwg.Done()
		return nil
	}
}

// Abort is only effective while bar progress is running, it means
// remove bar now without waiting for its completion. If bar is already
// completed, there is nothing to abort. If you need to remove bar
// after completion, use BarRemoveOnComplete BarOption.
func (p *Progress) Abort(b *Bar, remove bool) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		if remove {
			s.heapUpdated = heap.Remove(s.bHeap, b.index) != nil
		}
		s.barShutdownQueue = append(s.barShutdownQueue, b.shutdown)
	}:
	case <-p.done:
	}
}

// UpdateBarPriority provides a way to change bar's order position.
// Zero is highest priority, i.e. bar will be on top.
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) { s.bHeap.update(b, priority) }:
	case <-p.done:
	}
}

// BarCount returns bars count
func (p *Progress) BarCount() int {
	result := make(chan int, 1)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Wait waits far all bars to complete and finally shutdowns container.
// After this method has been called, there is no way to reuse *Progress
// instance.
func (p *Progress) Wait() {
	if p.uwg != nil {
		// wait for user wg
		p.uwg.Wait()
	}

	// wait for bars to quit, if any
	p.bwg.Wait()

	p.once.Do(p.shutdown)

	// wait for container to quit
	p.cwg.Wait()
}

func (p *Progress) shutdown() {
	close(p.done)
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer p.cwg.Done()

	manualOrTickCh, cleanUp := s.manualOrTick()
	defer cleanUp()

	refreshCh := fanInRefreshSrc(p.done, p.forceRefresh, manualOrTickCh)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case _, ok := <-refreshCh:
			if !ok {
				if s.shutdownNotifier != nil {
					close(s.shutdownNotifier)
				}
				return
			}
			if err := s.render(cw); err != nil {
				p.dlogger.Println(err)
			}
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	tw, err := cw.GetWidth()
	if err != nil {
		tw = s.width
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := (*s.bHeap)[i]
		go bar.render(tw)
	}

	return s.flush(cw)
}

func (s *pState) flush(cw *cwriter.Writer) error {
	var lineCount int
	for s.bHeap.Len() > 0 {
		bar := heap.Pop(s.bHeap).(*Bar)
		defer func() {
			if bar.toShutdown {
				// shutdown at next flush, in other words decrement underlying WaitGroup
				// only after the bar with completed state has been flushed. this
				// ensures no bar ends up with less than 100% rendered.
				s.barShutdownQueue = append(s.barShutdownQueue, bar.shutdown)
				if parkedBar := s.parkedBars[bar]; parkedBar != nil {
					heap.Push(s.bHeap, parkedBar)
					s.heapUpdated = true
					delete(s.parkedBars, bar)
				}
				if bar.dropOnComplete {
					s.heapUpdated = true
					return
				}
			}
			heap.Push(s.bHeap, bar)
		}()
		cw.ReadFrom(<-bar.frameCh)
		lineCount += bar.extendedLines + 1
	}

	for i := len(s.barShutdownQueue) - 1; i >= 0; i-- {
		close(s.barShutdownQueue[i])
		s.barShutdownQueue = s.barShutdownQueue[:i]
	}

	return cw.Flush(lineCount)
}

func (s *pState) manualOrTick() (<-chan time.Time, func()) {
	if s.manualRefreshCh != nil {
		return s.manualRefreshCh, func() {}
	}
	ticker := time.NewTicker(s.rr)
	return ticker.C, ticker.Stop
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := (*s.bHeap)[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		column := column
		go func() {
			var maxWidth int
			for _, ch := range column {
				w := <-ch
				if w > maxWidth {
					maxWidth = w
				}
			}
			for _, ch := range column {
				ch <- maxWidth
			}
		}()
	}
}

func fanInRefreshSrc(done <-chan struct{}, channels ...<-chan time.Time) <-chan time.Time {
	var wg sync.WaitGroup
	multiplexedStream := make(chan time.Time)

	multiplex := func(c <-chan time.Time) {
		defer wg.Done()
		// source channels are never closed (time.Ticker never closes associated
		// channel), so we cannot simply range over a c, instead we use select
		// inside infinite loop
		for {
			select {
			case v := <-c:
				select {
				case multiplexedStream <- v:
				case <-done:
					return
				}
			case <-done:
				return
			}
		}
	}

	wg.Add(len(channels))
	for _, c := range channels {
		go multiplex(c)
	}

	go func() {
		wg.Wait()
		close(multiplexedStream)
	}()

	return multiplexedStream
}
<|endoftext|>


<commit_before>// Copyright 2014 The Serviced 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 isvcs

import (
	"io/ioutil"
	"os"

	"github.com/Sirupsen/logrus"

	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"
)

const (
	ESRed ESHealth = iota
	ESYellow
	ESGreen
)

type ESHealth int

func GetHealth(health string) ESHealth {
	switch health {
	case "red":
		return ESRed
	case "yellow":
		return ESYellow
	case "green":
		return ESGreen
	}
	return ESHealth(-1)
}

func (health ESHealth) String() string {
	switch health {
	case ESRed:
		return "red"
	case ESYellow:
		return "yellow"
	case ESGreen:
		return "green"
	}
	return "unknown"
}

const DEFAULT_ES_STARTUP_TIMEOUT_SECONDS = 240 //default startup timeout in seconds (4 minutes)
const MIN_ES_STARTUP_TIMEOUT_SECONDS = 30      //minimum startup timeout in seconds

var elasticsearch_logstash *IService
var elasticsearch_serviced *IService

func initElasticSearch() {
	var serviceName string
	var err error

	serviceName = "elasticsearch-serviced"

	elasticsearch_servicedPortBinding := portBinding{
		HostIp:         "127.0.0.1",
		HostIpOverride: "SERVICED_ISVC_ELASTICSEARCH_SERVICED_PORT_9200_HOSTIP",
		HostPort:       9200,
	}

	defaultHealthCheck := healthCheckDefinition{
		healthCheck: esHealthCheck(getHostIp(elasticsearch_servicedPortBinding), 9200, ESYellow),
		Interval:    DEFAULT_HEALTHCHECK_INTERVAL,
		Timeout:     DEFAULT_HEALTHCHECK_TIMEOUT,
	}

	healthChecks := []map[string]healthCheckDefinition{
		{
			DEFAULT_HEALTHCHECK_NAME: defaultHealthCheck,
		},
	}

	elasticsearch_serviced, err = NewIService(
		IServiceDefinition{
			ID:             ElasticsearchServicedISVC.ID,
			Name:           serviceName,
			Repo:           IMAGE_REPO,
			Tag:            IMAGE_TAG,
			Command:        func() string { return "" },
			PortBindings:   []portBinding{elasticsearch_servicedPortBinding},
			Volumes:        map[string]string{"data": "/opt/elasticsearch-serviced/data"},
			Configuration:  make(map[string]interface{}),
			HealthChecks:   healthChecks,
			StartupTimeout: time.Duration(DEFAULT_ES_STARTUP_TIMEOUT_SECONDS) * time.Second,
			CustomStats:    GetElasticSearchCustomStats,
		},
	)
	if err != nil {
		log.WithFields(logrus.Fields{
			"isvc": elasticsearch_serviced.ID,
		}).WithError(err).Fatal("Unable to initialize internal service")
	}
	elasticsearch_serviced.Command = func() string {
		clusterArg := ""
		if clusterName, ok := elasticsearch_serviced.Configuration["cluster"]; ok {
			clusterArg = fmt.Sprintf(`-Ecluster.name="%s" `, clusterName)
		}
		cmd := fmt.Sprintf(`export JAVA_HOME=/usr/lib/jvm/jre-11; su elastic -c 'exec /opt/elasticsearch-serviced/bin/elasticsearch -Ecluster.initial_master_nodes="%s" -Enode.name="%s" %s'`,
			elasticsearch_serviced.Name, elasticsearch_serviced.Name, clusterArg)
		log.Infof("Build the command for running es-serviced: %s", cmd)
		return cmd
	}

	serviceName = "elasticsearch-logstash"

	elasticsearch_logstashPortBinding := portBinding{
		HostIp:         "127.0.0.1",
		HostIpOverride: "SERVICED_ISVC_ELASTICSEARCH_LOGSTASH_PORT_9100_HOSTIP",
		HostPort:       9100,
	}

	logStashHealthCheck := defaultHealthCheck
	logStashHealthCheck.healthCheck = esHealthCheck(getHostIp(elasticsearch_logstashPortBinding), 9100, ESYellow)

	healthChecks = []map[string]healthCheckDefinition{
		{
			DEFAULT_HEALTHCHECK_NAME: logStashHealthCheck,
		},
	}

	elasticsearch_logstash, err = NewIService(
		IServiceDefinition{
			ID:             ElasticsearchLogStashISVC.ID,
			Name:           serviceName,
			Repo:           IMAGE_REPO,
			Tag:            IMAGE_TAG,
			Command:        func() string { return "" },
			PortBindings:   []portBinding{elasticsearch_logstashPortBinding},
			Volumes:        map[string]string{"data": "/opt/elasticsearch-logstash/data"},
			Configuration:  make(map[string]interface{}),
			HealthChecks:   healthChecks,
			Recover:        recoverES,
			StartupTimeout: time.Duration(DEFAULT_ES_STARTUP_TIMEOUT_SECONDS) * time.Second,
			StartupFailed:  getESShardStatus,
		},
	)
	if err != nil {
		log.WithFields(logrus.Fields{
			"isvc": elasticsearch_logstash.ID,
		}).WithError(err).Fatal("Unable to initialize internal service")
	}

	elasticsearch_logstash.Command = func() string {
		nodeName := elasticsearch_logstash.Name
		clusterName := elasticsearch_logstash.Configuration["cluster"]
		return fmt.Sprintf("exec /opt/elasticsearch-logstash/bin/es-logstash-start.sh %s %s", nodeName, clusterName)
	}
}

func recoverES(path string) error {
	recoveryPath := path + "-backup"
	log := log.WithFields(logrus.Fields{
		"basepath":     path,
		"recoverypath": recoveryPath,
	})

	if _, err := os.Stat(recoveryPath); err == nil {
		log.Info("Overwriting existing recovery path")
		os.RemoveAll(recoveryPath)
	} else if !os.IsNotExist(err) {
		log.Debug("Could not stat recovery path")
		return err
	}

	if err := os.Rename(path, recoveryPath); err != nil {
		log.WithError(err).Debug("Could not recover elasticsearch")
		return err
	}
	log.Info("Moved and reset elasticsearch data")
	return nil
}

type esres struct {
	url      string
	response map[string]interface{}
	err      error
}

func getESShardStatus() {
	// try to get more information about how the shards are looking.
	// If some are 'UNASSIGNED', it may be possible to delete just those and restart
	host := elasticsearch_logstash.PortBindings[0].HostIp
	port := elasticsearch_logstash.PortBindings[0].HostPort
	url := fmt.Sprintf("http://%s:%d/_cat/shards", host, port)
	resp, err := http.Get(url)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		log.WithError(err).Error("Failed to get ES shard status.")
	}
	output, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.WithError(err).Error("Failed to get ES shard status.")
	} else {
		log.Warnf("Shard Status:\n%s", string(output))
	}
}

func getESHealth(url string) <-chan esres {
	esresC := make(chan esres, 1)
	go func() {
		resp, err := http.Get(url)
		if resp != nil {
			defer resp.Body.Close()
		}
		if err != nil {
			esresC <- esres{url, nil, err}
			return
		}
		if resp.StatusCode != 200 {
			esresC <- esres{url, nil, fmt.Errorf("received %d status code", resp.StatusCode)}
			return
		}

		var health map[string]interface{}
		if err := json.NewDecoder(resp.Body).Decode(&health); err != nil {
			esresC <- esres{url, nil, err}
			return
		}
		esresC <- esres{url, health, nil}

	}()
	return esresC
}

func esHealthCheck(host string, port int, minHealth ESHealth) HealthCheckFunction {
	return func(cancel <-chan struct{}) error {
		url := fmt.Sprintf("http://%s:%d/_cluster/health", host, port)
		log := log.WithFields(logrus.Fields{
			"url":       url,
			"minhealth": minHealth,
		})
		var r esres
		for {
			select {
			case r = <-getESHealth(url):
				if r.err != nil {
					log.WithError(r.err).Debugf("Unable to check Elastic health: %s", r.err)
					break
				}
				if status := GetHealth(r.response["status"].(string)); status < minHealth {
					log.WithFields(logrus.Fields{
						"reported":              r.response["status"],
						"cluster_name":          r.response["cluster_name"],
						"timed_out":             r.response["timed_out"],
						"number_of_nodes":       r.response["number_of_nodes"],
						"number_of_data_nodes":  r.response["number_of_data_nodes"],
						"active_primary_shards": r.response["active_primary_shards"],
						"active_shards":         r.response["active_shards"],
						"relocating_shards":     r.response["relocating_shards"],
						"initializing_shards":   r.response["initializing_shards"],
						"unassigned_shards":     r.response["unassigned_shards"],
					}).Warn("Elastic health reported below minimum")
					break
				}
				return nil
			case <-cancel:
				log.Debug("Canceled health check for Elastic")
				return nil
			}
			time.Sleep(time.Second)
		}
	}
}

func PurgeLogstashIndices(days int, gb int) {
	iservice := elasticsearch_logstash
	port := iservice.PortBindings[0].HostPort
	prefix := []string{"/usr/bin/curator", "--port", fmt.Sprintf("%d", port)}

	log := log.WithFields(logrus.Fields{
		"maxagedays": days,
		"maxsizegb":  gb,
	})

	log.Debug("Purging Logstash entries older than max age")
	indices := []string{"indices", "--older-than", fmt.Sprintf("%d", days), "--time-unit", "days", "--timestring", "%Y.%m.%d"}
	if output, err := iservice.Exec(append(append(prefix, "delete"), indices...)); err != nil {
		if !(strings.Contains(string(output), "No indices found in Elasticsearch") ||
			strings.Contains(string(output), "No indices matched provided args")) {
			log.WithError(err).Warn("Unable to purge logstash entries older than max age")
		}
	}
	log.Info("Purged Logstash entries older than max age")

	log.Debug("Purging Logstash entries to be below max size")
	indices = []string{"--disk-space", fmt.Sprintf("%d", gb), "indices", "--all-indices"}
	if output, err := iservice.Exec(append(append(prefix, "delete"), indices...)); err != nil {
		if !(strings.Contains(string(output), "No indices found in Elasticsearch") ||
			strings.Contains(string(output), "No indices matched provided args")) {
			log.WithError(err).Warn("Unable to purge logstash entries to be below max size")
		}
	}
	log.Info("Purged Logstash entries to be below max size")
}
<commit_msg>CC-4424 Update elasticsearch-curator to 5.8.3<commit_after>// Copyright 2014 The Serviced 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 isvcs

import (
	"io/ioutil"
	"os"

	"github.com/Sirupsen/logrus"

	"encoding/json"
	"fmt"
	"net/http"
	"time"
)

const (
	ESRed ESHealth = iota
	ESYellow
	ESGreen
)

type ESHealth int

func GetHealth(health string) ESHealth {
	switch health {
	case "red":
		return ESRed
	case "yellow":
		return ESYellow
	case "green":
		return ESGreen
	}
	return ESHealth(-1)
}

func (health ESHealth) String() string {
	switch health {
	case ESRed:
		return "red"
	case ESYellow:
		return "yellow"
	case ESGreen:
		return "green"
	}
	return "unknown"
}

const DEFAULT_ES_STARTUP_TIMEOUT_SECONDS = 240 //default startup timeout in seconds (4 minutes)
const MIN_ES_STARTUP_TIMEOUT_SECONDS = 30      //minimum startup timeout in seconds

var elasticsearch_logstash *IService
var elasticsearch_serviced *IService

func initElasticSearch() {
	var serviceName string
	var err error

	serviceName = "elasticsearch-serviced"

	elasticsearch_servicedPortBinding := portBinding{
		HostIp:         "127.0.0.1",
		HostIpOverride: "SERVICED_ISVC_ELASTICSEARCH_SERVICED_PORT_9200_HOSTIP",
		HostPort:       9200,
	}

	defaultHealthCheck := healthCheckDefinition{
		healthCheck: esHealthCheck(getHostIp(elasticsearch_servicedPortBinding), 9200, ESYellow),
		Interval:    DEFAULT_HEALTHCHECK_INTERVAL,
		Timeout:     DEFAULT_HEALTHCHECK_TIMEOUT,
	}

	healthChecks := []map[string]healthCheckDefinition{
		{
			DEFAULT_HEALTHCHECK_NAME: defaultHealthCheck,
		},
	}

	elasticsearch_serviced, err = NewIService(
		IServiceDefinition{
			ID:             ElasticsearchServicedISVC.ID,
			Name:           serviceName,
			Repo:           IMAGE_REPO,
			Tag:            IMAGE_TAG,
			Command:        func() string { return "" },
			PortBindings:   []portBinding{elasticsearch_servicedPortBinding},
			Volumes:        map[string]string{"data": "/opt/elasticsearch-serviced/data"},
			Configuration:  make(map[string]interface{}),
			HealthChecks:   healthChecks,
			StartupTimeout: time.Duration(DEFAULT_ES_STARTUP_TIMEOUT_SECONDS) * time.Second,
			CustomStats:    GetElasticSearchCustomStats,
		},
	)
	if err != nil {
		log.WithFields(logrus.Fields{
			"isvc": elasticsearch_serviced.ID,
		}).WithError(err).Fatal("Unable to initialize internal service")
	}
	elasticsearch_serviced.Command = func() string {
		clusterArg := ""
		if clusterName, ok := elasticsearch_serviced.Configuration["cluster"]; ok {
			clusterArg = fmt.Sprintf(`-Ecluster.name="%s" `, clusterName)
		}
		cmd := fmt.Sprintf(`export JAVA_HOME=/usr/lib/jvm/jre-11; su elastic -c 'exec /opt/elasticsearch-serviced/bin/elasticsearch -Ecluster.initial_master_nodes="%s" -Enode.name="%s" %s'`,
			elasticsearch_serviced.Name, elasticsearch_serviced.Name, clusterArg)
		log.Infof("Build the command for running es-serviced: %s", cmd)
		return cmd
	}

	serviceName = "elasticsearch-logstash"

	elasticsearch_logstashPortBinding := portBinding{
		HostIp:         "127.0.0.1",
		HostIpOverride: "SERVICED_ISVC_ELASTICSEARCH_LOGSTASH_PORT_9100_HOSTIP",
		HostPort:       9100,
	}

	logStashHealthCheck := defaultHealthCheck
	logStashHealthCheck.healthCheck = esHealthCheck(getHostIp(elasticsearch_logstashPortBinding), 9100, ESYellow)

	healthChecks = []map[string]healthCheckDefinition{
		{
			DEFAULT_HEALTHCHECK_NAME: logStashHealthCheck,
		},
	}

	elasticsearch_logstash, err = NewIService(
		IServiceDefinition{
			ID:             ElasticsearchLogStashISVC.ID,
			Name:           serviceName,
			Repo:           IMAGE_REPO,
			Tag:            IMAGE_TAG,
			Command:        func() string { return "" },
			PortBindings:   []portBinding{elasticsearch_logstashPortBinding},
			Volumes:        map[string]string{"data": "/opt/elasticsearch-logstash/data"},
			Configuration:  make(map[string]interface{}),
			HealthChecks:   healthChecks,
			Recover:        recoverES,
			StartupTimeout: time.Duration(DEFAULT_ES_STARTUP_TIMEOUT_SECONDS) * time.Second,
			StartupFailed:  getESShardStatus,
		},
	)
	if err != nil {
		log.WithFields(logrus.Fields{
			"isvc": elasticsearch_logstash.ID,
		}).WithError(err).Fatal("Unable to initialize internal service")
	}

	elasticsearch_logstash.Command = func() string {
		nodeName := elasticsearch_logstash.Name
		clusterName := elasticsearch_logstash.Configuration["cluster"]
		return fmt.Sprintf("exec /opt/elasticsearch-logstash/bin/es-logstash-start.sh %s %s", nodeName, clusterName)
	}
}

func recoverES(path string) error {
	recoveryPath := path + "-backup"
	log := log.WithFields(logrus.Fields{
		"basepath":     path,
		"recoverypath": recoveryPath,
	})

	if _, err := os.Stat(recoveryPath); err == nil {
		log.Info("Overwriting existing recovery path")
		os.RemoveAll(recoveryPath)
	} else if !os.IsNotExist(err) {
		log.Debug("Could not stat recovery path")
		return err
	}

	if err := os.Rename(path, recoveryPath); err != nil {
		log.WithError(err).Debug("Could not recover elasticsearch")
		return err
	}
	log.Info("Moved and reset elasticsearch data")
	return nil
}

type esres struct {
	url      string
	response map[string]interface{}
	err      error
}

func getESShardStatus() {
	// try to get more information about how the shards are looking.
	// If some are 'UNASSIGNED', it may be possible to delete just those and restart
	host := elasticsearch_logstash.PortBindings[0].HostIp
	port := elasticsearch_logstash.PortBindings[0].HostPort
	url := fmt.Sprintf("http://%s:%d/_cat/shards", host, port)
	resp, err := http.Get(url)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		log.WithError(err).Error("Failed to get ES shard status.")
	}
	output, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.WithError(err).Error("Failed to get ES shard status.")
	} else {
		log.Warnf("Shard Status:\n%s", string(output))
	}
}

func getESHealth(url string) <-chan esres {
	esresC := make(chan esres, 1)
	go func() {
		resp, err := http.Get(url)
		if resp != nil {
			defer resp.Body.Close()
		}
		if err != nil {
			esresC <- esres{url, nil, err}
			return
		}
		if resp.StatusCode != 200 {
			esresC <- esres{url, nil, fmt.Errorf("received %d status code", resp.StatusCode)}
			return
		}

		var health map[string]interface{}
		if err := json.NewDecoder(resp.Body).Decode(&health); err != nil {
			esresC <- esres{url, nil, err}
			return
		}
		esresC <- esres{url, health, nil}

	}()
	return esresC
}

func esHealthCheck(host string, port int, minHealth ESHealth) HealthCheckFunction {
	return func(cancel <-chan struct{}) error {
		url := fmt.Sprintf("http://%s:%d/_cluster/health", host, port)
		log := log.WithFields(logrus.Fields{
			"url":       url,
			"minhealth": minHealth,
		})
		var r esres
		for {
			select {
			case r = <-getESHealth(url):
				if r.err != nil {
					log.WithError(r.err).Debugf("Unable to check Elastic health: %s", r.err)
					break
				}
				if status := GetHealth(r.response["status"].(string)); status < minHealth {
					log.WithFields(logrus.Fields{
						"reported":              r.response["status"],
						"cluster_name":          r.response["cluster_name"],
						"timed_out":             r.response["timed_out"],
						"number_of_nodes":       r.response["number_of_nodes"],
						"number_of_data_nodes":  r.response["number_of_data_nodes"],
						"active_primary_shards": r.response["active_primary_shards"],
						"active_shards":         r.response["active_shards"],
						"relocating_shards":     r.response["relocating_shards"],
						"initializing_shards":   r.response["initializing_shards"],
						"unassigned_shards":     r.response["unassigned_shards"],
					}).Warn("Elastic health reported below minimum")
					break
				}
				return nil
			case <-cancel:
				log.Debug("Canceled health check for Elastic")
				return nil
			}
			time.Sleep(time.Second)
		}
	}
}

func PurgeLogstashIndices(days int, gb int) {
	iservice := elasticsearch_logstash
	port := iservice.PortBindings[0].HostPort

	curatorConfigPath := "/opt/curator/curator-config.yml"
	curatorActionsPath := "/opt/curator/curator-actions.yml"

	prefix := []string{
		fmt.Sprintf("ES_PORT=%d", port),
		fmt.Sprintf("MAX_AGE_DAYS=%d", days),
		fmt.Sprintf("MAX_SIZE_GB=%d", gb),
	}
	command := []string{"/usr/bin/curator", "--config", curatorConfigPath, curatorActionsPath}

	log := log.WithFields(logrus.Fields{
		"maxagedays": days,
		"maxsizegb":  gb,
	})

	log.Debug("Purging Logstash entries older than max age or bigger than max size")
	if _, err := iservice.Exec(append(prefix, command...)); err != nil {
		log.WithError(err).Warn("Unable to purge logstash entries")
	} else {
		log.Info("Purged Logstash entries older than max age or bigger than max size")
	}
}
<|endoftext|>


<commit_before>package agent

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"strings"

	"github.com/buildkite/agent/env"
	"github.com/buildkite/interpolate"

	// This is a fork of gopkg.in/yaml.v2 that fixes anchors with MapSlice
	yaml "github.com/buildkite/yaml"
)

type PipelineParser struct {
	Env      *env.Environment
	Filename string
	Pipeline []byte
}

func (p PipelineParser) Parse() (interface{}, error) {
	if p.Env == nil {
		p.Env = env.FromSlice(os.Environ())
	}

	var errPrefix string
	if p.Filename == "" {
		errPrefix = "Failed to parse pipeline"
	} else {
		errPrefix = fmt.Sprintf("Failed to parse %s", p.Filename)
	}

	var pipeline interface{}
	var pipelineAsSlice []interface{}

	// Historically we support uploading just steps, so we parse it as either a
	// slice, or if it's a map we need to do environment block processing
	if err := yaml.Unmarshal([]byte(p.Pipeline), &pipelineAsSlice); err == nil {
		pipeline = pipelineAsSlice
	} else {
		pipelineAsMap, err := p.parseWithEnv()
		if err != nil {
			return nil, fmt.Errorf("%s: %v", errPrefix, formatYAMLError(err))
		}
		pipeline = pipelineAsMap
	}

	// Recursively go through the entire pipeline and perform environment
	// variable interpolation on strings
	interpolated, err := p.interpolate(pipeline)
	if err != nil {
		return nil, err
	}

	// Now we roundtrip this back into YAML bytes and back into a generic interface{}
	// that works with all upstream code (which likes working with JSON). Specifically we
	// need to convert the map[interface{}]interface{}'s that YAML likes into JSON compatible
	// map[string]interface{}
	b, err := yaml.Marshal(interpolated)
	if err != nil {
		return nil, err
	}

	var result interface{}
	if err := unmarshalAsStringMap(b, &result); err != nil {
		return nil, err
	}

	return result, nil
}

func (p PipelineParser) parseWithEnv() (interface{}, error) {
	var pipeline yaml.MapSlice

	// Initially we unmarshal this into a yaml.MapSlice so that we preserve the order of maps
	if err := yaml.Unmarshal([]byte(p.Pipeline), &pipeline); err != nil {
		return nil, err
	}

	// Preprocess any env tat are defined in the top level block and place them into env for
	// later interpolation into env blocks
	if item, ok := mapSliceItem("env", pipeline); ok {
		if envMap, ok := item.Value.(yaml.MapSlice); ok {
			if err := p.interpolateEnvBlock(envMap); err != nil {
				return nil, err
			}
		} else {
			return nil, fmt.Errorf("Expected pipeline top-level env block to be a map, got %T", item)
		}
	}

	return pipeline, nil
}

func mapSliceItem(key string, s yaml.MapSlice) (yaml.MapItem, bool) {
	for _, item := range s {
		if k, ok := item.Key.(string); ok && k == key {
			return item, true
		}
	}
	return yaml.MapItem{}, false
}

func (p PipelineParser) interpolateEnvBlock(envMap yaml.MapSlice) error {
	for _, item := range envMap {
		k, ok := item.Key.(string)
		if !ok {
			return fmt.Errorf("Unexpected type of %T for env block key %v", item.Key, item.Key)
		}
		switch tv := item.Value.(type) {
		case string:
			interpolated, err := interpolate.Interpolate(p.Env, tv)
			if err != nil {
				return err
			}
			p.Env.Set(k, interpolated)
		}
	}
	return nil
}

func formatYAMLError(err error) error {
	return errors.New(strings.TrimPrefix(err.Error(), "yaml: "))
}

// interpolate function inspired from: https://gist.github.com/hvoecking/10772475

func (p PipelineParser) interpolate(obj interface{}) (interface{}, error) {
	// Make sure there's something actually to interpolate
	if obj == nil {
		return nil, nil
	}

	// Wrap the original in a reflect.Value
	original := reflect.ValueOf(obj)

	// Make a copy that we'll add the new values to
	copy := reflect.New(original.Type()).Elem()

	err := p.interpolateRecursive(copy, original)
	if err != nil {
		return nil, err
	}

	// Remove the reflection wrapper
	return copy.Interface(), nil
}

func (p PipelineParser) interpolateRecursive(copy, original reflect.Value) error {
	switch original.Kind() {
	// If it is a pointer we need to unwrap and call once again
	case reflect.Ptr:
		// To get the actual value of the original we have to call Elem()
		// At the same time this unwraps the pointer so we don't end up in
		// an infinite recursion
		originalValue := original.Elem()

		// Check if the pointer is nil
		if !originalValue.IsValid() {
			return nil
		}

		// Allocate a new object and set the pointer to it
		copy.Set(reflect.New(originalValue.Type()))

		// Unwrap the newly created pointer
		err := p.interpolateRecursive(copy.Elem(), originalValue)
		if err != nil {
			return err
		}

	// If it is an interface (which is very similar to a pointer), do basically the
	// same as for the pointer. Though a pointer is not the same as an interface so
	// note that we have to call Elem() after creating a new object because otherwise
	// we would end up with an actual pointer
	case reflect.Interface:
		// Get rid of the wrapping interface
		originalValue := original.Elem()

		// Check to make sure the interface isn't nil
		if !originalValue.IsValid() {
			return nil
		}

		// Create a new object. Now new gives us a pointer, but we want the value it
		// points to, so we have to call Elem() to unwrap it
		copyValue := reflect.New(originalValue.Type()).Elem()

		err := p.interpolateRecursive(copyValue, originalValue)
		if err != nil {
			return err
		}

		copy.Set(copyValue)

	// If it is a struct we interpolate each field
	case reflect.Struct:
		for i := 0; i < original.NumField(); i += 1 {
			err := p.interpolateRecursive(copy.Field(i), original.Field(i))
			if err != nil {
				return err
			}
		}

	// If it is a slice we create a new slice and interpolate each element
	case reflect.Slice:
		copy.Set(reflect.MakeSlice(original.Type(), original.Len(), original.Cap()))

		for i := 0; i < original.Len(); i += 1 {
			err := p.interpolateRecursive(copy.Index(i), original.Index(i))
			if err != nil {
				return err
			}
		}

	// If it is a map we create a new map and interpolate each value
	case reflect.Map:
		copy.Set(reflect.MakeMap(original.Type()))

		for _, key := range original.MapKeys() {
			originalValue := original.MapIndex(key)

			// New gives us a pointer, but again we want the value
			copyValue := reflect.New(originalValue.Type()).Elem()
			err := p.interpolateRecursive(copyValue, originalValue)
			if err != nil {
				return err
			}

			// Also interpolate the key if it's a string
			if key.Kind() == reflect.String {
				interpolatedKey, err := interpolate.Interpolate(p.Env, key.Interface().(string))
				if err != nil {
					return err
				}
				copy.SetMapIndex(reflect.ValueOf(interpolatedKey), copyValue)
			} else {
				copy.SetMapIndex(key, copyValue)
			}
		}

	// If it is a string interpolate it (yay finally we're doing what we came for)
	case reflect.String:
		interpolated, err := interpolate.Interpolate(p.Env, original.Interface().(string))
		if err != nil {
			return err
		}
		copy.SetString(interpolated)

	// And everything else will simply be taken from the original
	default:
		copy.Set(original)
	}

	return nil
}

// Unmarshal YAML to map[string]interface{} instead of map[interface{}]interface{}, such that
// we can Marshal cleanly into JSON
// Via https://github.com/go-yaml/yaml/issues/139#issuecomment-220072190
func unmarshalAsStringMap(in []byte, out interface{}) error {
	var res interface{}

	if err := yaml.Unmarshal(in, &res); err != nil {
		return err
	}
	*out.(*interface{}) = cleanupMapValue(res)

	return nil
}

func cleanupInterfaceArray(in []interface{}) []interface{} {
	res := make([]interface{}, len(in))
	for i, v := range in {
		res[i] = cleanupMapValue(v)
	}
	return res
}

func cleanupInterfaceMap(in map[interface{}]interface{}) map[string]interface{} {
	res := make(map[string]interface{})
	for k, v := range in {
		res[fmt.Sprintf("%v", k)] = cleanupMapValue(v)
	}
	return res
}

func cleanupMapValue(v interface{}) interface{} {
	switch v := v.(type) {
	case []interface{}:
		return cleanupInterfaceArray(v)
	case map[interface{}]interface{}:
		return cleanupInterfaceMap(v)
	case string:
		return v
	default:
		return fmt.Sprintf("%v", v)
	}
}
<commit_msg>Preserve types in YAML to JSON conversion<commit_after>package agent

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"strings"

	"github.com/buildkite/agent/env"
	"github.com/buildkite/interpolate"

	// This is a fork of gopkg.in/yaml.v2 that fixes anchors with MapSlice
	yaml "github.com/buildkite/yaml"
)

type PipelineParser struct {
	Env      *env.Environment
	Filename string
	Pipeline []byte
}

func (p PipelineParser) Parse() (interface{}, error) {
	if p.Env == nil {
		p.Env = env.FromSlice(os.Environ())
	}

	var errPrefix string
	if p.Filename == "" {
		errPrefix = "Failed to parse pipeline"
	} else {
		errPrefix = fmt.Sprintf("Failed to parse %s", p.Filename)
	}

	var pipeline interface{}
	var pipelineAsSlice []interface{}

	// Historically we support uploading just steps, so we parse it as either a
	// slice, or if it's a map we need to do environment block processing
	if err := yaml.Unmarshal([]byte(p.Pipeline), &pipelineAsSlice); err == nil {
		pipeline = pipelineAsSlice
	} else {
		pipelineAsMap, err := p.parseWithEnv()
		if err != nil {
			return nil, fmt.Errorf("%s: %v", errPrefix, formatYAMLError(err))
		}
		pipeline = pipelineAsMap
	}

	// Recursively go through the entire pipeline and perform environment
	// variable interpolation on strings
	interpolated, err := p.interpolate(pipeline)
	if err != nil {
		return nil, err
	}

	// Now we roundtrip this back into YAML bytes and back into a generic interface{}
	// that works with all upstream code (which likes working with JSON). Specifically we
	// need to convert the map[interface{}]interface{}'s that YAML likes into JSON compatible
	// map[string]interface{}
	b, err := yaml.Marshal(interpolated)
	if err != nil {
		return nil, err
	}

	var result interface{}
	if err := unmarshalAsStringMap(b, &result); err != nil {
		return nil, fmt.Errorf("%s: %v", errPrefix, formatYAMLError(err))
	}

	return result, nil
}

func (p PipelineParser) parseWithEnv() (interface{}, error) {
	var pipeline yaml.MapSlice

	// Initially we unmarshal this into a yaml.MapSlice so that we preserve the order of maps
	if err := yaml.Unmarshal([]byte(p.Pipeline), &pipeline); err != nil {
		return nil, err
	}

	// Preprocess any env tat are defined in the top level block and place them into env for
	// later interpolation into env blocks
	if item, ok := mapSliceItem("env", pipeline); ok {
		if envMap, ok := item.Value.(yaml.MapSlice); ok {
			if err := p.interpolateEnvBlock(envMap); err != nil {
				return nil, err
			}
		} else {
			return nil, fmt.Errorf("Expected pipeline top-level env block to be a map, got %T", item)
		}
	}

	return pipeline, nil
}

func mapSliceItem(key string, s yaml.MapSlice) (yaml.MapItem, bool) {
	for _, item := range s {
		if k, ok := item.Key.(string); ok && k == key {
			return item, true
		}
	}
	return yaml.MapItem{}, false
}

func (p PipelineParser) interpolateEnvBlock(envMap yaml.MapSlice) error {
	for _, item := range envMap {
		k, ok := item.Key.(string)
		if !ok {
			return fmt.Errorf("Unexpected type of %T for env block key %v", item.Key, item.Key)
		}
		switch tv := item.Value.(type) {
		case string:
			interpolated, err := interpolate.Interpolate(p.Env, tv)
			if err != nil {
				return err
			}
			p.Env.Set(k, interpolated)
		}
	}
	return nil
}

func formatYAMLError(err error) error {
	return errors.New(strings.TrimPrefix(err.Error(), "yaml: "))
}

// interpolate function inspired from: https://gist.github.com/hvoecking/10772475

func (p PipelineParser) interpolate(obj interface{}) (interface{}, error) {
	// Make sure there's something actually to interpolate
	if obj == nil {
		return nil, nil
	}

	// Wrap the original in a reflect.Value
	original := reflect.ValueOf(obj)

	// Make a copy that we'll add the new values to
	copy := reflect.New(original.Type()).Elem()

	err := p.interpolateRecursive(copy, original)
	if err != nil {
		return nil, err
	}

	// Remove the reflection wrapper
	return copy.Interface(), nil
}

func (p PipelineParser) interpolateRecursive(copy, original reflect.Value) error {
	switch original.Kind() {
	// If it is a pointer we need to unwrap and call once again
	case reflect.Ptr:
		// To get the actual value of the original we have to call Elem()
		// At the same time this unwraps the pointer so we don't end up in
		// an infinite recursion
		originalValue := original.Elem()

		// Check if the pointer is nil
		if !originalValue.IsValid() {
			return nil
		}

		// Allocate a new object and set the pointer to it
		copy.Set(reflect.New(originalValue.Type()))

		// Unwrap the newly created pointer
		err := p.interpolateRecursive(copy.Elem(), originalValue)
		if err != nil {
			return err
		}

	// If it is an interface (which is very similar to a pointer), do basically the
	// same as for the pointer. Though a pointer is not the same as an interface so
	// note that we have to call Elem() after creating a new object because otherwise
	// we would end up with an actual pointer
	case reflect.Interface:
		// Get rid of the wrapping interface
		originalValue := original.Elem()

		// Check to make sure the interface isn't nil
		if !originalValue.IsValid() {
			return nil
		}

		// Create a new object. Now new gives us a pointer, but we want the value it
		// points to, so we have to call Elem() to unwrap it
		copyValue := reflect.New(originalValue.Type()).Elem()

		err := p.interpolateRecursive(copyValue, originalValue)
		if err != nil {
			return err
		}

		copy.Set(copyValue)

	// If it is a struct we interpolate each field
	case reflect.Struct:
		for i := 0; i < original.NumField(); i += 1 {
			err := p.interpolateRecursive(copy.Field(i), original.Field(i))
			if err != nil {
				return err
			}
		}

	// If it is a slice we create a new slice and interpolate each element
	case reflect.Slice:
		copy.Set(reflect.MakeSlice(original.Type(), original.Len(), original.Cap()))

		for i := 0; i < original.Len(); i += 1 {
			err := p.interpolateRecursive(copy.Index(i), original.Index(i))
			if err != nil {
				return err
			}
		}

	// If it is a map we create a new map and interpolate each value
	case reflect.Map:
		copy.Set(reflect.MakeMap(original.Type()))

		for _, key := range original.MapKeys() {
			originalValue := original.MapIndex(key)

			// New gives us a pointer, but again we want the value
			copyValue := reflect.New(originalValue.Type()).Elem()
			err := p.interpolateRecursive(copyValue, originalValue)
			if err != nil {
				return err
			}

			// Also interpolate the key if it's a string
			if key.Kind() == reflect.String {
				interpolatedKey, err := interpolate.Interpolate(p.Env, key.Interface().(string))
				if err != nil {
					return err
				}
				copy.SetMapIndex(reflect.ValueOf(interpolatedKey), copyValue)
			} else {
				copy.SetMapIndex(key, copyValue)
			}
		}

	// If it is a string interpolate it (yay finally we're doing what we came for)
	case reflect.String:
		interpolated, err := interpolate.Interpolate(p.Env, original.Interface().(string))
		if err != nil {
			return err
		}
		copy.SetString(interpolated)

	// And everything else will simply be taken from the original
	default:
		copy.Set(original)
	}

	return nil
}

// Unmarshal YAML to map[string]interface{} instead of map[interface{}]interface{}, such that
// we can Marshal cleanly into JSON
// Via https://github.com/go-yaml/yaml/issues/139#issuecomment-220072190
func unmarshalAsStringMap(in []byte, out interface{}) error {
	var res interface{}

	if err := yaml.Unmarshal(in, &res); err != nil {
		return err
	}
	*out.(*interface{}) = cleanupMapValue(res)

	return nil
}

func cleanupInterfaceArray(in []interface{}) []interface{} {
	res := make([]interface{}, len(in))
	for i, v := range in {
		res[i] = cleanupMapValue(v)
	}
	return res
}

func cleanupInterfaceMap(in map[interface{}]interface{}) map[string]interface{} {
	res := make(map[string]interface{})
	for k, v := range in {
		res[fmt.Sprintf("%v", k)] = cleanupMapValue(v)
	}
	return res
}

func cleanupMapValue(v interface{}) interface{} {
	switch v := v.(type) {
	case []interface{}:
		return cleanupInterfaceArray(v)
	case map[interface{}]interface{}:
		return cleanupInterfaceMap(v)
	case nil, bool, string, int, float64:
		return v
	default:
		panic("Unhandled map type " + fmt.Sprintf("%T", v))
	}
}
<|endoftext|>


<commit_before>package comb

import "testing"

func TestMany(t *testing.T) {
	for _, str := range []string{"", "  "} {
		s := NewState(str)
		result, err := s.Many(s.Char(' '))()

		if err != nil {
			t.Error(err.Error())
		}

		t.Logf("%#v", result)
	}
}

func testMany1Space(str string) (interface{}, error) {
	s := NewState(str)
	return s.Many1(s.Char(' '))()
}

func TestMany1(t *testing.T) {
	result, err := testMany1Space(" ")

	if err != nil {
		t.Error(err.Error())
	}

	t.Logf("%#v", result)
}

func TestMany1Fail(t *testing.T) {
	result, err := testMany1Space("")

	if result != nil {
		t.Errorf("`result` should be nil but %#v.", result)
	}

	t.Log(err.Error())
}
<commit_msg>Test nested Many1<commit_after>package comb

import "testing"

func TestMany(t *testing.T) {
	for _, str := range []string{"", "  "} {
		s := NewState(str)
		result, err := s.Many(s.Char(' '))()

		if err != nil {
			t.Error(err.Error())
		}

		t.Logf("%#v", result)
	}
}

func testMany1Space(str string) (interface{}, error) {
	s := NewState(str)
	return s.Many1(s.Char(' '))()
}

func TestMany1(t *testing.T) {
	result, err := testMany1Space(" ")

	if err != nil {
		t.Error(err.Error())
	}

	t.Logf("%#v", result)
}

func TestMany1Fail(t *testing.T) {
	result, err := testMany1Space("")

	if result != nil {
		t.Errorf("`result` should be nil but %#v.", result)
	}

	t.Log(err.Error())
}

func TestMany1Nest(t *testing.T) {
	s := NewState("    ")
	result, err := s.Many1(s.Many1(s.Char(' ')))()

	if err != nil {
		t.Error(err.Error())
	}

	t.Logf("%#v", result)
}
<|endoftext|>


<commit_before>package lazycache

import "fmt"
import "net/http"
import "encoding/json"

// RootHandler is the default HTTP handler, registered at "/"
// It returns a JSON structure giving the relative path to
// each of the registered mirrors.
//
// e.g.
//
//   {
//     "https://rawdata.oceanobservatories.org/files/": {
//       "APIPath": {
//         "V1": "/v1/org/oceanobservatories/rawdata/files/"
//       }
//     }
//   }
func RootHandler(w http.ResponseWriter, req *http.Request) {

	// Temporary structures which define the output JSON structure
	type APIPathOut struct {
		V1 string
	}

	type RootMapOut struct {
		APIPath APIPathOut
	}

	jsonRootMap := make(map[string]RootMapOut)

	for key, root := range RootMap {
		jsonRootMap[key] = RootMapOut{
			APIPath: APIPathOut{
				V1: root.node.trimPath,
			},
		}
	}

	//	if jsonExtension.MatchString(req.URL.Path) {

	b, err := json.MarshalIndent(jsonRootMap, "", "  ")
	if err != nil {
		fmt.Fprintln(w, "JSON error:", err)
	}

	w.Write(b)
	//	} else {

	// 	fmt.Fprintf(w, "
")
	// 	for key, val := range RootMap {
	// 		fmt.Fprintf(w, "
%s
\n", val, key)
	// 	}
	// 	fmt.Fprintf(w, "
")
	// 	//fmt.Println("Indexing from ", req.URL.String())
	// }
}
<commit_msg>Removed some extraneous code from root_handler.<commit_after>package lazycache

import "fmt"
import "net/http"
import "encoding/json"

// RootHandler is the default HTTP handler, registered at "/"
// It returns a JSON structure giving the relative path to
// each of the registered mirrors.
//
// e.g.
//
//   {
//     "https://rawdata.oceanobservatories.org/files/": {
//       "APIPath": {
//         "V1": "/v1/org/oceanobservatories/rawdata/files/"
//       }
//     }
//   }
func RootHandler(w http.ResponseWriter, req *http.Request) {

	// Temporary structures which define the output JSON structure
	type APIPathOut struct {
		V1 string
	}

	type RootMapOut struct {
		APIPath APIPathOut
	}

	jsonRootMap := make(map[string]RootMapOut)

	for key, root := range RootMap {
		jsonRootMap[key] = RootMapOut{
			APIPath: APIPathOut{
				V1: root.node.trimPath,
			},
		}
	}

	//	if jsonExtension.MatchString(req.URL.Path) {

	b, err := json.MarshalIndent(jsonRootMap, "", "  ")
	if err != nil {
		fmt.Fprintln(w, "JSON error:", err)
	}

	w.Write(b)
}
<|endoftext|>


<commit_before>// Copyright 2017, TCN Inc.
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of TCN Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package utils_test

import (
	"testing"
	"time"
	utils "github.com/tcncloud/protoc-gen-persist/utils"
	google_protobuf "github.com/golang/protobuf/ptypes/timestamp"
	"github.com/lib/pq"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

)

func TestConversion(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Conversion methods Suite")
}

var _ = Describe("protobuf message value conversion functions", func() {
	var _ = Describe("ToSafeType", func() {
		It("can convert []string => *pq.StringArray", func() {
			in := []string{"hello", "world"}
			out := utils.ToSafeType(in)

			Expect(out).To(BeAssignableToTypeOf(&pq.StringArray{}))

			back, ok := out.(*pq.StringArray)

			Expect(ok).To(Equal(true))

			Expect([]string(*back)).To(BeEquivalentTo([]string{"hello", "world"}))
		})

		It("can convert []int64 => *pq.Int64Array", func() {
			in := []int64{1, 2}
			out := utils.ToSafeType(in)

			Expect(out).To(BeAssignableToTypeOf(&pq.Int64Array{}))

			back, ok := out.(*pq.Int64Array)

			Expect(ok).To(Equal(true))
			Expect([]int64(*back)).To(BeEquivalentTo([]int64{1, 2}))
		})

		It("can convert []float64 => *pq.Float64Array", func() {
			in := []float64{1.123, 2.234}
			out := utils.ToSafeType(in)

			Expect(out).To(BeAssignableToTypeOf(&pq.Float64Array{}))

			back, ok := out.(*pq.Float64Array)

			Expect(ok).To(Equal(true))
			Expect([]float64(*back)).To(BeEquivalentTo([]float64{1.123, 2.234}))
		})

		It("can convert *google protobuf Timestamp => *time.Time", func() {
			in := &google_protobuf.Timestamp{
				Seconds: int64(1234567),
				Nanos:   12345,
			}
			out := utils.ToSafeType(in)

			con, ok := out.(*time.Time)

			Expect(ok).To(Equal(true))

			back := utils.ToProtobufTime(con)

			Expect(back).To(BeEquivalentTo(&google_protobuf.Timestamp{
				Seconds: int64(1234567),
				Nanos:   12345,
			}))
		})

		It("can convert default Valuer types", func() {
			if _, ok := utils.ToSafeType("hello").(*string); !ok {
				Fail("failed on string")
			}
			if _, ok := utils.ToSafeType(int64(1234)).(*int64); !ok {
				Fail("failed on int64")
			}
			if _, ok := utils.ToSafeType(int32(1234)).(*int32); !ok {
				Fail("failed on int32")
			}
			if _, ok := utils.ToSafeType(float64(1234.1234)).(*float64); !ok {
				Fail("failed on float64")
			}
		})
	})

	var _ = Describe("AssignTo", func() {
		It("can assign a pq.StringArray value to a *[]string key", func() {
			key := new([]string)
			val := pq.StringArray{"first", "second"}

			utils.AssignTo(key, val)
			Expect(*key).To(Equal([]string{"first", "second"}))
		})

		It(" can assign a pq.Int64Array value to a *[]int64 key", func() {
			key := new([]int64)
			val := pq.Int64Array{int64(1), int64(2)}

			utils.AssignTo(key, val)
			Expect(*key).To(Equal([]int64{int64(1), int64(2)}))
		})

		It("can assign a pq.Float64Array value to a *[]float64 key", func() {
			key := new([]float64)
			val := pq.Float64Array{1.1234, 2.2345}

			utils.AssignTo(key, val)
			Expect(*key).To(Equal([]float64{1.1234, 2.2345}))
		})

		It("can assign a int64 value to a *int64 key", func() {
			key := new(int64)
			val := int64(1234)

			utils.AssignTo(key, val)
			Expect(*key).To(Equal(int64(1234)))
		})

		It("can assign a int32 value to a *int32 key", func() {
			key := new(int32)
			val := int32(1234)

			utils.AssignTo(key, val)
			Expect(*key).To(Equal(int32(1234)))
		})

		It("can assign a float64 value to a *float64 key", func() {
			key := new(float64)
			val := 100.1234

			utils.AssignTo(key, val)
			Expect(*key).To(Equal(100.1234))
		})

		It("can assign a string value to a *string key", func() {
			key := new(string)
			val := "hello world"

			utils.AssignTo(key, val)
			Expect(*key).To(Equal("hello world"))
		})

		It("can assign a *time.Time value to a **google Timestamp key", func() {
			key := new(*google_protobuf.Timestamp)
			val := new(time.Time)
			*val = time.Now().UTC().Truncate(time.Microsecond)

			utils.AssignTo(key, val)

			convertedBack := ((*utils.ToTime(*key)).Truncate(time.Microsecond))
			Expect(val.String()).To(Equal(convertedBack.String()))
		})

		It("returns true if value is assigned to key", func() {
			key := new(string)
			val := ""

			res := utils.AssignTo(key, val)
			Expect(res).To(Equal(true))
		})

		var _ = Context("when no supported value is placed as val", func() {
			It("returns false", func() {
				key := new([]string)
				val := []string{"not", "supported"}

				res := utils.AssignTo(key, val)

				Expect(res).To(Equal(false))
			})

			It("does not assign to key", func() {
				key := new([]string)
				val := []string{"not", "supported"}

				utils.AssignTo(key, val)
				k := new([]string)
				Expect(key).To(BeEquivalentTo(k))
			})
		})

		var _ = Context("when key cannot convert to correct type", func() {
			It("returns false", func() {
				key := new(int64)
				val := []string{"not", "integers"}

				res := utils.AssignTo(key, val)
				Expect(res).To(Equal(false))
			})

			It("does not assign to key", func() {
				key := new(int64)
				val := []string{"not", "integers"}

				utils.AssignTo(key, val)
				Expect(key).To(BeEquivalentTo(new(int64)))
			})
		})

		var _ = Context("when key is nil", func() {
			It("returns false", func() {
				var key *int64
				val := int64(12345)

				res := utils.AssignTo(key, val)
				Expect(res).To(Equal(false))
			})

			It("does not assign to key", func() {
				var key *int64
				val := int64(12345)

				utils.AssignTo(key, val)
				Expect(key).To(BeZero())
			})
		})
	})
})
<commit_msg>formatting<commit_after>// Copyright 2017, TCN Inc.
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of TCN Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package utils_test

import (
	"testing"
	"time"
	utils "github.com/tcncloud/protoc-gen-persist/utils"
	google_protobuf "github.com/golang/protobuf/ptypes/timestamp"
	"github.com/lib/pq"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

)

func TestConversion(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Conversion methods Suite")
}

var _ = Describe("protobuf message value conversion functions", func() {
	var _ = Describe("ToSafeType", func() {
		It("can convert []string => *pq.StringArray", func() {
			in := []string{"hello", "world"}
			out := utils.ToSafeType(in)

			Expect(out).To(BeAssignableToTypeOf(&pq.StringArray{}))

			back, ok := out.(*pq.StringArray)

			Expect(ok).To(Equal(true))

			Expect([]string(*back)).To(BeEquivalentTo([]string{"hello", "world"}))
		})

		It("can convert []int64 => *pq.Int64Array", func() {
			in := []int64{1, 2}
			out := utils.ToSafeType(in)

			Expect(out).To(BeAssignableToTypeOf(&pq.Int64Array{}))

			back, ok := out.(*pq.Int64Array)

			Expect(ok).To(Equal(true))
			Expect([]int64(*back)).To(BeEquivalentTo([]int64{1, 2}))
		})

		It("can convert []float64 => *pq.Float64Array", func() {
			in := []float64{1.123, 2.234}
			out := utils.ToSafeType(in)

			Expect(out).To(BeAssignableToTypeOf(&pq.Float64Array{}))

			back, ok := out.(*pq.Float64Array)

			Expect(ok).To(Equal(true))
			Expect([]float64(*back)).To(BeEquivalentTo([]float64{1.123, 2.234}))
		})

		It("can convert *google protobuf Timestamp => *time.Time", func() {
			in := &google_protobuf.Timestamp{
				Seconds: int64(1234567),
				Nanos:   12345,
			}
			out := utils.ToSafeType(in)

			con, ok := out.(*time.Time)

			Expect(ok).To(Equal(true))

			back := utils.ToProtobufTime(con)

			Expect(back).To(BeEquivalentTo(&google_protobuf.Timestamp{
				Seconds: int64(1234567),
				Nanos:   12345,
			}))
		})

		It("can convert default Valuer types", func() {
			if _, ok := utils.ToSafeType("hello").(*string); !ok {
				Fail("failed on string")
			}
			if _, ok := utils.ToSafeType(int64(1234)).(*int64); !ok {
				Fail("failed on int64")
			}
			if _, ok := utils.ToSafeType(int32(1234)).(*int32); !ok {
				Fail("failed on int32")
			}
			if _, ok := utils.ToSafeType(float64(1234.1234)).(*float64); !ok {
				Fail("failed on float64")
			}
		})
	})

	var _ = Describe("AssignTo", func() {
		It("can assign a pq.StringArray value to a *[]string key", func() {
			key := new([]string)
			val := pq.StringArray{"first", "second"}

			utils.AssignTo(key, val)
			Expect(*key).To(Equal([]string{"first", "second"}))
		})

		It(" can assign a pq.Int64Array value to a *[]int64 key", func() {
			key := new([]int64)
			val := pq.Int64Array{int64(1), int64(2)}

			utils.AssignTo(key, val)
			Expect(*key).To(Equal([]int64{int64(1), int64(2)}))
		})

		It("can assign a pq.Float64Array value to a *[]float64 key", func() {
			key := new([]float64)
			val := pq.Float64Array{1.1234, 2.2345}

			utils.AssignTo(key, val)
			Expect(*key).To(Equal([]float64{1.1234, 2.2345}))
		})

		It("can assign a int64 value to a *int64 key", func() {
			key := new(int64)
			val := int64(1234)

			utils.AssignTo(key, val)
			Expect(*key).To(Equal(int64(1234)))
		})

		It("can assign a int32 value to a *int32 key", func() {
			key := new(int32)
			val := int32(1234)

			utils.AssignTo(key, val)
			Expect(*key).To(Equal(int32(1234)))
		})

		It("can assign a float64 value to a *float64 key", func() {
			key := new(float64)
			val := 100.1234

			utils.AssignTo(key, val)
			Expect(*key).To(Equal(100.1234))
		})

		It("can assign a string value to a *string key", func() {
			key := new(string)
			val := "hello world"

			utils.AssignTo(key, val)
			Expect(*key).To(Equal("hello world"))
		})

		It("can assign a *time.Time value to a **google Timestamp key", func() {
			key := new(*google_protobuf.Timestamp)
			val := new(time.Time)
			*val = time.Now().UTC().Truncate(time.Microsecond)

			utils.AssignTo(key, val)

			convertedBack := ((*utils.ToTime(*key)).Truncate(time.Microsecond))
			Expect(val.String()).To(Equal(convertedBack.String()))
		})

		It("returns true if value is assigned to key", func() {
			key := new(string)
			val := ""

			res := utils.AssignTo(key, val)
			Expect(res).To(Equal(true))
		})

		var _ = Context("when no supported value is placed as val", func() {
			It("returns false", func() {
				key := new([]string)
				val := []string{"not", "supported"}

				res := utils.AssignTo(key, val)

				Expect(res).To(Equal(false))
			})

			It("does not assign to key", func() {
				key := new([]string)
				val := []string{"not", "supported"}

				utils.AssignTo(key, val)
				k := new([]string)
				Expect(key).To(BeEquivalentTo(k))
			})
		})

		var _ = Context("when key cannot convert to correct type", func() {
			It("returns false", func() {
				key := new(int64)
				val := []string{"not", "integers"}

				res := utils.AssignTo(key, val)
				Expect(res).To(Equal(false))
			})

			It("does not assign to key", func() {
				key := new(int64)
				val := []string{"not", "integers"}

				utils.AssignTo(key, val)
				Expect(key).To(BeEquivalentTo(new(int64)))
			})
		})

		var _ = Context("when key is nil", func() {
			It("returns false", func() {
				var key *int64
				val := int64(12345)

				res := utils.AssignTo(key, val)
				Expect(res).To(Equal(false))
			})

			It("does not assign to key", func() {
				var key *int64
				val := int64(12345)

				utils.AssignTo(key, val)
				Expect(key).To(BeZero())
			})
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sources

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/OWASP/Amass/amass/core"
	"github.com/OWASP/Amass/amass/utils"
	"github.com/dghubble/go-twitter/twitter"
	"golang.org/x/oauth2"
)

// Twitter is the Service that handles access to the Twitter data source.
type Twitter struct {
	core.BaseService

	API        *core.APIKey
	SourceType string
	RateLimit  time.Duration
	client     *twitter.Client
}

// NewTwitter returns he object initialized, but not yet started.
func NewTwitter(config *core.Config, bus *core.EventBus) *Twitter {
	t := &Twitter{
		SourceType: core.API,
		RateLimit:  3 * time.Second,
	}

	t.BaseService = *core.NewBaseService(t, "Twitter", config, bus)
	return t
}

// OnStart implements the Service interface
func (t *Twitter) OnStart() error {
	t.BaseService.OnStart()

	t.API = t.Config().GetAPIKey(t.String())
	if t.API == nil || t.API.Key == "" || t.API.Secret == "" {
		t.Config().Log.Printf("%s: API key data was not provided", t.String())
	}
	if t.API != nil && t.API.Key != "" && t.API.Secret != "" {
		if bearer, err := t.getBearerToken(); err == nil {
			config := &oauth2.Config{}
			token := &oauth2.Token{AccessToken: bearer}
			// OAuth2 http.Client will automatically authorize Requests
			httpClient := config.Client(oauth2.NoContext, token)
			// Twitter client
			t.client = twitter.NewClient(httpClient)
		}
	}

	go t.processRequests()
	return nil
}

func (t *Twitter) processRequests() {
	last := time.Now()

	for {
		select {
		case <-t.Quit():
			return
		case req := <-t.DNSRequestChan():
			if t.Config().IsDomainInScope(req.Domain) {
				if time.Now().Sub(last) < t.RateLimit {
					time.Sleep(t.RateLimit)
				}
				last = time.Now()
				t.executeQuery(req.Domain)
				last = time.Now()
			}
		case <-t.AddrRequestChan():
		case <-t.ASNRequestChan():
		case <-t.WhoisRequestChan():
		}
	}
}

func (t *Twitter) executeQuery(domain string) {
	re := t.Config().DomainRegex(domain)
	if t.client == nil || re == nil {
		return
	}

	searchParams := &twitter.SearchTweetParams{
		Query: domain,
		Count: 100,
	}
	t.SetActive()
	search, _, err := t.client.Search.Tweets(searchParams)
	if err != nil {
		t.Config().Log.Printf("%s: %v", t.String(), err)
		return
	}

	for _, tweet := range search.Statuses {

		// Urls in the tweet body
		for _, url_entity := range tweet.Entities.Urls {
			for _, name := range re.FindAllString(url_entity.ExpandedURL, -1) {
				t.Bus().Publish(core.NewNameTopic, &core.DNSRequest{
					Name:   name,
					Domain: domain,
					Tag:    t.SourceType,
					Source: t.String(),
				})
			}
		}

		// Source of the tweet
		for _, name := range re.FindAllString(tweet.Source, -1) {
			t.Bus().Publish(core.NewNameTopic, &core.DNSRequest{
				Name:   name,
				Domain: domain,
				Tag:    t.SourceType,
				Source: t.String(),
			})
		}
	}
}

func (t *Twitter) getBearerToken() (string, error) {
	headers := map[string]string{"Content-Type": "application/x-www-form-urlencoded;charset=UTF-8"}
	page, err := utils.RequestWebPage(
		"https://api.twitter.com/oauth2/token",
		strings.NewReader("grant_type=client_credentials"),
		headers, t.API.Key, t.API.Secret)
	if err != nil {
		return "", fmt.Errorf("token request failed: %+v", err)
	}

	var v struct {
		AccessToken string `json:"access_token"`
	}
	if err := json.Unmarshal([]byte(page), &v); err != nil {
		return "", fmt.Errorf("error parsing json in token response: %+v", err)
	}
	if v.AccessToken == "" {
		return "", fmt.Errorf("token response does not have access_token")
	}
	return v.AccessToken, nil
}
<commit_msg>changed variables names according to the case convention<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sources

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/OWASP/Amass/amass/core"
	"github.com/OWASP/Amass/amass/utils"
	"github.com/dghubble/go-twitter/twitter"
	"golang.org/x/oauth2"
)

// Twitter is the Service that handles access to the Twitter data source.
type Twitter struct {
	core.BaseService

	API        *core.APIKey
	SourceType string
	RateLimit  time.Duration
	client     *twitter.Client
}

// NewTwitter returns he object initialized, but not yet started.
func NewTwitter(config *core.Config, bus *core.EventBus) *Twitter {
	t := &Twitter{
		SourceType: core.API,
		RateLimit:  3 * time.Second,
	}

	t.BaseService = *core.NewBaseService(t, "Twitter", config, bus)
	return t
}

// OnStart implements the Service interface
func (t *Twitter) OnStart() error {
	t.BaseService.OnStart()

	t.API = t.Config().GetAPIKey(t.String())
	if t.API == nil || t.API.Key == "" || t.API.Secret == "" {
		t.Config().Log.Printf("%s: API key data was not provided", t.String())
	}
	if t.API != nil && t.API.Key != "" && t.API.Secret != "" {
		if bearer, err := t.getBearerToken(); err == nil {
			config := &oauth2.Config{}
			token := &oauth2.Token{AccessToken: bearer}
			// OAuth2 http.Client will automatically authorize Requests
			httpClient := config.Client(oauth2.NoContext, token)
			// Twitter client
			t.client = twitter.NewClient(httpClient)
		}
	}

	go t.processRequests()
	return nil
}

func (t *Twitter) processRequests() {
	last := time.Now()

	for {
		select {
		case <-t.Quit():
			return
		case req := <-t.DNSRequestChan():
			if t.Config().IsDomainInScope(req.Domain) {
				if time.Now().Sub(last) < t.RateLimit {
					time.Sleep(t.RateLimit)
				}
				last = time.Now()
				t.executeQuery(req.Domain)
				last = time.Now()
			}
		case <-t.AddrRequestChan():
		case <-t.ASNRequestChan():
		case <-t.WhoisRequestChan():
		}
	}
}

func (t *Twitter) executeQuery(domain string) {
	re := t.Config().DomainRegex(domain)
	if t.client == nil || re == nil {
		return
	}

	searchParams := &twitter.SearchTweetParams{
		Query: domain,
		Count: 100,
	}
	t.SetActive()
	search, _, err := t.client.Search.Tweets(searchParams)
	if err != nil {
		t.Config().Log.Printf("%s: %v", t.String(), err)
		return
	}

	for _, tweet := range search.Statuses {

		// Urls in the tweet body
		for _, urlEntity := range tweet.Entities.Urls {
			for _, name := range re.FindAllString(urlEntity.ExpandedURL, -1) {
				t.Bus().Publish(core.NewNameTopic, &core.DNSRequest{
					Name:   name,
					Domain: domain,
					Tag:    t.SourceType,
					Source: t.String(),
				})
			}
		}

		// Source of the tweet
		for _, name := range re.FindAllString(tweet.Source, -1) {
			t.Bus().Publish(core.NewNameTopic, &core.DNSRequest{
				Name:   name,
				Domain: domain,
				Tag:    t.SourceType,
				Source: t.String(),
			})
		}
	}
}

func (t *Twitter) getBearerToken() (string, error) {
	headers := map[string]string{"Content-Type": "application/x-www-form-urlencoded;charset=UTF-8"}
	page, err := utils.RequestWebPage(
		"https://api.twitter.com/oauth2/token",
		strings.NewReader("grant_type=client_credentials"),
		headers, t.API.Key, t.API.Secret)
	if err != nil {
		return "", fmt.Errorf("token request failed: %+v", err)
	}

	var v struct {
		AccessToken string `json:"access_token"`
	}
	if err := json.Unmarshal([]byte(page), &v); err != nil {
		return "", fmt.Errorf("error parsing json in token response: %+v", err)
	}
	if v.AccessToken == "" {
		return "", fmt.Errorf("token response does not have access_token")
	}
	return v.AccessToken, nil
}
<|endoftext|>


<commit_before>package config

import (
	"context"
	"fmt"
	"log"
	"net/url"
	"time"

	"github.com/grafeas/voucher/v2/metrics"
	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc"
	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp"
	"go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/resource"
	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
)

func MetricsClient(secrets *Secrets) (metrics.Client, error) {
	tags := viper.GetStringSlice("statsd.tags")

	switch backend := viper.GetString("statsd.backend"); backend {
	case "statsd", "":
		if statsdAddr := viper.GetString("statsd.addr"); statsdAddr != "" {
			sampleRate := viper.GetFloat64("statsd.sample_rate")
			return metrics.NewStatsdClient(statsdAddr, sampleRate, tags)
		}

		log.Printf("No metrics client configured")
		return &metrics.NoopClient{}, nil
	case "otel", "opentelemetry":
		ctx := context.Background()
		exporter, err := otelExporter(ctx)
		if err != nil {
			return nil, fmt.Errorf("creating otel exporter: %w", err)
		}
		interval := viper.GetDuration("statsd.interval")
		if interval == 0 {
			interval = time.Minute
		}
		res, err := resource.New(ctx, resource.WithAttributes(
			semconv.ServiceNameKey.String("voucher"),
		))
		if err != nil {
			return nil, fmt.Errorf("creating otel resource: %w", err)
		}

		mp := metric.NewMeterProvider(
			metric.WithResource(res),
			metric.WithReader(metric.NewPeriodicReader(exporter, metric.WithInterval(interval))),
		)
		return metrics.NewOpenTelemetryClient(mp, exporter)
	case "datadog":
		if secrets != nil && secrets.Datadog.APIKey != "" && secrets.Datadog.AppKey != "" {
			return metrics.NewDatadogClient(secrets.Datadog.APIKey, secrets.Datadog.AppKey, metrics.WithDatadogTags(tags)), nil
		}
		return &metrics.NoopClient{}, fmt.Errorf("missing secrets for datadog")
	default:
		return &metrics.NoopClient{}, fmt.Errorf("unknown statsd backend: %s", backend)
	}
}

func otelExporter(ctx context.Context) (metric.Exporter, error) {
	insecure := viper.GetBool("statsd.insecure")

	addr := viper.GetString("statsd.addr")
	otelURL, err := url.Parse(addr)
	if err != nil {
		return nil, fmt.Errorf("parsing otel url: %w", err)
	}

	log := logrus.WithFields(logrus.Fields{
		"otel_addr": addr,
		"insecure":  insecure,
		"scheme":    otelURL.Scheme,
		"host":      otelURL.Host,
	})

	switch otelURL.Scheme {
	case "grpc":
		opts := []otlpmetricgrpc.Option{
			otlpmetricgrpc.WithEndpoint(otelURL.Host),
		}
		if insecure {
			opts = append(opts, otlpmetricgrpc.WithInsecure())
		}
		log.Info("creating otel exporter")
		return otlpmetricgrpc.New(ctx, opts...)
	case "http", "https":
		opts := []otlpmetrichttp.Option{
			otlpmetrichttp.WithEndpoint(otelURL.Host),
		}
		if insecure {
			opts = append(opts, otlpmetrichttp.WithInsecure())
		}
		log.Info("creating otel exporter")
		return otlpmetrichttp.New(ctx, opts...)
	default:
		return nil, fmt.Errorf("unknown otel scheme: %s", otelURL.Scheme)
	}
}
<commit_msg>otel-metrics: tags too<commit_after>package config

import (
	"context"
	"fmt"
	"log"
	"net/url"
	"strings"
	"time"

	"github.com/grafeas/voucher/v2/metrics"
	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc"
	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp"
	"go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/resource"
	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
)

func MetricsClient(secrets *Secrets) (metrics.Client, error) {
	tags := viper.GetStringSlice("statsd.tags")

	switch backend := viper.GetString("statsd.backend"); backend {
	case "statsd", "":
		if statsdAddr := viper.GetString("statsd.addr"); statsdAddr != "" {
			sampleRate := viper.GetFloat64("statsd.sample_rate")
			return metrics.NewStatsdClient(statsdAddr, sampleRate, tags)
		}

		log.Printf("No metrics client configured")
		return &metrics.NoopClient{}, nil
	case "otel", "opentelemetry":
		ctx := context.Background()
		exporter, err := otelExporter(ctx)
		if err != nil {
			return nil, fmt.Errorf("creating otel exporter: %w", err)
		}
		interval := viper.GetDuration("statsd.interval")
		if interval == 0 {
			interval = time.Minute
		}

		attrs := make([]attribute.KeyValue, 0, len(tags)+1)
		for _, tag := range tags {
			s := strings.SplitN(tag, ":", 2)
			attrs = append(attrs, attribute.String(s[0], s[1]))
		}
		attrs = append(attrs, semconv.ServiceNameKey.String("voucher"))

		res, err := resource.New(ctx, resource.WithAttributes(attrs...))
		if err != nil {
			return nil, fmt.Errorf("creating otel resource: %w", err)
		}

		mp := metric.NewMeterProvider(
			metric.WithResource(res),
			metric.WithReader(metric.NewPeriodicReader(exporter, metric.WithInterval(interval))),
		)
		return metrics.NewOpenTelemetryClient(mp, exporter)
	case "datadog":
		if secrets != nil && secrets.Datadog.APIKey != "" && secrets.Datadog.AppKey != "" {
			return metrics.NewDatadogClient(secrets.Datadog.APIKey, secrets.Datadog.AppKey, metrics.WithDatadogTags(tags)), nil
		}
		return &metrics.NoopClient{}, fmt.Errorf("missing secrets for datadog")
	default:
		return &metrics.NoopClient{}, fmt.Errorf("unknown statsd backend: %s", backend)
	}
}

func otelExporter(ctx context.Context) (metric.Exporter, error) {
	insecure := viper.GetBool("statsd.insecure")

	addr := viper.GetString("statsd.addr")
	otelURL, err := url.Parse(addr)
	if err != nil {
		return nil, fmt.Errorf("parsing otel url: %w", err)
	}

	log := logrus.WithFields(logrus.Fields{
		"otel_addr": addr,
		"insecure":  insecure,
		"scheme":    otelURL.Scheme,
		"host":      otelURL.Host,
	})

	switch otelURL.Scheme {
	case "grpc":
		opts := []otlpmetricgrpc.Option{
			otlpmetricgrpc.WithEndpoint(otelURL.Host),
		}
		if insecure {
			opts = append(opts, otlpmetricgrpc.WithInsecure())
		}
		log.Info("creating otel exporter")
		return otlpmetricgrpc.New(ctx, opts...)
	case "http", "https":
		opts := []otlpmetrichttp.Option{
			otlpmetrichttp.WithEndpoint(otelURL.Host),
		}
		if insecure {
			opts = append(opts, otlpmetrichttp.WithInsecure())
		}
		log.Info("creating otel exporter")
		return otlpmetrichttp.New(ctx, opts...)
	default:
		return nil, fmt.Errorf("unknown otel scheme: %s", otelURL.Scheme)
	}
}
<|endoftext|>


<commit_before>package models

import "github.com/astaxie/beego/orm"

type Dashboard struct {
	BookNumber int64                `json:"book_number"`
	DocumentNumber int64        `json:"document_number"`
	MemberNumber int64        `json:"member_number"`
	CommentNumber int64        `json:"comment_number"`
	AttachmentNumber int64        `json:"attachment_number"`
}

func NewDashboard() *Dashboard {
	return &Dashboard{}
}

func (m *Dashboard) Query() (*Dashboard) {
	o := orm.NewOrm()

	book_number,_ := o.QueryTable(NewBook().TableNameWithPrefix()).Count()

	m.BookNumber = book_number

	document_count,_ := o.QueryTable(NewDocument().TableNameWithPrefix()).Count()
	m.DocumentNumber = document_count

	member_number,_ := o.QueryTable(NewMember().TableNameWithPrefix()).Count()
	m.MemberNumber = member_number

	comment_number,_ := o.QueryTable(NewComment().TableNameWithPrefix()).Count()
	m.CommentNumber = comment_number

	attachment_number,_ := o.QueryTable(NewAttachment().TableNameWithPrefix()).Count()

	m.AttachmentNumber = attachment_number

	return m
}
<commit_msg>屏蔽评论功能<commit_after>package models

import "github.com/astaxie/beego/orm"

type Dashboard struct {
	BookNumber int64                `json:"book_number"`
	DocumentNumber int64        `json:"document_number"`
	MemberNumber int64        `json:"member_number"`
	CommentNumber int64        `json:"comment_number"`
	AttachmentNumber int64        `json:"attachment_number"`
}

func NewDashboard() *Dashboard {
	return &Dashboard{}
}

func (m *Dashboard) Query() (*Dashboard) {
	o := orm.NewOrm()

	book_number,_ := o.QueryTable(NewBook().TableNameWithPrefix()).Count()

	m.BookNumber = book_number

	document_count,_ := o.QueryTable(NewDocument().TableNameWithPrefix()).Count()
	m.DocumentNumber = document_count

	member_number,_ := o.QueryTable(NewMember().TableNameWithPrefix()).Count()
	m.MemberNumber = member_number

	//comment_number,_ := o.QueryTable(NewComment().TableNameWithPrefix()).Count()
	m.CommentNumber = 0

	attachment_number,_ := o.QueryTable(NewAttachment().TableNameWithPrefix()).Count()

	m.AttachmentNumber = attachment_number

	return m
}
<|endoftext|>


<commit_before>package migrations

func m27_create_table_youtube() {
	CreateTableIfNotExists("youtube")
}
<commit_msg>[youtube] fixes typo<commit_after>package migrations

func m29_create_table_youtube() {
	CreateTableIfNotExists("youtube")
}
<|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 system

import (
	"fmt"
)

// Validator is the interface for all validators.
type Validator interface {
	// Name is the name of the validator.
	Name() string
	// Validate is the validate function.
	Validate(SysSpec) ([]error, []error)
}

// Reporter is the interface for the reporters for the validators.
type Reporter interface {
	// Report reports the results of the system verification
	Report(string, string, ValidationResultType) error
}

// Validate uses validators to validate the system and returns a warning or error.
func Validate(spec SysSpec, validators []Validator) ([]error, []error) {
	var errs []error
	var warns []error

	for _, v := range validators {
		fmt.Printf("Validating %s...\n", v.Name())
		warn, err := v.Validate(spec)
		if len(err) != 0 {
			errs = append(errs, err...)
		}
		if len(warn) != 0 {
			warns = append(warns, warn...)
		}
	}
	return warns, errs
}

// ValidateSpec uses all default validators to validate the system and writes to stdout.
func ValidateSpec(spec SysSpec, runtime string) ([]error, []error) {
	// OS-level validators.
	var osValidators = []Validator{
		&OSValidator{Reporter: DefaultReporter},
		&KernelValidator{Reporter: DefaultReporter},
		&CgroupsValidator{Reporter: DefaultReporter},
		&packageValidator{reporter: DefaultReporter},
	}
	// Docker-specific validators.
	var dockerValidators = []Validator{
		&DockerValidator{Reporter: DefaultReporter},
	}

	validators := osValidators
	switch runtime {
	case "docker":
		validators = append(validators, dockerValidators...)
	}
	return Validate(spec, validators)
}
<commit_msg>ValidateSpec: don't include cgroups and package validators on non-Linux<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 system

import (
	"fmt"
	"runtime"
)

// Validator is the interface for all validators.
type Validator interface {
	// Name is the name of the validator.
	Name() string
	// Validate is the validate function.
	Validate(SysSpec) ([]error, []error)
}

// Reporter is the interface for the reporters for the validators.
type Reporter interface {
	// Report reports the results of the system verification
	Report(string, string, ValidationResultType) error
}

// Validate uses validators to validate the system and returns a warning or error.
func Validate(spec SysSpec, validators []Validator) ([]error, []error) {
	var errs []error
	var warns []error

	for _, v := range validators {
		fmt.Printf("Validating %s...\n", v.Name())
		warn, err := v.Validate(spec)
		if len(err) != 0 {
			errs = append(errs, err...)
		}
		if len(warn) != 0 {
			warns = append(warns, warn...)
		}
	}
	return warns, errs
}

// ValidateSpec uses all default validators to validate the system and writes to stdout.
func ValidateSpec(spec SysSpec, containerRuntime string) ([]error, []error) {
	// OS-level validators.
	var osValidators = []Validator{
		&OSValidator{Reporter: DefaultReporter},
		&KernelValidator{Reporter: DefaultReporter},
	}

	// Docker-specific validators.
	var dockerValidators = []Validator{
		&DockerValidator{Reporter: DefaultReporter},
	}

	validators := osValidators
	switch containerRuntime {
	case "docker":
		validators = append(validators, dockerValidators...)
	}

	// Linux-specific validators.
	if runtime.GOOS == "linux" {
		validators = append(validators,
			&CgroupsValidator{Reporter: DefaultReporter},
			&packageValidator{reporter: DefaultReporter},
		)
	}

	return Validate(spec, validators)
}
<|endoftext|>