<commit_before>// Copyright 2014 The Prometheus 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 storage

import (
	"encoding/gob"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/model"

	dto "github.com/prometheus/client_model/go"
)

const (
	writeQueueCapacity = 1000
)

// DiskMetricStore is an implementation of MetricStore that persists metrics to
// disk.
type DiskMetricStore struct {
	lock            sync.RWMutex // Protects metricFamilies.
	writeQueue      chan WriteRequest
	drain           chan struct{}
	done            chan error
	metricGroups    GroupingKeyToMetricGroup
	persistenceFile string
}

type mfStat struct {
	pos    int  // Where in the result slice is the MetricFamily?
	copied bool // Has the MetricFamily already been copied?
}

// NewDiskMetricStore returns a DiskMetricStore ready to use. To cleanly shut it
// down and free resources, the Shutdown() method has to be called.  If
// persistenceFile is the empty string, no persisting to disk will
// happen. Otherwise, a file of that name is used for persisting metrics to
// disk. If the file already exists, metrics are read from it as part of the
// start-up. Persisting is happening upon shutdown and after every write action,
// but the latter will only happen persistenceDuration after the previous
// persisting.
func NewDiskMetricStore(
	persistenceFile string,
	persistenceInterval time.Duration,
) *DiskMetricStore {
	// TODO: Do that outside of the constructor to allow the HTTP server to
	//  serve /-/healthy and /-/ready earlier.
	dms := &DiskMetricStore{
		writeQueue:      make(chan WriteRequest, writeQueueCapacity),
		drain:           make(chan struct{}),
		done:            make(chan error),
		metricGroups:    GroupingKeyToMetricGroup{},
		persistenceFile: persistenceFile,
	}
	if err := dms.restore(); err != nil {
		log.Errorln("Could not load persisted metrics:", err)
		log.Info("Retrying assuming legacy format for persisted metrics...")
		if err := dms.legacyRestore(); err != nil {
			log.Errorln("Could not load persisted metrics in legacy format: ", err)
		}
	}

	go dms.loop(persistenceInterval)
	return dms
}

// SubmitWriteRequest implements the MetricStore interface.
func (dms *DiskMetricStore) SubmitWriteRequest(req WriteRequest) {
	dms.writeQueue <- req
}

// GetMetricFamilies implements the MetricStore interface.
func (dms *DiskMetricStore) GetMetricFamilies() []*dto.MetricFamily {
	result := []*dto.MetricFamily{}
	mfStatByName := map[string]mfStat{}

	dms.lock.RLock()
	defer dms.lock.RUnlock()

	for _, group := range dms.metricGroups {
		for name, tmf := range group.Metrics {
			mf := tmf.MetricFamily
			stat, exists := mfStatByName[name]
			if exists {
				existingMF := result[stat.pos]
				if !stat.copied {
					mfStatByName[name] = mfStat{
						pos:    stat.pos,
						copied: true,
					}
					existingMF = copyMetricFamily(existingMF)
					result[stat.pos] = existingMF
				}
				if mf.GetHelp() != existingMF.GetHelp() || mf.GetType() != existingMF.GetType() {
					log.Infof(
						"Metric families '%s' and '%s' are inconsistent, help and type of the latter will have priority. This is bad. Fix your pushed metrics!",
						mf, existingMF,
					)
				}
				for _, metric := range mf.Metric {
					existingMF.Metric = append(existingMF.Metric, metric)
				}
			} else {
				mfStatByName[name] = mfStat{
					pos:    len(result),
					copied: false,
				}
				result = append(result, mf)
			}
		}
	}
	return result
}

// Shutdown implements the MetricStore interface.
func (dms *DiskMetricStore) Shutdown() error {
	close(dms.drain)
	return <-dms.done
}

// Healthy implements the MetricStore interface.
func (dms *DiskMetricStore) Healthy() error {
	// By taking the lock we check that there is no deadlock.
	dms.lock.Lock()
	defer dms.lock.Unlock()

	// A pushgateway that cannot be written to should not be
	// considered as healthy.
	if len(dms.writeQueue) == cap(dms.writeQueue) {
		return fmt.Errorf("write queue is full")
	}

	return nil
}

// Ready implements the MetricStore interface.
func (dms *DiskMetricStore) Ready() error {
	return dms.Healthy()
}

func (dms *DiskMetricStore) loop(persistenceInterval time.Duration) {
	lastPersist := time.Now()
	persistScheduled := false
	lastWrite := time.Time{}
	persistDone := make(chan time.Time)
	var persistTimer *time.Timer

	checkPersist := func() {
		if !persistScheduled && lastWrite.After(lastPersist) {
			persistTimer = time.AfterFunc(
				persistenceInterval-lastWrite.Sub(lastPersist),
				func() {
					persistStarted := time.Now()
					if err := dms.persist(); err != nil {
						log.Errorln("Error persisting metrics:", err)
					} else {
						log.Infof(
							"Metrics persisted to '%s'.",
							dms.persistenceFile,
						)
					}
					persistDone <- persistStarted
				},
			)
			persistScheduled = true
		}
	}

	for {
		select {
		case wr := <-dms.writeQueue:
			dms.processWriteRequest(wr)
			lastWrite = time.Now()
			checkPersist()
		case lastPersist = <-persistDone:
			persistScheduled = false
			checkPersist() // In case something has been written in the meantime.
		case <-dms.drain:
			// Prevent a scheduled persist from firing later.
			if persistTimer != nil {
				persistTimer.Stop()
			}
			// Now draining...
			for {
				select {
				case wr := <-dms.writeQueue:
					dms.processWriteRequest(wr)
				default:
					dms.done <- dms.persist()
					return
				}
			}
		}
	}
}

func (dms *DiskMetricStore) processWriteRequest(wr WriteRequest) {
	dms.lock.Lock()
	defer dms.lock.Unlock()

	key := model.LabelsToSignature(wr.Labels)

	if wr.MetricFamilies == nil {
		// Delete.
		delete(dms.metricGroups, key)
		return
	}
	// Update.
	for name, mf := range wr.MetricFamilies {
		group, ok := dms.metricGroups[key]
		if !ok {
			group = MetricGroup{
				Labels:  wr.Labels,
				Metrics: NameToTimestampedMetricFamilyMap{},
			}
			dms.metricGroups[key] = group
		}
		group.Metrics[name] = TimestampedMetricFamily{
			Timestamp:    wr.Timestamp,
			MetricFamily: mf,
		}
	}
}

// GetMetricFamiliesMap implements the MetricStore interface.
func (dms *DiskMetricStore) GetMetricFamiliesMap() GroupingKeyToMetricGroup {
	dms.lock.RLock()
	defer dms.lock.RUnlock()
	groupsCopy := make(GroupingKeyToMetricGroup, len(dms.metricGroups))
	for k, g := range dms.metricGroups {
		metricsCopy := make(NameToTimestampedMetricFamilyMap, len(g.Metrics))
		groupsCopy[k] = MetricGroup{Labels: g.Labels, Metrics: metricsCopy}
		for n, tmf := range g.Metrics {
			metricsCopy[n] = tmf
		}
	}
	return groupsCopy
}

func (dms *DiskMetricStore) persist() error {
	if dms.persistenceFile == "" {
		return nil
	}
	f, err := ioutil.TempFile(
		path.Dir(dms.persistenceFile),
		path.Base(dms.persistenceFile)+".in_progress.",
	)
	if err != nil {
		return err
	}
	inProgressFileName := f.Name()
	e := gob.NewEncoder(f)

	dms.lock.RLock()
	err = e.Encode(dms.metricGroups)
	dms.lock.RUnlock()
	if err != nil {
		f.Close()
		os.Remove(inProgressFileName)
		return err
	}
	if err := f.Close(); err != nil {
		os.Remove(inProgressFileName)
		return err
	}
	return os.Rename(inProgressFileName, dms.persistenceFile)
}

func (dms *DiskMetricStore) restore() error {
	if dms.persistenceFile == "" {
		return nil
	}
	f, err := os.Open(dms.persistenceFile)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	defer f.Close()
	d := gob.NewDecoder(f)
	return d.Decode(&dms.metricGroups)
}

func (dms *DiskMetricStore) legacyRestore() error {
	if dms.persistenceFile == "" {
		return nil
	}
	f, err := os.Open(dms.persistenceFile)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	defer f.Close()

	var tmf TimestampedMetricFamily
	for d := gob.NewDecoder(f); err == nil; tmf, err = legacyReadTimestampedMetricFamily(d) {
		if len(tmf.MetricFamily.GetMetric()) == 0 {
			continue // No metric in this MetricFamily.
		}
		name := tmf.MetricFamily.GetName()
		var job, instance string
		for _, lp := range tmf.MetricFamily.GetMetric()[0].GetLabel() {
			// With the way the pushgateway persists things, all
			// metrics in a single MetricFamily proto message share
			// the same job and instance label. So we only have to
			// peek at the first metric to find it.
			switch lp.GetName() {
			case "job":
				job = lp.GetValue()
			case "instance":
				instance = lp.GetValue()
			}
			if job != "" && instance != "" {
				break
			}
		}
		labels := map[string]string{
			"job":      job,
			"instance": instance,
		}
		key := model.LabelsToSignature(labels)
		group, ok := dms.metricGroups[key]
		if !ok {
			group = MetricGroup{
				Labels:  labels,
				Metrics: NameToTimestampedMetricFamilyMap{},
			}
			dms.metricGroups[key] = group
		}
		group.Metrics[name] = tmf
	}
	if err == io.EOF {
		return nil
	}
	return err
}

func legacyReadTimestampedMetricFamily(d *gob.Decoder) (TimestampedMetricFamily, error) {
	var buffer []byte
	if err := d.Decode(&buffer); err != nil {
		return TimestampedMetricFamily{}, err
	}
	mf := &dto.MetricFamily{}
	if err := proto.Unmarshal(buffer, mf); err != nil {
		return TimestampedMetricFamily{}, err
	}
	var timestamp time.Time
	if err := d.Decode(&timestamp); err != nil {
		return TimestampedMetricFamily{}, err
	}
	return TimestampedMetricFamily{MetricFamily: mf, Timestamp: timestamp}, nil
}

func copyMetricFamily(mf *dto.MetricFamily) *dto.MetricFamily {
	return &dto.MetricFamily{
		Name:   mf.Name,
		Help:   mf.Help,
		Type:   mf.Type,
		Metric: append([]*dto.Metric{}, mf.Metric...),
	}
}
<commit_msg>Don't schedule persisting if persistence is not configured<commit_after>// Copyright 2014 The Prometheus 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 storage

import (
	"encoding/gob"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/model"

	dto "github.com/prometheus/client_model/go"
)

const (
	writeQueueCapacity = 1000
)

// DiskMetricStore is an implementation of MetricStore that persists metrics to
// disk.
type DiskMetricStore struct {
	lock            sync.RWMutex // Protects metricFamilies.
	writeQueue      chan WriteRequest
	drain           chan struct{}
	done            chan error
	metricGroups    GroupingKeyToMetricGroup
	persistenceFile string
}

type mfStat struct {
	pos    int  // Where in the result slice is the MetricFamily?
	copied bool // Has the MetricFamily already been copied?
}

// NewDiskMetricStore returns a DiskMetricStore ready to use. To cleanly shut it
// down and free resources, the Shutdown() method has to be called.  If
// persistenceFile is the empty string, no persisting to disk will
// happen. Otherwise, a file of that name is used for persisting metrics to
// disk. If the file already exists, metrics are read from it as part of the
// start-up. Persisting is happening upon shutdown and after every write action,
// but the latter will only happen persistenceDuration after the previous
// persisting.
func NewDiskMetricStore(
	persistenceFile string,
	persistenceInterval time.Duration,
) *DiskMetricStore {
	// TODO: Do that outside of the constructor to allow the HTTP server to
	//  serve /-/healthy and /-/ready earlier.
	dms := &DiskMetricStore{
		writeQueue:      make(chan WriteRequest, writeQueueCapacity),
		drain:           make(chan struct{}),
		done:            make(chan error),
		metricGroups:    GroupingKeyToMetricGroup{},
		persistenceFile: persistenceFile,
	}
	if err := dms.restore(); err != nil {
		log.Errorln("Could not load persisted metrics:", err)
		log.Info("Retrying assuming legacy format for persisted metrics...")
		if err := dms.legacyRestore(); err != nil {
			log.Errorln("Could not load persisted metrics in legacy format: ", err)
		}
	}

	go dms.loop(persistenceInterval)
	return dms
}

// SubmitWriteRequest implements the MetricStore interface.
func (dms *DiskMetricStore) SubmitWriteRequest(req WriteRequest) {
	dms.writeQueue <- req
}

// GetMetricFamilies implements the MetricStore interface.
func (dms *DiskMetricStore) GetMetricFamilies() []*dto.MetricFamily {
	result := []*dto.MetricFamily{}
	mfStatByName := map[string]mfStat{}

	dms.lock.RLock()
	defer dms.lock.RUnlock()

	for _, group := range dms.metricGroups {
		for name, tmf := range group.Metrics {
			mf := tmf.MetricFamily
			stat, exists := mfStatByName[name]
			if exists {
				existingMF := result[stat.pos]
				if !stat.copied {
					mfStatByName[name] = mfStat{
						pos:    stat.pos,
						copied: true,
					}
					existingMF = copyMetricFamily(existingMF)
					result[stat.pos] = existingMF
				}
				if mf.GetHelp() != existingMF.GetHelp() || mf.GetType() != existingMF.GetType() {
					log.Infof(
						"Metric families '%s' and '%s' are inconsistent, help and type of the latter will have priority. This is bad. Fix your pushed metrics!",
						mf, existingMF,
					)
				}
				for _, metric := range mf.Metric {
					existingMF.Metric = append(existingMF.Metric, metric)
				}
			} else {
				mfStatByName[name] = mfStat{
					pos:    len(result),
					copied: false,
				}
				result = append(result, mf)
			}
		}
	}
	return result
}

// Shutdown implements the MetricStore interface.
func (dms *DiskMetricStore) Shutdown() error {
	close(dms.drain)
	return <-dms.done
}

// Healthy implements the MetricStore interface.
func (dms *DiskMetricStore) Healthy() error {
	// By taking the lock we check that there is no deadlock.
	dms.lock.Lock()
	defer dms.lock.Unlock()

	// A pushgateway that cannot be written to should not be
	// considered as healthy.
	if len(dms.writeQueue) == cap(dms.writeQueue) {
		return fmt.Errorf("write queue is full")
	}

	return nil
}

// Ready implements the MetricStore interface.
func (dms *DiskMetricStore) Ready() error {
	return dms.Healthy()
}

func (dms *DiskMetricStore) loop(persistenceInterval time.Duration) {
	lastPersist := time.Now()
	persistScheduled := false
	lastWrite := time.Time{}
	persistDone := make(chan time.Time)
	var persistTimer *time.Timer

	checkPersist := func() {
		if dms.persistenceFile != "" && !persistScheduled && lastWrite.After(lastPersist) {
			persistTimer = time.AfterFunc(
				persistenceInterval-lastWrite.Sub(lastPersist),
				func() {
					persistStarted := time.Now()
					if err := dms.persist(); err != nil {
						log.Errorln("Error persisting metrics:", err)
					} else {
						log.Infof(
							"Metrics persisted to '%s'.",
							dms.persistenceFile,
						)
					}
					persistDone <- persistStarted
				},
			)
			persistScheduled = true
		}
	}

	for {
		select {
		case wr := <-dms.writeQueue:
			dms.processWriteRequest(wr)
			lastWrite = time.Now()
			checkPersist()
		case lastPersist = <-persistDone:
			persistScheduled = false
			checkPersist() // In case something has been written in the meantime.
		case <-dms.drain:
			// Prevent a scheduled persist from firing later.
			if persistTimer != nil {
				persistTimer.Stop()
			}
			// Now draining...
			for {
				select {
				case wr := <-dms.writeQueue:
					dms.processWriteRequest(wr)
				default:
					dms.done <- dms.persist()
					return
				}
			}
		}
	}
}

func (dms *DiskMetricStore) processWriteRequest(wr WriteRequest) {
	dms.lock.Lock()
	defer dms.lock.Unlock()

	key := model.LabelsToSignature(wr.Labels)

	if wr.MetricFamilies == nil {
		// Delete.
		delete(dms.metricGroups, key)
		return
	}
	// Update.
	for name, mf := range wr.MetricFamilies {
		group, ok := dms.metricGroups[key]
		if !ok {
			group = MetricGroup{
				Labels:  wr.Labels,
				Metrics: NameToTimestampedMetricFamilyMap{},
			}
			dms.metricGroups[key] = group
		}
		group.Metrics[name] = TimestampedMetricFamily{
			Timestamp:    wr.Timestamp,
			MetricFamily: mf,
		}
	}
}

// GetMetricFamiliesMap implements the MetricStore interface.
func (dms *DiskMetricStore) GetMetricFamiliesMap() GroupingKeyToMetricGroup {
	dms.lock.RLock()
	defer dms.lock.RUnlock()
	groupsCopy := make(GroupingKeyToMetricGroup, len(dms.metricGroups))
	for k, g := range dms.metricGroups {
		metricsCopy := make(NameToTimestampedMetricFamilyMap, len(g.Metrics))
		groupsCopy[k] = MetricGroup{Labels: g.Labels, Metrics: metricsCopy}
		for n, tmf := range g.Metrics {
			metricsCopy[n] = tmf
		}
	}
	return groupsCopy
}

func (dms *DiskMetricStore) persist() error {
	// Check (again) if persistence is configured because some code paths
	// will call this method even if it is not.
	if dms.persistenceFile == "" {
		return nil
	}
	f, err := ioutil.TempFile(
		path.Dir(dms.persistenceFile),
		path.Base(dms.persistenceFile)+".in_progress.",
	)
	if err != nil {
		return err
	}
	inProgressFileName := f.Name()
	e := gob.NewEncoder(f)

	dms.lock.RLock()
	err = e.Encode(dms.metricGroups)
	dms.lock.RUnlock()
	if err != nil {
		f.Close()
		os.Remove(inProgressFileName)
		return err
	}
	if err := f.Close(); err != nil {
		os.Remove(inProgressFileName)
		return err
	}
	return os.Rename(inProgressFileName, dms.persistenceFile)
}

func (dms *DiskMetricStore) restore() error {
	if dms.persistenceFile == "" {
		return nil
	}
	f, err := os.Open(dms.persistenceFile)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	defer f.Close()
	d := gob.NewDecoder(f)
	return d.Decode(&dms.metricGroups)
}

func (dms *DiskMetricStore) legacyRestore() error {
	if dms.persistenceFile == "" {
		return nil
	}
	f, err := os.Open(dms.persistenceFile)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	defer f.Close()

	var tmf TimestampedMetricFamily
	for d := gob.NewDecoder(f); err == nil; tmf, err = legacyReadTimestampedMetricFamily(d) {
		if len(tmf.MetricFamily.GetMetric()) == 0 {
			continue // No metric in this MetricFamily.
		}
		name := tmf.MetricFamily.GetName()
		var job, instance string
		for _, lp := range tmf.MetricFamily.GetMetric()[0].GetLabel() {
			// With the way the pushgateway persists things, all
			// metrics in a single MetricFamily proto message share
			// the same job and instance label. So we only have to
			// peek at the first metric to find it.
			switch lp.GetName() {
			case "job":
				job = lp.GetValue()
			case "instance":
				instance = lp.GetValue()
			}
			if job != "" && instance != "" {
				break
			}
		}
		labels := map[string]string{
			"job":      job,
			"instance": instance,
		}
		key := model.LabelsToSignature(labels)
		group, ok := dms.metricGroups[key]
		if !ok {
			group = MetricGroup{
				Labels:  labels,
				Metrics: NameToTimestampedMetricFamilyMap{},
			}
			dms.metricGroups[key] = group
		}
		group.Metrics[name] = tmf
	}
	if err == io.EOF {
		return nil
	}
	return err
}

func legacyReadTimestampedMetricFamily(d *gob.Decoder) (TimestampedMetricFamily, error) {
	var buffer []byte
	if err := d.Decode(&buffer); err != nil {
		return TimestampedMetricFamily{}, err
	}
	mf := &dto.MetricFamily{}
	if err := proto.Unmarshal(buffer, mf); err != nil {
		return TimestampedMetricFamily{}, err
	}
	var timestamp time.Time
	if err := d.Decode(&timestamp); err != nil {
		return TimestampedMetricFamily{}, err
	}
	return TimestampedMetricFamily{MetricFamily: mf, Timestamp: timestamp}, nil
}

func copyMetricFamily(mf *dto.MetricFamily) *dto.MetricFamily {
	return &dto.MetricFamily{
		Name:   mf.Name,
		Help:   mf.Help,
		Type:   mf.Type,
		Metric: append([]*dto.Metric{}, mf.Metric...),
	}
}
<|endoftext|>


<commit_before>/*
   Copyright (c) 2015, Percona LLC and/or its affiliates. All rights reserved.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package install

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"testing"

	"github.com/percona/percona-agent/pct"
	. "gopkg.in/check.v1"
)

func Test(t *testing.T) { TestingT(t) }

type MainTestSuite struct {
	cmd        *exec.Cmd
	basedir    string
	bin        string
	anotherbin string
	initscript string
	username   string
}

const (
	MOCKED_PERCONA_AGENT = "github.com/percona/percona-agent/install/mock"
)

var _ = Suite(&MainTestSuite{})

func (s *MainTestSuite) SetUpSuite(t *C) {
	var err error

	// We can't/shouldn't use /usr/local/percona/ (the default basedir),
	// so use a tmpdir instead with only a bin dir inside
	s.basedir, err = ioutil.TempDir("/tmp", "percona-agent-init-test-")
	t.Assert(err, IsNil)
	err = os.Mkdir(filepath.Join(s.basedir, pct.BIN_DIR), 0777)
	t.Assert(err, IsNil)

	// Lets compile and place the mocked percona-agent on the tmp basedir
	s.bin = filepath.Join(s.basedir, pct.BIN_DIR, "percona-agent")
	cmd := exec.Command("go", "build", "-o", s.bin, MOCKED_PERCONA_AGENT)
	err = cmd.Run()
	// Failed to compile mocked percona-agent
	t.Assert(err, IsNil, Commentf("Failed to build mocked percona agent: %v", err))

	user, erruser := user.Current()
	t.Assert(erruser, IsNil, Commentf("Failed to obtain current user: %v", err))
	s.username = user.Username
	initscript, err := filepath.Abs("./percona-agent")
	// Check if absolute path resolving succedeed
	t.Assert(err, IsNil)
	// Check if init script is there
	t.Assert(pct.FileExists(initscript), Equals, true)
	s.initscript = filepath.Join(s.basedir, pct.BIN_DIR, "init-script")
	cmd = exec.Command("cp", initscript, s.initscript)
	err = cmd.Run()
	t.Assert(err, IsNil, Commentf("Failed to copy init script to tmp dir: %v", err))
	resetTestEnvVars(s)
}

func (s *MainTestSuite) TearDownTest(t *C) {
	// Delete any left pid file and set mocked agent start delay to 0
	resetTestEnvVars(s)
	// Kill any remaining process before deleting pidfile
	if pid, err := readPidFile(filepath.Join(s.basedir, "percona-agent.pid")); pid != "" && err == nil {
		if numPid, err := strconv.ParseInt(pid, 10, 0); err == nil {
			syscall.Kill(int(numPid), syscall.SIGTERM)
		}
	}
	// Delete if pidFile exists
	os.Remove(filepath.Join(s.basedir, "percona-agent.pid"))
}

func (s *MainTestSuite) TearDownSuite(t *C) {
	// Delete tmp
	if err := os.RemoveAll(s.basedir); err != nil {
		t.Error(err)
	}
}

func resetTestEnvVars(s *MainTestSuite) {
	// Sadly no os.Unsetenv in Go 1.3.x
	os.Setenv("TEST_PERCONA_AGENT_START_DELAY", "")
	os.Setenv("TEST_PERCONA_AGENT_STOP_DELAY", "")
	os.Setenv("PERCONA_AGENT_START_TIMEOUT", "")
	os.Setenv("PERCONA_AGENT_STOP_TIMEOUT", "")
	os.Setenv("PERCONA_AGENT_USER", s.username)
	os.Setenv("PERCONA_AGENT_DIR", s.basedir)
}

func writePidFile(filePath, pid string) error {
	flags := os.O_CREATE | os.O_EXCL | os.O_WRONLY
	file, err := os.OpenFile(filePath, flags, 0644)
	if err != nil {
		//Could not create pidfile
		return err
	}
	// Write PID to file
	if _, err := file.WriteString(pid); err != nil {
		// Could not write to stale pidfile
		return err
	}
	file.Close()
	return nil
}

func readPidFile(pidFilePath string) (pid string, err error) {
	if bytes, err := ioutil.ReadFile(pidFilePath); err != nil {
		return "", err
	} else {
		// Remove any \n
		return strings.Replace(string(bytes), "\n", "", -1), nil
	}
}

//-----------------------------------------------------------------------------

func (s *MainTestSuite) TestStatusNoAgent(t *C) {
	cmd := exec.Command(s.initscript, "status")
	output, err := cmd.Output()
	// status exit code should be 1
	t.Check(err, NotNil)
	// script should output a message
	t.Assert(string(output), Equals, "percona-agent is not running.\n")
}

func (s *MainTestSuite) TestStopNoAgent(t *C) {
	cmd := exec.Command(s.initscript, "stop")
	output, err := cmd.Output()
	// stop exit code should be 0
	t.Check(err, IsNil)
	// script should output a message
	t.Assert(string(output), Equals, "Stopping percona-agent...\npercona-agent is not running.\n")
}

func (s *MainTestSuite) TestStartStop(t *C) {
	// Start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// Start exit code should be 0
	t.Check(err, IsNil)
	// Script should output a message
	t.Assert(string(output), Equals, "Starting percona-agent...\nWaiting for percona-agent to start...\nOK\n")

	// Check status
	cmd = exec.Command(s.initscript, "status")
	output, err = cmd.Output()
	// Status exit code should be 0
	t.Check(err, IsNil)
	// Extract PID from command output
	rePID := regexp.MustCompile(`^percona-agent\ is\ running\ \((\d+)\)\.`)
	found := rePID.FindStringSubmatch(string(output))
	// Check if the command provided a PID
	var pid string
	if len(found) == 2 {
		pid = found[1]
	} else {
		t.Error("Could not get pid for mocked percona-agent")
	}

	pidbinary, err := os.Readlink(fmt.Sprintf("/proc/%v/exe", pid))
	// Check that PID actually points to our mocked percona-agent binary
	t.Assert(pidbinary, Equals, s.bin)

	// Now try to stop
	cmd = exec.Command(s.initscript, "stop")
	output, err = cmd.Output()
	// stop exit code should be 0
	t.Check(err, IsNil)
	// script should output a message
	t.Assert(string(output), Equals, "Stopping percona-agent...\nWaiting for percona-agent to exit...\nStopped percona-agent.\n")
}

func (s *MainTestSuite) TestDoubleStart(t *C) {
	// Start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)
	// Script should output a message
	t.Assert(string(output), Equals, "Starting percona-agent...\nWaiting for percona-agent to start...\nOK\n")

	// Start service again
	cmd = exec.Command(s.initscript, "start")
	output, err = cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)
	// script should output a message
	t.Assert(string(output), Equals, "Starting percona-agent...\npercona-agent is already running.\n")
}

func (s *MainTestSuite) TestWrongBin(t *C) {
	pidFilePath := filepath.Join(s.basedir, "percona-agent.pid")
	// Create pidfile with valid PID but not corresponding to a mocked percona-agent
	if err := writePidFile(pidFilePath, string(os.Getpid())); err != nil {
		t.Errorf("Could not create pidfile: %v", err)
	}

	// Now start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)
	// Script should output a message
	t.Assert(string(output), Equals, fmt.Sprintf("Starting percona-agent...\nRemoved stale pid file: %v\nWaiting for "+
		"percona-agent to start...\nOK\n", pidFilePath))
}

func (s *MainTestSuite) TestStalePIDFile(t *C) {
	// Create pidfile with non valid PID
	pidFilePath := filepath.Join(s.basedir, "percona-agent.pid")
	if err := writePidFile(pidFilePath, string(rand.Uint32())); err != nil {
		t.Errorf("Could not create pidfile: %v", err)
	}

	// Now start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)
	// script should output a message
	t.Assert(string(output), Equals, fmt.Sprintf("Starting percona-agent...\nRemoved stale pid file: %v\nWaiting for "+
		"percona-agent to start...\nOK\n", pidFilePath))
}

func (s *MainTestSuite) TestEnvVariables(t *C) {
	// Set percona-agent user to run with nobody
	os.Setenv("PERCONA_AGENT_USER", "nobody")
	// Now start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 1
	t.Check(err, NotNil)
	// script should output no message TODO: check this
	t.Check(string(output), Equals, "")

	// Set to percona-agent basedir to non existant directory
	os.Setenv("PERCONA_AGENT_DIR", filepath.Join("/tmp", string(rand.Uint32())))
	// Try to start service
	cmd = exec.Command(s.initscript, "start")
	output, err = cmd.Output()
	// start exit code should be 1
	t.Check(err, NotNil)
	// script should output no message TODO: check this
	t.Check(string(output), Equals, "")
}

func (s *MainTestSuite) TestDelayedStart(t *C) {
	// Set init script timeout to 1 second
	os.Setenv("PERCONA_AGENT_START_TIMEOUT", "1")
	// Set percona-agent start delay to 2 seconds
	os.Setenv("TEST_PERCONA_AGENT_START_DELAY", "2")
	// Now try to start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 1
	t.Check(err, NotNil)
	// path to log file, its part of the output
	perconaLogPath := filepath.Join(s.basedir, "percona-agent.log")
	// script should output message
	t.Check(string(output), Equals, fmt.Sprintf("Starting percona-agent...\nWaiting for percona-agent to start...\nFail.  "+
		"Check %v for details.\n", perconaLogPath))
}

func (s *MainTestSuite) TestDelayedStop(t *C) {
	// Set init script stop timeout to 1 second
	os.Setenv("PERCONA_AGENT_STOP_TIMEOUT", "1")
	// Set percona-agent stop delay to 2 seconds
	os.Setenv("TEST_PERCONA_AGENT_STOP_DELAY", "2")
	// Now try to start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)

	// Get the PID from the pidfile
	pid, err := readPidFile(filepath.Join(s.basedir, "percona-agent.pid"))
	// Check if we could read the pidfile
	t.Check(err, IsNil)
	// pid should be non empty
	t.Check(pid, Not(Equals), "")

	stop_cmd := exec.Command(s.initscript, "stop")
	output, err = stop_cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)

	// Script should output message
	t.Check(string(output), Equals, fmt.Sprintf("Stopping percona-agent...\nWaiting for percona-agent to exit...\n"+
		"Time out waiting for percona-agent to exit.  Trying kill -9 %v...\nStopped percona-agent.\n", pid))
	// Make sure the process was killed
	t.Assert(pct.FileExists(fmt.Sprintf("/proc/%v/stat", pid)), Equals, false)
}
<commit_msg>Fixes hardcoded /tmp dir in test<commit_after>/*
   Copyright (c) 2015, Percona LLC and/or its affiliates. All rights reserved.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package install

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"testing"

	"github.com/percona/percona-agent/pct"
	. "gopkg.in/check.v1"
)

func Test(t *testing.T) { TestingT(t) }

type MainTestSuite struct {
	cmd        *exec.Cmd
	basedir    string
	bin        string
	anotherbin string
	initscript string
	username   string
}

const (
	MOCKED_PERCONA_AGENT = "github.com/percona/percona-agent/install/mock"
)

var _ = Suite(&MainTestSuite{})

func (s *MainTestSuite) SetUpSuite(t *C) {
	var err error

	// We can't/shouldn't use /usr/local/percona/ (the default basedir),
	// so use a tmpdir instead with only a bin dir inside
	s.basedir, err = ioutil.TempDir("/tmp", "percona-agent-init-test-")
	t.Assert(err, IsNil)
	err = os.Mkdir(filepath.Join(s.basedir, pct.BIN_DIR), 0777)
	t.Assert(err, IsNil)

	// Lets compile and place the mocked percona-agent on the tmp basedir
	s.bin = filepath.Join(s.basedir, pct.BIN_DIR, "percona-agent")
	cmd := exec.Command("go", "build", "-o", s.bin, MOCKED_PERCONA_AGENT)
	err = cmd.Run()
	// Failed to compile mocked percona-agent
	t.Assert(err, IsNil, Commentf("Failed to build mocked percona agent: %v", err))

	user, erruser := user.Current()
	t.Assert(erruser, IsNil, Commentf("Failed to obtain current user: %v", err))
	s.username = user.Username
	initscript, err := filepath.Abs("./percona-agent")
	// Check if absolute path resolving succedeed
	t.Assert(err, IsNil)
	// Check if init script is there
	t.Assert(pct.FileExists(initscript), Equals, true)
	s.initscript = filepath.Join(s.basedir, pct.BIN_DIR, "init-script")
	cmd = exec.Command("cp", initscript, s.initscript)
	err = cmd.Run()
	t.Assert(err, IsNil, Commentf("Failed to copy init script to tmp dir: %v", err))
	resetTestEnvVars(s)
}

func (s *MainTestSuite) TearDownTest(t *C) {
	// Delete any left pid file and set mocked agent start delay to 0
	resetTestEnvVars(s)
	// Kill any remaining process before deleting pidfile
	if pid, err := readPidFile(filepath.Join(s.basedir, "percona-agent.pid")); pid != "" && err == nil {
		if numPid, err := strconv.ParseInt(pid, 10, 0); err == nil {
			syscall.Kill(int(numPid), syscall.SIGTERM)
		}
	}
	// Delete if pidFile exists
	os.Remove(filepath.Join(s.basedir, "percona-agent.pid"))
}

func (s *MainTestSuite) TearDownSuite(t *C) {
	// Delete tmp
	if err := os.RemoveAll(s.basedir); err != nil {
		t.Error(err)
	}
}

func resetTestEnvVars(s *MainTestSuite) {
	// Sadly no os.Unsetenv in Go 1.3.x
	os.Setenv("TEST_PERCONA_AGENT_START_DELAY", "")
	os.Setenv("TEST_PERCONA_AGENT_STOP_DELAY", "")
	os.Setenv("PERCONA_AGENT_START_TIMEOUT", "")
	os.Setenv("PERCONA_AGENT_STOP_TIMEOUT", "")
	os.Setenv("PERCONA_AGENT_USER", s.username)
	os.Setenv("PERCONA_AGENT_DIR", s.basedir)
}

func writePidFile(filePath, pid string) error {
	flags := os.O_CREATE | os.O_EXCL | os.O_WRONLY
	file, err := os.OpenFile(filePath, flags, 0644)
	if err != nil {
		//Could not create pidfile
		return err
	}
	// Write PID to file
	if _, err := file.WriteString(pid); err != nil {
		// Could not write to stale pidfile
		return err
	}
	file.Close()
	return nil
}

func readPidFile(pidFilePath string) (pid string, err error) {
	if bytes, err := ioutil.ReadFile(pidFilePath); err != nil {
		return "", err
	} else {
		// Remove any \n
		return strings.Replace(string(bytes), "\n", "", -1), nil
	}
}

//-----------------------------------------------------------------------------

func (s *MainTestSuite) TestStatusNoAgent(t *C) {
	cmd := exec.Command(s.initscript, "status")
	output, err := cmd.Output()
	// status exit code should be 1
	t.Check(err, NotNil)
	// script should output a message
	t.Assert(string(output), Equals, "percona-agent is not running.\n")
}

func (s *MainTestSuite) TestStopNoAgent(t *C) {
	cmd := exec.Command(s.initscript, "stop")
	output, err := cmd.Output()
	// stop exit code should be 0
	t.Check(err, IsNil)
	// script should output a message
	t.Assert(string(output), Equals, "Stopping percona-agent...\npercona-agent is not running.\n")
}

func (s *MainTestSuite) TestStartStop(t *C) {
	// Start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// Start exit code should be 0
	t.Check(err, IsNil)
	// Script should output a message
	t.Assert(string(output), Equals, "Starting percona-agent...\nWaiting for percona-agent to start...\nOK\n")

	// Check status
	cmd = exec.Command(s.initscript, "status")
	output, err = cmd.Output()
	// Status exit code should be 0
	t.Check(err, IsNil)
	// Extract PID from command output
	rePID := regexp.MustCompile(`^percona-agent\ is\ running\ \((\d+)\)\.`)
	found := rePID.FindStringSubmatch(string(output))
	// Check if the command provided a PID
	var pid string
	if len(found) == 2 {
		pid = found[1]
	} else {
		t.Error("Could not get pid for mocked percona-agent")
	}

	pidbinary, err := os.Readlink(fmt.Sprintf("/proc/%v/exe", pid))
	// Check that PID actually points to our mocked percona-agent binary
	t.Assert(pidbinary, Equals, s.bin)

	// Now try to stop
	cmd = exec.Command(s.initscript, "stop")
	output, err = cmd.Output()
	// stop exit code should be 0
	t.Check(err, IsNil)
	// script should output a message
	t.Assert(string(output), Equals, "Stopping percona-agent...\nWaiting for percona-agent to exit...\nStopped percona-agent.\n")
}

func (s *MainTestSuite) TestDoubleStart(t *C) {
	// Start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)
	// Script should output a message
	t.Assert(string(output), Equals, "Starting percona-agent...\nWaiting for percona-agent to start...\nOK\n")

	// Start service again
	cmd = exec.Command(s.initscript, "start")
	output, err = cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)
	// script should output a message
	t.Assert(string(output), Equals, "Starting percona-agent...\npercona-agent is already running.\n")
}

func (s *MainTestSuite) TestWrongBin(t *C) {
	pidFilePath := filepath.Join(s.basedir, "percona-agent.pid")
	// Create pidfile with valid PID but not corresponding to a mocked percona-agent
	if err := writePidFile(pidFilePath, string(os.Getpid())); err != nil {
		t.Errorf("Could not create pidfile: %v", err)
	}

	// Now start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)
	// Script should output a message
	t.Assert(string(output), Equals, fmt.Sprintf("Starting percona-agent...\nRemoved stale pid file: %v\nWaiting for "+
		"percona-agent to start...\nOK\n", pidFilePath))
}

func (s *MainTestSuite) TestStalePIDFile(t *C) {
	// Create pidfile with non valid PID
	pidFilePath := filepath.Join(s.basedir, "percona-agent.pid")
	if err := writePidFile(pidFilePath, string(rand.Uint32())); err != nil {
		t.Errorf("Could not create pidfile: %v", err)
	}

	// Now start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)
	// script should output a message
	t.Assert(string(output), Equals, fmt.Sprintf("Starting percona-agent...\nRemoved stale pid file: %v\nWaiting for "+
		"percona-agent to start...\nOK\n", pidFilePath))
}

func (s *MainTestSuite) TestEnvVariables(t *C) {
	// Set percona-agent user to run with nobody
	os.Setenv("PERCONA_AGENT_USER", "nobody")
	// Now start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 1
	t.Check(err, NotNil)
	// script should output no message TODO: check this
	t.Check(string(output), Equals, "")

	// Set to percona-agent basedir to non existant directory
	os.Setenv("PERCONA_AGENT_DIR", filepath.Join(s.basedir, string(rand.Uint32())))
	// Try to start service
	cmd = exec.Command(s.initscript, "start")
	output, err = cmd.Output()
	// start exit code should be 1
	t.Check(err, NotNil)
	// script should output no message TODO: check this
	t.Check(string(output), Equals, "")
}

func (s *MainTestSuite) TestDelayedStart(t *C) {
	// Set init script timeout to 1 second
	os.Setenv("PERCONA_AGENT_START_TIMEOUT", "1")
	// Set percona-agent start delay to 2 seconds
	os.Setenv("TEST_PERCONA_AGENT_START_DELAY", "2")
	// Now try to start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 1
	t.Check(err, NotNil)
	// path to log file, its part of the output
	perconaLogPath := filepath.Join(s.basedir, "percona-agent.log")
	// script should output message
	t.Check(string(output), Equals, fmt.Sprintf("Starting percona-agent...\nWaiting for percona-agent to start...\nFail.  "+
		"Check %v for details.\n", perconaLogPath))
}

func (s *MainTestSuite) TestDelayedStop(t *C) {
	// Set init script stop timeout to 1 second
	os.Setenv("PERCONA_AGENT_STOP_TIMEOUT", "1")
	// Set percona-agent stop delay to 2 seconds
	os.Setenv("TEST_PERCONA_AGENT_STOP_DELAY", "2")
	// Now try to start service
	cmd := exec.Command(s.initscript, "start")
	output, err := cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)

	// Get the PID from the pidfile
	pid, err := readPidFile(filepath.Join(s.basedir, "percona-agent.pid"))
	// Check if we could read the pidfile
	t.Check(err, IsNil)
	// pid should be non empty
	t.Check(pid, Not(Equals), "")

	stop_cmd := exec.Command(s.initscript, "stop")
	output, err = stop_cmd.Output()
	// start exit code should be 0
	t.Check(err, IsNil)

	// Script should output message
	t.Check(string(output), Equals, fmt.Sprintf("Stopping percona-agent...\nWaiting for percona-agent to exit...\n"+
		"Time out waiting for percona-agent to exit.  Trying kill -9 %v...\nStopped percona-agent.\n", pid))
	// Make sure the process was killed
	t.Assert(pct.FileExists(fmt.Sprintf("/proc/%v/stat", pid)), Equals, false)
}
<|endoftext|>


<commit_before>// Package server implements waterfall service
package server

import (
	"context"
	"io"
	"net"
	"os"
	"os/exec"
	"syscall"

	empty_pb "github.com/golang/protobuf/ptypes/empty"
	"github.com/waterfall"
	"github.com/waterfall/forward"
	waterfall_grpc "github.com/waterfall/proto/waterfall_go_grpc"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const bufSize = 32 * 1024

// WaterfallServer implements the waterfall gRPC service
type WaterfallServer struct {
	ctx context.Context
}

// Echo exists solely for test purposes. It's a utility function to
// create integration tests between grpc and custom network transports
// like qemu_pipe and usb
func (s *WaterfallServer) Echo(stream waterfall_grpc.Waterfall_EchoServer) error {
	for {
		msg, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		if err := stream.Send(msg); err != nil {
			return err
		}
	}
}

// Push untars a file/dir that it receives over a gRPC stream
func (s *WaterfallServer) Push(rpc waterfall_grpc.Waterfall_PushServer) error {
	xfer, err := rpc.Recv()
	if err != nil {
		return err
	}

	r, w := io.Pipe()

	// Connect the gRPC stream with a reader that untars
	// the contents of the stream into the desired path
	eg, _ := errgroup.WithContext(s.ctx)
	eg.Go(func() error {
		defer r.Close()
		return waterfall.Untar(r, xfer.Path)
	})

	eg.Go(func() error {
		defer w.Close()
		for {
			if err == io.EOF {
				return rpc.SendAndClose(
					&waterfall_grpc.Transfer{Success: true})
			}
			if err != nil {
				return err
			}
			if _, err := w.Write(xfer.Payload); err != nil {
				return rpc.SendAndClose(
					&waterfall_grpc.Transfer{
						Success: false,
						Err:     []byte(err.Error())})
			}
			xfer, err = rpc.Recv()
		}
	})
	return eg.Wait()
}

// Pull tars up a file/dir and sends it over a gRPC stream
func (s *WaterfallServer) Pull(t *waterfall_grpc.Transfer, ps waterfall_grpc.Waterfall_PullServer) error {
	if _, err := os.Stat(t.Path); os.IsNotExist(err) {
		return err
	}

	r, w := io.Pipe()

	// Connect a writer that traverses a the filesystem
	// tars the contents to the gRPC stream
	eg, _ := errgroup.WithContext(s.ctx)
	eg.Go(func() error {
		defer w.Close()
		return waterfall.Tar(w, t.Path)
	})

	buff := make([]byte, bufSize)
	eg.Go(func() error {
		defer r.Close()

		for {
			n, err := r.Read(buff)
			if err != nil && err != io.EOF {
				return err
			}

			if n > 0 {
				// No need to populate anything else. Sending the path
				// everytime is just overhead.
				xfer := &waterfall_grpc.Transfer{Payload: buff[0:n]}
				if err := ps.Send(xfer); err != nil {
					return err
				}
			}

			if err == io.EOF {
				return nil
			}
		}
	})
	err := eg.Wait()
	return err
}

const (
	sendOut = iota
	sendErr
)

// execWriter wraps the exec grpc server around the Writer interface
type execWriter struct {
	es     waterfall_grpc.Waterfall_ExecServer
	sendTo int32
}

func (ew *execWriter) Write(b []byte) (int, error) {
	cp := &waterfall_grpc.CmdProgress{}
	switch ew.sendTo {
	case sendOut:
		cp.Stdout = b
	case sendErr:
		cp.Stderr = b
	}

	if err := ew.es.Send(cp); err != nil {
		return 0, err
	}
	return len(b), nil
}

// Exec forks a new process with the desired command and pipes its output to the gRPC stream
func (s *WaterfallServer) Exec(es waterfall_grpc.Waterfall_ExecServer) error {

	// The first message contains the actual command to execute.
	// Implmented as a streaming method in order to support stdin redirection in the future.
	cmdMsg, err := es.Recv()
	if err != nil {
		return err
	}

	if cmdMsg.Cmd.PipeIn {
		// stdin redirection is not implemented. Field exists to avoid future incompatibilities.
		return status.Error(codes.Unimplemented, "stdin redirection is not implemented")
	}

	// Avoid doing any sort of input check, e.g. check that the path exists
	// this way we can forward the exact shell output and exit code.
	cmd := exec.CommandContext(s.ctx, cmdMsg.Cmd.Path, cmdMsg.Cmd.Args...)
	cmd.Dir = cmdMsg.Cmd.Dir
	if cmd.Dir == "" {
		cmd.Dir = "/"
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	if err := cmd.Start(); err != nil {
		return err
	}
	// Kill the process in case somehting went wrong with the connection
	defer cmd.Process.Kill()
	defer cmd.Process.Release()

	outWriter := &execWriter{es: es, sendTo: sendOut}
	errWriter := &execWriter{es: es, sendTo: sendErr}

	eg, _ := errgroup.WithContext(s.ctx)
	eg.Go(func() error {
		_, err := io.Copy(outWriter, stdout)
		return err

	})
	eg.Go(func() error {
		_, err := io.Copy(errWriter, stderr)
		return err
	})

	if err := eg.Wait(); err != nil {
		return err
	}

	if err := cmd.Wait(); err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			if stat, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				return es.Send(&waterfall_grpc.CmdProgress{
					ExitCode: uint32(stat.ExitStatus())})
			}
			// the server always runs on android so the type assertion will aways succeed
			panic("this never happens")
		}
		// we got some other kind of error
		return err
	}
	return es.Send(&waterfall_grpc.CmdProgress{ExitCode: 0})
}

// Forward forwards the grpc stream to the requested port.
func (s *WaterfallServer) Forward(stream waterfall_grpc.Waterfall_ForwardServer) error {
	fwd, err := stream.Recv()
	if err != nil {
		return err
	}

	var kind string
	switch fwd.Kind {
	case waterfall_grpc.ForwardMessage_TCP:
		kind = "tcp"
	case waterfall_grpc.ForwardMessage_UDP:
		kind = "udp"
	case waterfall_grpc.ForwardMessage_UNIX:
		kind = "unix"
	default:
		return status.Error(codes.InvalidArgument, "unsupported network type")
	}

	conn, err := net.Dial(kind, fwd.Addr)
	if err != nil {
		return err
	}
	sf := forward.NewStreamForwarder(stream, conn.(forward.HalfReadWriteCloser))
	return sf.Forward()
}

// Version returns the version of the server.
func (s *WaterfallServer) Version(context.Context, *empty_pb.Empty) (*waterfall_grpc.VersionMessage, error) {
	return &waterfall_grpc.VersionMessage{Version: "0.0"}, nil
}

// NewWaterfallServer returns a gRPC server that implements the waterfall service
func NewWaterfallServer(ctx context.Context) *WaterfallServer {
	return &WaterfallServer{ctx: ctx}
}
<commit_msg>Serialize stdout/stderr before Send to stream underlying connection is safe but this is not officially supported in the documentation and grpc migth be doing extra stuff<commit_after>// Package server implements waterfall service
package server

import (
	"context"
	"io"
	"net"
	"os"
	"os/exec"
	"syscall"

	empty_pb "github.com/golang/protobuf/ptypes/empty"
	"github.com/waterfall"
	"github.com/waterfall/forward"
	waterfall_grpc "github.com/waterfall/proto/waterfall_go_grpc"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const bufSize = 32 * 1024

// WaterfallServer implements the waterfall gRPC service
type WaterfallServer struct {
	ctx context.Context
}

// Echo exists solely for test purposes. It's a utility function to
// create integration tests between grpc and custom network transports
// like qemu_pipe and usb
func (s *WaterfallServer) Echo(stream waterfall_grpc.Waterfall_EchoServer) error {
	for {
		msg, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		if err := stream.Send(msg); err != nil {
			return err
		}
	}
}

// Push untars a file/dir that it receives over a gRPC stream
func (s *WaterfallServer) Push(rpc waterfall_grpc.Waterfall_PushServer) error {
	xfer, err := rpc.Recv()
	if err != nil {
		return err
	}

	r, w := io.Pipe()

	// Connect the gRPC stream with a reader that untars
	// the contents of the stream into the desired path
	eg, _ := errgroup.WithContext(s.ctx)
	eg.Go(func() error {
		defer r.Close()
		return waterfall.Untar(r, xfer.Path)
	})

	eg.Go(func() error {
		defer w.Close()
		for {
			if err == io.EOF {
				return rpc.SendAndClose(
					&waterfall_grpc.Transfer{Success: true})
			}
			if err != nil {
				return err
			}
			if _, err := w.Write(xfer.Payload); err != nil {
				return rpc.SendAndClose(
					&waterfall_grpc.Transfer{
						Success: false,
						Err:     []byte(err.Error())})
			}
			xfer, err = rpc.Recv()
		}
	})
	return eg.Wait()
}

// Pull tars up a file/dir and sends it over a gRPC stream
func (s *WaterfallServer) Pull(t *waterfall_grpc.Transfer, ps waterfall_grpc.Waterfall_PullServer) error {
	if _, err := os.Stat(t.Path); os.IsNotExist(err) {
		return err
	}

	r, w := io.Pipe()

	// Connect a writer that traverses a the filesystem
	// tars the contents to the gRPC stream
	eg, _ := errgroup.WithContext(s.ctx)
	eg.Go(func() error {
		defer w.Close()
		return waterfall.Tar(w, t.Path)
	})

	buff := make([]byte, bufSize)
	eg.Go(func() error {
		defer r.Close()

		for {
			n, err := r.Read(buff)
			if err != nil && err != io.EOF {
				return err
			}

			if n > 0 {
				// No need to populate anything else. Sending the path
				// everytime is just overhead.
				xfer := &waterfall_grpc.Transfer{Payload: buff[0:n]}
				if err := ps.Send(xfer); err != nil {
					return err
				}
			}

			if err == io.EOF {
				return nil
			}
		}
	})
	err := eg.Wait()
	return err
}

// chanWriter implements Writer and allows us to use the channel for Copy calls
type chanWriter chan []byte

func (c chanWriter) Write(b []byte) (int, error) {
	// make a copy of the slice, the caller might reuse the slice
	// before we have a chance to flush the bytes to the stream.
	nb := make([]byte, len(b))
	copy(nb, b)
	c <- nb
	return len(nb), nil
}

// Exec forks a new process with the desired command and pipes its output to the gRPC stream
func (s *WaterfallServer) Exec(es waterfall_grpc.Waterfall_ExecServer) error {

	// The first message contains the actual command to execute.
	// Implmented as a streaming method in order to support stdin redirection in the future.
	cmdMsg, err := es.Recv()
	if err != nil {
		return err
	}

	if cmdMsg.Cmd.PipeIn {
		// stdin redirection is not implemented. Field exists to avoid future incompatibilities.
		return status.Error(codes.Unimplemented, "stdin redirection is not implemented")
	}

	// Avoid doing any sort of input check, e.g. check that the path exists
	// this way we can forward the exact shell output and exit code.
	cmd := exec.CommandContext(s.ctx, cmdMsg.Cmd.Path, cmdMsg.Cmd.Args...)
	cmd.Dir = cmdMsg.Cmd.Dir
	if cmd.Dir == "" {
		cmd.Dir = "/"
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	if err := cmd.Start(); err != nil {
		return err
	}
	// Kill the process in case something went wrong with the connection
	defer cmd.Process.Kill()
	defer cmd.Process.Release()

	// We want to read concurrently from stdout/stderr
	// but we only have one stream to send it to so we need to merge streams before writing.
	// Note that order is not necessarily preserverd.
	stdoutCh := make(chan []byte, 1)
	stderrCh := make(chan []byte, 1)

	eg, _ := errgroup.WithContext(s.ctx)

	// Serialize and multiplex stdout/stderr channels to the grpc stream
	eg.Go(func() error {
		_, err := io.Copy(chanWriter(stdoutCh), stdout)
		close(stdoutCh)
		return err

	})
	eg.Go(func() error {
		_, err := io.Copy(chanWriter(stderrCh), stderr)
		close(stderrCh)
		return err
	})
	eg.Go(func() error {
		var o []byte
		var e []byte
		oo := true
		eo := true
		for oo || eo {
			var msg *waterfall_grpc.CmdProgress
			select {
			case o, oo = <-stdoutCh:
				if !oo {
					break
				}
				msg = &waterfall_grpc.CmdProgress{Stdout: o}
			case e, eo = <-stderrCh:
				if !eo {
					break
				}
				msg = &waterfall_grpc.CmdProgress{Stderr: e}
			}
			if msg != nil {
				if err := es.Send(msg); err != nil {
					return err
				}
			}
		}
		return nil
	})

	if err := eg.Wait(); err != nil {
		return err
	}

	if err := cmd.Wait(); err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			if stat, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				return es.Send(&waterfall_grpc.CmdProgress{
					ExitCode: uint32(stat.ExitStatus())})
			}
			// the server always runs on android so the type assertion will aways succeed
			panic("this never happens")
		}
		// we got some other kind of error
		return err
	}
	return es.Send(&waterfall_grpc.CmdProgress{ExitCode: 0})
}

// Forward forwards the grpc stream to the requested port.
func (s *WaterfallServer) Forward(stream waterfall_grpc.Waterfall_ForwardServer) error {
	fwd, err := stream.Recv()
	if err != nil {
		return err
	}

	var kind string
	switch fwd.Kind {
	case waterfall_grpc.ForwardMessage_TCP:
		kind = "tcp"
	case waterfall_grpc.ForwardMessage_UDP:
		kind = "udp"
	case waterfall_grpc.ForwardMessage_UNIX:
		kind = "unix"
	default:
		return status.Error(codes.InvalidArgument, "unsupported network type")
	}

	conn, err := net.Dial(kind, fwd.Addr)
	if err != nil {
		return err
	}
	sf := forward.NewStreamForwarder(stream, conn.(forward.HalfReadWriteCloser))
	return sf.Forward()
}

// Version returns the version of the server.
func (s *WaterfallServer) Version(context.Context, *empty_pb.Empty) (*waterfall_grpc.VersionMessage, error) {
	return &waterfall_grpc.VersionMessage{Version: "0.0"}, nil
}

// NewWaterfallServer returns a gRPC server that implements the waterfall service
func NewWaterfallServer(ctx context.Context) *WaterfallServer {
	return &WaterfallServer{ctx: ctx}
}
<|endoftext|>


<commit_before>// Copyright 2017 The LUCI 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 buildbucket

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/luci/buildbucket"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	bbv1 "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"
	"google.golang.org/genproto/protobuf/field_mask"
)

var (
	buildCounter = metric.NewCounter(
		"luci/milo/buildbucket_pubsub/builds",
		"The number of buildbucket builds received by Milo from PubSub",
		nil,
		field.String("bucket"),
		// True for luci build, False for non-luci (ie buildbot) build.
		field.Bool("luci"),
		// Status can be "COMPLETED", "SCHEDULED", or "STARTED"
		field.String("status"),
		// Action can be one of 3 options.
		//   * "Created" - This is the first time Milo heard about this build
		//   * "Modified" - Milo updated some information about this build vs. what
		//     it knew before.
		//   * "Rejected" - Milo was unable to accept this build.
		field.String("action"))
)

// PubSubHandler is a webhook that stores the builds coming in from pubsub.
func PubSubHandler(ctx *router.Context) {
	err := pubSubHandlerImpl(ctx.Context, ctx.Request)
	if err != nil {
		logging.Errorf(ctx.Context, "error while handling pubsub event")
		errors.Log(ctx.Context, err)
	}
	if transient.Tag.In(err) {
		// Transient errors are 500 so that PubSub retries them.
		ctx.Writer.WriteHeader(http.StatusInternalServerError)
		return
	}
	// No errors or non-transient errors are 200s so that PubSub does not retry
	// them.
	ctx.Writer.WriteHeader(http.StatusOK)
}

func mustTimestamp(ts *timestamp.Timestamp) time.Time {
	if t, err := ptypes.Timestamp(ts); err == nil {
		return t
	}
	return time.Time{}
}

var summaryBuildMask = &field_mask.FieldMask{
	Paths: []string{
		"id",
		"builder",
		"number",
		"create_time",
		"start_time",
		"end_time",
		"update_time",
		"status",
		"summary_markdown",
		"tags",
		"infra.swarming",
	},
}

// getSummary returns a model.BuildSummary representing a buildbucket build.
func getSummary(c context.Context, host string, project string, id int64) (*model.BuildSummary, error) {
	client, err := buildbucketClient(c, host, auth.AsProject, auth.WithProject(project))
	if err != nil {
		return nil, err
	}
	b, err := client.GetBuild(c, &buildbucketpb.GetBuildRequest{
		Id:     id,
		Fields: summaryBuildMask,
	})
	if err != nil {
		return nil, err
	}
	buildAddress := fmt.Sprintf("%d", b.Id)
	if b.Number != 0 {
		buildAddress = fmt.Sprintf("luci.%s.%s/%s/%d", b.Builder.Project, b.Builder.Bucket, b.Builder.Builder, b.Number)
	}

	// Note: The parent for buildbucket build summaries is currently a fake entity.
	// In the future, builds can be cached here, but we currently don't do that.
	buildKey := datastore.MakeKey(c, "buildbucket.Build", fmt.Sprintf("%s:%s", host, buildAddress))
	swarming := b.GetInfra().GetSwarming()

	return &model.BuildSummary{
		ProjectID: b.Builder.Project,
		BuildKey:  buildKey,
		BuilderID: BuilderID{*b.Builder}.String(),
		BuildID:   "buildbucket/" + buildAddress,
		BuildSet:  b.Buildsets(),
		ContextURI: []string{
			fmt.Sprintf("buildbucket://%s/build/%d", host, id),
			fmt.Sprintf("swarming://%s/task/%s", swarming.GetHostname(), swarming.GetTaskId()),
		},
		Created: mustTimestamp(b.CreateTime),
		Summary: model.Summary{
			Start:  mustTimestamp(b.StartTime),
			End:    mustTimestamp(b.EndTime),
			Status: statusMap[b.Status],
		},
		Version: mustTimestamp(b.UpdateTime).UnixNano(),
	}, nil
}

// generateSummary takes a decoded buildbucket event and generates
// a model.BuildSummary from it.
//
// This is the portion of the summarization process which cannot fail (i.e. is
// pure-data).
func generateSummary(c context.Context, hostname string, build buildbucket.Build) (*model.BuildSummary, error) {
	bs, err := getSummary(c, hostname, build.Project, build.ID)
	if err != nil {
		logging.WithError(err).Errorf(c, "got error while getting summary")
	}

	buildset := build.Tags[bbv1.TagBuildSet]
	if buildset == nil {
		buildset = []string{}
	}
	ret := &model.BuildSummary{
		ProjectID: build.Project,
		BuildKey:  MakeBuildKey(c, hostname, build.Address()),
		BuilderID: NewBuilderID(build.Bucket, build.Builder).String(),
		BuildID:   "buildbucket/" + build.Address(),
		BuildSet:  buildset,
		ContextURI: []string{
			fmt.Sprintf("buildbucket://%s/build/%d", hostname, build.ID),
		},

		Created: build.CreationTime,
		Summary: model.Summary{
			Start:  build.StartTime,
			End:    build.CompletionTime,
			Status: parseStatus(build.Status),
		},

		Version: build.UpdateTime.UnixNano(),

		Experimental: build.Experimental,
	}

	if shost, sid := build.Tags.Get("swarming_hostname"), build.Tags.Get("swarming_task_id"); shost != "" && sid != "" {
		ret.ContextURI = append(ret.ContextURI, fmt.Sprintf("swarming://%s/task/%s", shost, sid))
	}

	// Informational compare of v1 and v2 API result.
	if diff := cmp.Diff(ret, bs, cmpopts.IgnoreUnexported(model.BuildSummary{})); diff != "" {
		logging.Errorf(c, "BuildSummary of v2 has Diff (-v1, +v2)\n%s", diff)
	} else {
		logging.Debugf(c, "BuildSummary between v1 and v2 are the same.")
	}
	// TODO(iannucci,nodir): get the bot context too

	// TODO(iannucci,nodir): support manifests/got_revision
	return ret, ret.AddManifestKeysFromBuildSets(c)
}

// pubSubHandlerImpl takes the http.Request, expects to find
// a common.PubSubSubscription JSON object in the Body, containing a bbPSEvent,
// and handles the contents with generateSummary.
func pubSubHandlerImpl(c context.Context, r *http.Request) error {
	// This is the default action. The code below will modify the values of some
	// or all of these parameters.
	isLUCI, bucket, status, action := false, "UNKNOWN", "UNKNOWN", "Rejected"

	defer func() {
		// closure for late binding
		buildCounter.Add(c, 1, bucket, isLUCI, status, action)
	}()

	msg := common.PubSubSubscription{}
	if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
		// This might be a transient error, e.g. when the json format changes
		// and Milo isn't updated yet.
		return errors.Annotate(err, "could not decode message").Tag(transient.Tag).Err()
	}
	if v, ok := msg.Message.Attributes["version"].(string); ok && v != "v1" {
		// TODO(nodir): switch to v2, crbug.com/826006
		logging.Debugf(c, "unsupported pubsub message version %q. Ignoring", v)
		return nil
	}
	bData, err := msg.GetData()
	if err != nil {
		return errors.Annotate(err, "could not parse pubsub message string").Err()
	}

	event := struct {
		Build    bbv1.ApiCommonBuildMessage `json:"build"`
		Hostname string                     `json:"hostname"`
	}{}
	if err := json.Unmarshal(bData, &event); err != nil {
		return errors.Annotate(err, "could not parse pubsub message data").Err()
	}

	build := buildbucket.Build{}
	if err := build.ParseMessage(&event.Build); err != nil {
		return errors.Annotate(err, "could not parse buildbucket.Build").Err()
	}

	bucket = build.Bucket
	status = build.Status.String()
	isLUCI = strings.HasPrefix(bucket, "luci.")

	logging.Debugf(c, "Received from %s: build %s/%s (%s)\n%v",
		event.Hostname, bucket, build.Builder, status, build)

	if !isLUCI || build.Builder == "" {
		logging.Infof(c, "This is not an ingestable build, ignoring")
		return nil
	}

	bs, err := generateSummary(c, event.Hostname, build)
	if err != nil {
		return err
	}

	return transient.Tag.Apply(datastore.RunInTransaction(c, func(c context.Context) error {
		curBS := &model.BuildSummary{BuildKey: bs.BuildKey}
		switch err := datastore.Get(c, curBS); err {
		case datastore.ErrNoSuchEntity:
			action = "Created"
		case nil:
			action = "Modified"
		default:
			return errors.Annotate(err, "reading current BuildSummary").Err()
		}

		if build.UpdateTime.UnixNano() <= curBS.Version {
			logging.Warningf(c, "current BuildSummary is newer: %d <= %d",
				build.UpdateTime.UnixNano(), curBS.Version)
			return nil
		}

		if err := datastore.Put(c, bs); err != nil {
			return err
		}

		return model.UpdateBuilderForBuild(c, bs)
	}, &datastore.TransactionOptions{XG: true}))
}

// MakeBuildKey returns a new datastore Key for a buildbucket.Build.
//
// There's currently no model associated with this key, but it's used as
// a parent for a model.BuildSummary.
func MakeBuildKey(c context.Context, host, buildAddress string) *datastore.Key {
	return datastore.MakeKey(c,
		"buildbucket.Build", fmt.Sprintf("%s:%s", host, buildAddress))
}
<commit_msg>[milo] Add experimental bit for pubsub pipeline v2<commit_after>// Copyright 2017 The LUCI 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 buildbucket

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"go.chromium.org/gae/service/datastore"
	"go.chromium.org/luci/buildbucket"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	bbv1 "go.chromium.org/luci/common/api/buildbucket/buildbucket/v1"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	"go.chromium.org/luci/milo/common"
	"go.chromium.org/luci/milo/common/model"
	"go.chromium.org/luci/server/auth"
	"go.chromium.org/luci/server/router"
	"google.golang.org/genproto/protobuf/field_mask"
)

var (
	buildCounter = metric.NewCounter(
		"luci/milo/buildbucket_pubsub/builds",
		"The number of buildbucket builds received by Milo from PubSub",
		nil,
		field.String("bucket"),
		// True for luci build, False for non-luci (ie buildbot) build.
		field.Bool("luci"),
		// Status can be "COMPLETED", "SCHEDULED", or "STARTED"
		field.String("status"),
		// Action can be one of 3 options.
		//   * "Created" - This is the first time Milo heard about this build
		//   * "Modified" - Milo updated some information about this build vs. what
		//     it knew before.
		//   * "Rejected" - Milo was unable to accept this build.
		field.String("action"))
)

// PubSubHandler is a webhook that stores the builds coming in from pubsub.
func PubSubHandler(ctx *router.Context) {
	err := pubSubHandlerImpl(ctx.Context, ctx.Request)
	if err != nil {
		logging.Errorf(ctx.Context, "error while handling pubsub event")
		errors.Log(ctx.Context, err)
	}
	if transient.Tag.In(err) {
		// Transient errors are 500 so that PubSub retries them.
		ctx.Writer.WriteHeader(http.StatusInternalServerError)
		return
	}
	// No errors or non-transient errors are 200s so that PubSub does not retry
	// them.
	ctx.Writer.WriteHeader(http.StatusOK)
}

func mustTimestamp(ts *timestamp.Timestamp) time.Time {
	if t, err := ptypes.Timestamp(ts); err == nil {
		return t
	}
	return time.Time{}
}

var summaryBuildMask = &field_mask.FieldMask{
	Paths: []string{
		"id",
		"builder",
		"number",
		"create_time",
		"start_time",
		"end_time",
		"update_time",
		"status",
		"summary_markdown",
		"tags",
		"infra.swarming",
	},
}

// getSummary returns a model.BuildSummary representing a buildbucket build.
func getSummary(c context.Context, host string, project string, id int64) (*model.BuildSummary, error) {
	client, err := buildbucketClient(c, host, auth.AsProject, auth.WithProject(project))
	if err != nil {
		return nil, err
	}
	b, err := client.GetBuild(c, &buildbucketpb.GetBuildRequest{
		Id:     id,
		Fields: summaryBuildMask,
	})
	if err != nil {
		return nil, err
	}
	buildAddress := fmt.Sprintf("%d", b.Id)
	if b.Number != 0 {
		buildAddress = fmt.Sprintf("luci.%s.%s/%s/%d", b.Builder.Project, b.Builder.Bucket, b.Builder.Builder, b.Number)
	}

	// Note: The parent for buildbucket build summaries is currently a fake entity.
	// In the future, builds can be cached here, but we currently don't do that.
	buildKey := datastore.MakeKey(c, "buildbucket.Build", fmt.Sprintf("%s:%s", host, buildAddress))
	swarming := b.GetInfra().GetSwarming()

	return &model.BuildSummary{
		ProjectID: b.Builder.Project,
		BuildKey:  buildKey,
		BuilderID: BuilderID{*b.Builder}.String(),
		BuildID:   "buildbucket/" + buildAddress,
		BuildSet:  b.Buildsets(),
		ContextURI: []string{
			fmt.Sprintf("buildbucket://%s/build/%d", host, id),
			fmt.Sprintf("swarming://%s/task/%s", swarming.GetHostname(), swarming.GetTaskId()),
		},
		Created: mustTimestamp(b.CreateTime),
		Summary: model.Summary{
			Start:  mustTimestamp(b.StartTime),
			End:    mustTimestamp(b.EndTime),
			Status: statusMap[b.Status],
		},
		Version:      mustTimestamp(b.UpdateTime).UnixNano(),
		Experimental: b.GetInput().GetExperimental(),
	}, nil
}

// generateSummary takes a decoded buildbucket event and generates
// a model.BuildSummary from it.
//
// This is the portion of the summarization process which cannot fail (i.e. is
// pure-data).
func generateSummary(c context.Context, hostname string, build buildbucket.Build) (*model.BuildSummary, error) {
	bs, err := getSummary(c, hostname, build.Project, build.ID)
	if err != nil {
		logging.WithError(err).Errorf(c, "got error while getting summary")
	}

	buildset := build.Tags[bbv1.TagBuildSet]
	if buildset == nil {
		buildset = []string{}
	}
	ret := &model.BuildSummary{
		ProjectID: build.Project,
		BuildKey:  MakeBuildKey(c, hostname, build.Address()),
		BuilderID: NewBuilderID(build.Bucket, build.Builder).String(),
		BuildID:   "buildbucket/" + build.Address(),
		BuildSet:  buildset,
		ContextURI: []string{
			fmt.Sprintf("buildbucket://%s/build/%d", hostname, build.ID),
		},

		Created: build.CreationTime,
		Summary: model.Summary{
			Start:  build.StartTime,
			End:    build.CompletionTime,
			Status: parseStatus(build.Status),
		},

		Version: build.UpdateTime.UnixNano(),

		Experimental: build.Experimental,
	}

	if shost, sid := build.Tags.Get("swarming_hostname"), build.Tags.Get("swarming_task_id"); shost != "" && sid != "" {
		ret.ContextURI = append(ret.ContextURI, fmt.Sprintf("swarming://%s/task/%s", shost, sid))
	}

	// Informational compare of v1 and v2 API result.
	if diff := cmp.Diff(ret, bs, cmpopts.IgnoreUnexported(model.BuildSummary{})); diff != "" {
		logging.Errorf(c, "BuildSummary of v2 has Diff (-v1, +v2)\n%s", diff)
	} else {
		logging.Debugf(c, "BuildSummary between v1 and v2 are the same.")
	}
	// TODO(iannucci,nodir): get the bot context too

	// TODO(iannucci,nodir): support manifests/got_revision
	return ret, ret.AddManifestKeysFromBuildSets(c)
}

// pubSubHandlerImpl takes the http.Request, expects to find
// a common.PubSubSubscription JSON object in the Body, containing a bbPSEvent,
// and handles the contents with generateSummary.
func pubSubHandlerImpl(c context.Context, r *http.Request) error {
	// This is the default action. The code below will modify the values of some
	// or all of these parameters.
	isLUCI, bucket, status, action := false, "UNKNOWN", "UNKNOWN", "Rejected"

	defer func() {
		// closure for late binding
		buildCounter.Add(c, 1, bucket, isLUCI, status, action)
	}()

	msg := common.PubSubSubscription{}
	if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
		// This might be a transient error, e.g. when the json format changes
		// and Milo isn't updated yet.
		return errors.Annotate(err, "could not decode message").Tag(transient.Tag).Err()
	}
	if v, ok := msg.Message.Attributes["version"].(string); ok && v != "v1" {
		// TODO(nodir): switch to v2, crbug.com/826006
		logging.Debugf(c, "unsupported pubsub message version %q. Ignoring", v)
		return nil
	}
	bData, err := msg.GetData()
	if err != nil {
		return errors.Annotate(err, "could not parse pubsub message string").Err()
	}

	event := struct {
		Build    bbv1.ApiCommonBuildMessage `json:"build"`
		Hostname string                     `json:"hostname"`
	}{}
	if err := json.Unmarshal(bData, &event); err != nil {
		return errors.Annotate(err, "could not parse pubsub message data").Err()
	}

	build := buildbucket.Build{}
	if err := build.ParseMessage(&event.Build); err != nil {
		return errors.Annotate(err, "could not parse buildbucket.Build").Err()
	}

	bucket = build.Bucket
	status = build.Status.String()
	isLUCI = strings.HasPrefix(bucket, "luci.")

	logging.Debugf(c, "Received from %s: build %s/%s (%s)\n%v",
		event.Hostname, bucket, build.Builder, status, build)

	if !isLUCI || build.Builder == "" {
		logging.Infof(c, "This is not an ingestable build, ignoring")
		return nil
	}

	bs, err := generateSummary(c, event.Hostname, build)
	if err != nil {
		return err
	}

	return transient.Tag.Apply(datastore.RunInTransaction(c, func(c context.Context) error {
		curBS := &model.BuildSummary{BuildKey: bs.BuildKey}
		switch err := datastore.Get(c, curBS); err {
		case datastore.ErrNoSuchEntity:
			action = "Created"
		case nil:
			action = "Modified"
		default:
			return errors.Annotate(err, "reading current BuildSummary").Err()
		}

		if build.UpdateTime.UnixNano() <= curBS.Version {
			logging.Warningf(c, "current BuildSummary is newer: %d <= %d",
				build.UpdateTime.UnixNano(), curBS.Version)
			return nil
		}

		if err := datastore.Put(c, bs); err != nil {
			return err
		}

		return model.UpdateBuilderForBuild(c, bs)
	}, &datastore.TransactionOptions{XG: true}))
}

// MakeBuildKey returns a new datastore Key for a buildbucket.Build.
//
// There's currently no model associated with this key, but it's used as
// a parent for a model.BuildSummary.
func MakeBuildKey(c context.Context, host, buildAddress string) *datastore.Key {
	return datastore.MakeKey(c,
		"buildbucket.Build", fmt.Sprintf("%s:%s", host, buildAddress))
}
<|endoftext|>


<commit_before>/*
 * Copyright 2020 Arm Limited.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 bpwriter

import (
	"fmt"
	"sort"
	"strings"
	"sync"

	"github.com/ARM-software/bob-build/internal/utils"
)

// Implement types and helpers to record all the modules that we want
// to write into an Android.bp file.
//
// This is a really basic implementation that allows us to add key
// value pairs to a module, and store string, bools and string lists.
// Only a single level of nesting for properties is supported.

func indentString(depth int) string {
	return strings.Repeat(" ", depth*4)
}

var bpEscaper = strings.NewReplacer("\"", "\\\"")

func Escape(s string) string {
	return bpEscaper.Replace(s)
}

func EscapeList(slice []string) []string {
	slice = append([]string(nil), slice...)
	for i, s := range slice {
		slice[i] = Escape(s)
	}
	return slice
}

// Key-value pair for a property
type property struct {
	// Property key
	key string
	// Property value represented as a string
	value string
}

// Adjacent properties within an Android.bp file
type Group interface {
	AddString(name, value string)
	AddBool(name string, value bool)
	AddStringList(name string, list []string)
	AddStringCmd(name string, argLists ...[]string)
}

type group struct {
	// The name of this group. "" for the top level group.
	// This is necessary as we want to store group as an array
	// rather than a map, to maintain ordering.
	name string
	// Depth
	depth int
	// Properties in the group.
	// This is a list rather than a map to maintain ordering.
	props []property
}

var _ Group = (*group)(nil)

func (g *group) addProp(key, value string) {
	p := property{
		key:   key,
		value: value,
	}
	g.props = append(g.props, p)
}

// Add a string property to the property group
func (g *group) AddString(name, value string) {
	g.addProp(name, "\""+Escape(value)+"\"")
}

// Add a boolean property to the property group
func (g *group) AddBool(name string, value bool) {
	if value {
		g.addProp(name, "true")
	} else {
		g.addProp(name, "false")
	}
}

// Add a string list property to the property group
func (g *group) AddStringList(name string, list []string) {
	if len(list) == 0 {
		return
	}

	s := ""

	if len(list) > 1 {
		// Put each entry on a new line, indented
		s += "[\n"
		indent := indentString(g.depth + 1)
		for _, v := range list {
			s += indent + "\"" + Escape(v) + "\",\n"
		}
		// The list close is back-indented one tab
		s += indentString(g.depth) + "]"
	} else {
		// One entry. Put on a single line.
		s += "[\"" + Escape(list[0]) + "\"]"
	}
	g.addProp(name, s)
}

// Add a string property containing a command, or list of shell-style
// arguments, generated by joining all the array elements provided. Assume that
// quoting and escaping is handled at source.
func (g *group) AddStringCmd(name string, argLists ...[]string) {
	g.AddString(name, utils.Join(argLists...))
}

// Render the property group into a string
func (g *group) render(b *strings.Builder) {
	indent := indentString(g.depth)
	for _, p := range g.props {
		b.WriteString(indent + p.key + ": " + p.value + ",\n")
	}
}

// Arbitrary Android.bp module
//
// No locking for modules, as the creation of each module is done
// in a single thread.
type Module interface {
	Group
	NewGroup(name string) Group
}

type module struct {
	// Module type
	modType string

	// Module name
	name string

	// Top level properties of module
	group

	// Nested properties (1 level deep only)
	groups []group
}

var _ Module = (*module)(nil)

// Add a string property as a top level module property
func (m *module) AddString(name, value string) {
	m.group.AddString(name, value)
}

// Add a boolean property as a top level module property
func (m *module) AddBool(name string, value bool) {
	m.group.AddBool(name, value)
}

// Add a string list property as a top level module property
func (m *module) AddStringList(name string, list []string) {
	m.group.AddStringList(name, list)
}

// Add a string property containing a command, or list of shell-style
// arguments, generated by joining all the array elements provided. Assume that
// quoting and escaping is handled at source.
func (m *module) AddStringCmd(name string, argLists ...[]string) {
	m.group.AddStringCmd(name, argLists...)
}

// Create a property group in the module
func (m *module) NewGroup(name string) Group {
	g := group{}
	g.name = name
	g.depth = 2
	m.groups = append(m.groups, g)
	return &g
}

// Render the module into a string
func (m *module) render(b *strings.Builder) {
	indent := indentString(1)
	b.WriteString(m.modType + " {\n" + indent + "name: \"" + m.name + "\",\n")
	m.group.render(b)
	for _, group := range m.groups {
		b.WriteString(indent + group.name + ": {\n")
		group.render(b)
		b.WriteString(indent + "},\n")
	}
	b.WriteString("}\n\n")
}

func moduleFactory(modType, name string) *module {
	m := module{}
	m.modType = modType
	m.name = name
	m.group.name = ""
	m.group.depth = 1
	return &m
}

// Content for an Android.bp file
type File interface {
	NewModule(modType, name string) (Module, error)
	Render(b *strings.Builder)
}

type file struct {
	sync.Mutex
	modules map[string]*module
}

var _ File = (*file)(nil)

// Create a module
func (f *file) NewModule(modType, name string) (Module, error) {
	m := moduleFactory(modType, name)

	// Lock the addition to ensure parallel build actions can add
	// modules to the file.
	f.Lock()
	defer f.Unlock()

	if _, dup := f.modules[name]; !dup {
		f.modules[name] = m
	} else {
		err := fmt.Errorf("Duplicate module name (%s)", name)
		return nil, err
	}

	return m, nil
}

// Render all the modules in the file
func (f *file) Render(b *strings.Builder) {
	modNames := make([]string, len(f.modules))
	i := 0
	for name, _ := range f.modules {
		modNames[i] = name
		i = i + 1
	}
	sort.Strings(modNames)
	for _, name := range modNames {
		f.modules[name].render(b)
	}
}

func FileFactory() File {
	f := file{}
	f.modules = make(map[string]*module)
	return &f
}
<commit_msg>Fix bpwriter groups not being rendered<commit_after>/*
 * Copyright 2020 Arm Limited.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 bpwriter

import (
	"fmt"
	"sort"
	"strings"
	"sync"

	"github.com/ARM-software/bob-build/internal/utils"
)

// Implement types and helpers to record all the modules that we want
// to write into an Android.bp file.
//
// This is a really basic implementation that allows us to add key
// value pairs to a module, and store string, bools and string lists.
// Only a single level of nesting for properties is supported.

func indentString(depth int) string {
	return strings.Repeat(" ", depth*4)
}

var bpEscaper = strings.NewReplacer("\"", "\\\"")

func Escape(s string) string {
	return bpEscaper.Replace(s)
}

func EscapeList(slice []string) []string {
	slice = append([]string(nil), slice...)
	for i, s := range slice {
		slice[i] = Escape(s)
	}
	return slice
}

// Key-value pair for a property
type property struct {
	// Property key
	key string
	// Property value represented as a string
	value string
}

// Adjacent properties within an Android.bp file
type Group interface {
	AddString(name, value string)
	AddBool(name string, value bool)
	AddStringList(name string, list []string)
	AddStringCmd(name string, argLists ...[]string)
}

type group struct {
	// The name of this group. "" for the top level group.
	// This is necessary as we want to store group as an array
	// rather than a map, to maintain ordering.
	name string
	// Depth
	depth int
	// Properties in the group.
	// This is a list rather than a map to maintain ordering.
	props []property
}

var _ Group = (*group)(nil)

func (g *group) addProp(key, value string) {
	p := property{
		key:   key,
		value: value,
	}
	g.props = append(g.props, p)
}

// Add a string property to the property group
func (g *group) AddString(name, value string) {
	g.addProp(name, "\""+Escape(value)+"\"")
}

// Add a boolean property to the property group
func (g *group) AddBool(name string, value bool) {
	if value {
		g.addProp(name, "true")
	} else {
		g.addProp(name, "false")
	}
}

// Add a string list property to the property group
func (g *group) AddStringList(name string, list []string) {
	if len(list) == 0 {
		return
	}

	s := ""

	if len(list) > 1 {
		// Put each entry on a new line, indented
		s += "[\n"
		indent := indentString(g.depth + 1)
		for _, v := range list {
			s += indent + "\"" + Escape(v) + "\",\n"
		}
		// The list close is back-indented one tab
		s += indentString(g.depth) + "]"
	} else {
		// One entry. Put on a single line.
		s += "[\"" + Escape(list[0]) + "\"]"
	}
	g.addProp(name, s)
}

// Add a string property containing a command, or list of shell-style
// arguments, generated by joining all the array elements provided. Assume that
// quoting and escaping is handled at source.
func (g *group) AddStringCmd(name string, argLists ...[]string) {
	g.AddString(name, utils.Join(argLists...))
}

// Render the property group into a string
func (g *group) render(b *strings.Builder) {
	indent := indentString(g.depth)
	for _, p := range g.props {
		b.WriteString(indent + p.key + ": " + p.value + ",\n")
	}
}

// Arbitrary Android.bp module
//
// No locking for modules, as the creation of each module is done
// in a single thread.
type Module interface {
	Group
	NewGroup(name string) Group
}

type module struct {
	// Module type
	modType string

	// Module name
	name string

	// Top level properties of module
	group

	// Nested properties (1 level deep only)
	groups []*group
}

var _ Module = (*module)(nil)

// Add a string property as a top level module property
func (m *module) AddString(name, value string) {
	m.group.AddString(name, value)
}

// Add a boolean property as a top level module property
func (m *module) AddBool(name string, value bool) {
	m.group.AddBool(name, value)
}

// Add a string list property as a top level module property
func (m *module) AddStringList(name string, list []string) {
	m.group.AddStringList(name, list)
}

// Add a string property containing a command, or list of shell-style
// arguments, generated by joining all the array elements provided. Assume that
// quoting and escaping is handled at source.
func (m *module) AddStringCmd(name string, argLists ...[]string) {
	m.group.AddStringCmd(name, argLists...)
}

// Create a property group in the module
func (m *module) NewGroup(name string) Group {
	g := group{}
	g.name = name
	g.depth = 2
	m.groups = append(m.groups, &g)
	return &g
}

// Render the module into a string
func (m *module) render(b *strings.Builder) {
	indent := indentString(1)
	b.WriteString(m.modType + " {\n" + indent + "name: \"" + m.name + "\",\n")
	m.group.render(b)
	for _, group := range m.groups {
		b.WriteString(indent + group.name + ": {\n")
		group.render(b)
		b.WriteString(indent + "},\n")
	}
	b.WriteString("}\n\n")
}

func moduleFactory(modType, name string) *module {
	m := module{}
	m.modType = modType
	m.name = name
	m.group.name = ""
	m.group.depth = 1
	return &m
}

// Content for an Android.bp file
type File interface {
	NewModule(modType, name string) (Module, error)
	Render(b *strings.Builder)
}

type file struct {
	sync.Mutex
	modules map[string]*module
}

var _ File = (*file)(nil)

// Create a module
func (f *file) NewModule(modType, name string) (Module, error) {
	m := moduleFactory(modType, name)

	// Lock the addition to ensure parallel build actions can add
	// modules to the file.
	f.Lock()
	defer f.Unlock()

	if _, dup := f.modules[name]; !dup {
		f.modules[name] = m
	} else {
		err := fmt.Errorf("Duplicate module name (%s)", name)
		return nil, err
	}

	return m, nil
}

// Render all the modules in the file
func (f *file) Render(b *strings.Builder) {
	modNames := make([]string, len(f.modules))
	i := 0
	for name, _ := range f.modules {
		modNames[i] = name
		i = i + 1
	}
	sort.Strings(modNames)
	for _, name := range modNames {
		f.modules[name].render(b)
	}
}

func FileFactory() File {
	f := file{}
	f.modules = make(map[string]*module)
	return &f
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2019 the Octant contributors. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package cluster

import (
	"context"
	"github.com/pkg/errors"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/rest"
)

//go:generate mockgen -source=namespace.go -destination=./fake/mock_namespace_interface.go -package=fake github.com/vmware-tanzu/octant/internal/cluster NamespaceInterface

// NamespaceInterface is an interface for querying namespace details.
type NamespaceInterface interface {
	Names() ([]string, error)
	InitialNamespace() string
	ProvidedNamespaces() []string
	HasNamespace(namespace string) bool
}

type namespaceClient struct {
	restClient         rest.Interface
	dynamicClient      dynamic.Interface
	initialNamespace   string
	providedNamespaces []string
}

var _ NamespaceInterface = (*namespaceClient)(nil)

func newNamespaceClient(dynamicClient dynamic.Interface, restClient rest.Interface, initialNamespace string, providedNamespaces []string) *namespaceClient {
	return &namespaceClient{
		restClient:         restClient,
		dynamicClient:      dynamicClient,
		initialNamespace:   initialNamespace,
		providedNamespaces: providedNamespaces,
	}
}

func (n *namespaceClient) Names() ([]string, error) {
	namespaces, err := namespaces(n.dynamicClient)
	if err != nil {
		return nil, err
	}

	var names []string
	for _, namespace := range namespaces {
		names = append(names, namespace.GetName())
	}

	return names, nil
}

func (n *namespaceClient) HasNamespace(namespace string) bool {
	ns := &corev1.Namespace{}
	err := n.restClient.Get().Resource("namespaces").Name(namespace).Do().Into(ns)
	if err != nil {
		return false
	}
	return true
}

// Namespaces returns available namespaces.
func namespaces(dc dynamic.Interface) ([]corev1.Namespace, error) {
	res := schema.GroupVersionResource{
		Version:  "v1",
		Resource: "namespaces",
	}

	nri := dc.Resource(res)

	list, err := nri.List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, errors.Wrap(err, "list namespaces")
	}

	var nsList corev1.NamespaceList
	err = runtime.DefaultUnstructuredConverter.FromUnstructured(list.UnstructuredContent(), &nsList)
	if err != nil {
		return nil, errors.Wrap(err, "convert object to namespace list")
	}

	return nsList.Items, nil
}

func (n *namespaceClient) InitialNamespace() string {
	return n.initialNamespace
}

func (n *namespaceClient) ProvidedNamespaces() []string {
	return n.providedNamespaces
}
<commit_msg>Fix missing context from namespace get<commit_after>/*
Copyright (c) 2019 the Octant contributors. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package cluster

import (
	"context"

	"github.com/pkg/errors"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/rest"
)

//go:generate mockgen -source=namespace.go -destination=./fake/mock_namespace_interface.go -package=fake github.com/vmware-tanzu/octant/internal/cluster NamespaceInterface

// NamespaceInterface is an interface for querying namespace details.
type NamespaceInterface interface {
	Names() ([]string, error)
	InitialNamespace() string
	ProvidedNamespaces() []string
	HasNamespace(namespace string) bool
}

type namespaceClient struct {
	restClient         rest.Interface
	dynamicClient      dynamic.Interface
	initialNamespace   string
	providedNamespaces []string
}

var _ NamespaceInterface = (*namespaceClient)(nil)

func newNamespaceClient(dynamicClient dynamic.Interface, restClient rest.Interface, initialNamespace string, providedNamespaces []string) *namespaceClient {
	return &namespaceClient{
		restClient:         restClient,
		dynamicClient:      dynamicClient,
		initialNamespace:   initialNamespace,
		providedNamespaces: providedNamespaces,
	}
}

func (n *namespaceClient) Names() ([]string, error) {
	namespaces, err := namespaces(n.dynamicClient)
	if err != nil {
		return nil, err
	}

	var names []string
	for _, namespace := range namespaces {
		names = append(names, namespace.GetName())
	}

	return names, nil
}

func (n *namespaceClient) HasNamespace(namespace string) bool {
	ns := &corev1.Namespace{}
	err := n.restClient.Get().Resource("namespaces").Name(namespace).Do(context.TODO()).Into(ns)
	if err != nil {
		return false
	}
	return true
}

// Namespaces returns available namespaces.
func namespaces(dc dynamic.Interface) ([]corev1.Namespace, error) {
	res := schema.GroupVersionResource{
		Version:  "v1",
		Resource: "namespaces",
	}

	nri := dc.Resource(res)

	list, err := nri.List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		return nil, errors.Wrap(err, "list namespaces")
	}

	var nsList corev1.NamespaceList
	err = runtime.DefaultUnstructuredConverter.FromUnstructured(list.UnstructuredContent(), &nsList)
	if err != nil {
		return nil, errors.Wrap(err, "convert object to namespace list")
	}

	return nsList.Items, nil
}

func (n *namespaceClient) InitialNamespace() string {
	return n.initialNamespace
}

func (n *namespaceClient) ProvidedNamespaces() []string {
	return n.providedNamespaces
}
<|endoftext|>