<commit_before>// Copyright 2013 Prometheus Team
// 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 metric

import (
	"github.com/prometheus/prometheus/model"
	"github.com/prometheus/prometheus/utility"
	"sort"
	"sync"
	"time"
)

const (
	// Assuming sample rate of 1 / 15Hz, this allows for one hour's worth of
	// storage per metric without any major reallocations.
	initialSeriesArenaSize = 4 * 60
)

// Models a given sample entry stored in the in-memory arena.
type value interface {
	// Gets the given value.
	get() model.SampleValue
}

// Models a single sample value.  It presumes that there is either no subsequent
// value seen or that any subsequent values are of a different value.
type singletonValue model.SampleValue

func (v singletonValue) get() model.SampleValue {
	return model.SampleValue(v)
}

type stream struct {
	sync.RWMutex

	metric model.Metric
	values model.Values
}

func (s *stream) add(timestamp time.Time, value model.SampleValue) {
	s.Lock()
	defer s.Unlock()

	// BUG(all): https://github.com/prometheus/prometheus/pull/265/files#r4336435.

	s.values = append(s.values, model.SamplePair{
		Timestamp: timestamp,
		Value:     value,
	})
}

func (s *stream) clone() model.Values {
	s.RLock()
	defer s.RUnlock()

	// BUG(all): Examine COW technique.

	clone := make(model.Values, len(s.values))
	copy(clone, s.values)

	return clone
}

func (s *stream) getValueAtTime(t time.Time) model.Values {
	s.RLock()
	defer s.RUnlock()

	// BUG(all): May be avenues for simplification.
	l := len(s.values)
	switch l {
	case 0:
		return model.Values{}
	case 1:
		return model.Values{s.values[0]}
	default:
		index := sort.Search(l, func(i int) bool {
			return !s.values[i].Timestamp.Before(t)
		})

		if index == 0 {
			return model.Values{s.values[0]}
		}
		if index == l {
			return model.Values{s.values[l-1]}
		}

		if s.values[index].Timestamp.Equal(t) {
			return model.Values{s.values[index]}
		}
		return model.Values{s.values[index-1], s.values[index]}
	}
}

func (s *stream) getBoundaryValues(i model.Interval) (model.Values, model.Values) {
	return s.getValueAtTime(i.OldestInclusive), s.getValueAtTime(i.NewestInclusive)
}

func (s *stream) getRangeValues(in model.Interval) model.Values {
	s.RLock()
	defer s.RUnlock()

	oldest := sort.Search(len(s.values), func(i int) bool {
		return !s.values[i].Timestamp.Before(in.OldestInclusive)
	})

	newest := sort.Search(len(s.values), func(i int) bool {
		return s.values[i].Timestamp.After(in.NewestInclusive)
	})

	result := make(model.Values, newest-oldest)
	copy(result, s.values[oldest:newest])

	return result
}

func newStream(metric model.Metric) *stream {
	return &stream{
		metric: metric,
		values: make(model.Values, 0, initialSeriesArenaSize),
	}
}

type memorySeriesStorage struct {
	sync.RWMutex

	fingerprintToSeries     map[model.Fingerprint]*stream
	labelPairToFingerprints map[model.LabelPair]model.Fingerprints
	labelNameToFingerprints map[model.LabelName]model.Fingerprints
}

func (s *memorySeriesStorage) AppendSamples(samples model.Samples) error {
	for _, sample := range samples {
		s.AppendSample(sample)
	}

	return nil
}

func (s *memorySeriesStorage) AppendSample(sample model.Sample) error {
	metric := sample.Metric
	fingerprint := model.NewFingerprintFromMetric(metric)
	s.RLock()
	series, ok := s.fingerprintToSeries[*fingerprint]
	s.RUnlock()

	if !ok {
		series = newStream(metric)
		s.Lock()
		s.fingerprintToSeries[*fingerprint] = series

		for k, v := range metric {
			labelPair := model.LabelPair{
				Name:  k,
				Value: v,
			}
			labelPairValues := s.labelPairToFingerprints[labelPair]
			labelPairValues = append(labelPairValues, fingerprint)
			s.labelPairToFingerprints[labelPair] = labelPairValues

			labelNameValues := s.labelNameToFingerprints[k]
			labelNameValues = append(labelNameValues, fingerprint)
			s.labelNameToFingerprints[k] = labelNameValues
		}

		s.Unlock()
	}

	series.add(sample.Timestamp, sample.Value)

	return nil
}

// Append raw sample, bypassing indexing. Only used to add data to views, which
// don't need to lookup by metric.
func (s *memorySeriesStorage) appendSampleWithoutIndexing(f *model.Fingerprint, timestamp time.Time, value model.SampleValue) {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()

	if !ok {
		series = newStream(model.Metric{})
		s.Lock()
		s.fingerprintToSeries[*f] = series
		s.Unlock()
	}

	series.add(timestamp, value)
}

func (s *memorySeriesStorage) GetFingerprintsForLabelSet(l model.LabelSet) (fingerprints model.Fingerprints, err error) {
	sets := []utility.Set{}

	s.RLock()
	for k, v := range l {
		values := s.labelPairToFingerprints[model.LabelPair{
			Name:  k,
			Value: v,
		}]
		set := utility.Set{}
		for _, fingerprint := range values {
			set.Add(*fingerprint)
		}
		sets = append(sets, set)
	}
	s.RUnlock()

	setCount := len(sets)
	if setCount == 0 {
		return fingerprints, nil
	}

	base := sets[0]
	for i := 1; i < setCount; i++ {
		base = base.Intersection(sets[i])
	}
	for _, e := range base.Elements() {
		fingerprint := e.(model.Fingerprint)
		fingerprints = append(fingerprints, &fingerprint)
	}

	return fingerprints, nil
}

func (s *memorySeriesStorage) GetFingerprintsForLabelName(l model.LabelName) (model.Fingerprints, error) {
	s.RLock()
	values, ok := s.labelNameToFingerprints[l]
	s.RUnlock()
	if !ok {
		return nil, nil
	}

	fingerprints := make(model.Fingerprints, len(values))
	copy(fingerprints, values)

	return fingerprints, nil
}

func (s *memorySeriesStorage) GetMetricForFingerprint(f *model.Fingerprint) (model.Metric, error) {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()
	if !ok {
		return nil, nil
	}

	metric := model.Metric{}
	for label, value := range series.metric {
		metric[label] = value
	}

	return metric, nil
}

func (s *memorySeriesStorage) CloneSamples(f *model.Fingerprint) model.Values {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()
	if !ok {
		return nil
	}

	return series.clone()
}

func (s *memorySeriesStorage) GetValueAtTime(f *model.Fingerprint, t time.Time) model.Values {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()
	if !ok {
		return nil
	}

	return series.getValueAtTime(t)
}

func (s *memorySeriesStorage) GetBoundaryValues(f *model.Fingerprint, i model.Interval) (model.Values, model.Values) {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()
	if !ok {
		return nil, nil
	}

	return series.getBoundaryValues(i)
}

func (s *memorySeriesStorage) GetRangeValues(f *model.Fingerprint, i model.Interval) model.Values {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()

	if !ok {
		return nil
	}

	return series.getRangeValues(i)
}

func (s *memorySeriesStorage) Close() {
	s.fingerprintToSeries = map[model.Fingerprint]*stream{}
	s.labelPairToFingerprints = map[model.LabelPair]model.Fingerprints{}
	s.labelNameToFingerprints = map[model.LabelName]model.Fingerprints{}
}

func (s *memorySeriesStorage) GetAllValuesForLabel(labelName model.LabelName) (values model.LabelValues, err error) {
	valueSet := map[model.LabelValue]bool{}
	for _, series := range s.fingerprintToSeries {
		if value, ok := series.metric[labelName]; ok {
			if !valueSet[value] {
				values = append(values, value)
				valueSet[value] = true
			}
		}
	}
	return
}

func NewMemorySeriesStorage() *memorySeriesStorage {
	return &memorySeriesStorage{
		fingerprintToSeries:     make(map[model.Fingerprint]*stream),
		labelPairToFingerprints: make(map[model.LabelPair]model.Fingerprints),
		labelNameToFingerprints: make(map[model.LabelName]model.Fingerprints),
	}
}
<commit_msg>Code Review: Extend lock.<commit_after>// Copyright 2013 Prometheus Team
// 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 metric

import (
	"github.com/prometheus/prometheus/model"
	"github.com/prometheus/prometheus/utility"
	"sort"
	"sync"
	"time"
)

const (
	// Assuming sample rate of 1 / 15Hz, this allows for one hour's worth of
	// storage per metric without any major reallocations.
	initialSeriesArenaSize = 4 * 60
)

// Models a given sample entry stored in the in-memory arena.
type value interface {
	// Gets the given value.
	get() model.SampleValue
}

// Models a single sample value.  It presumes that there is either no subsequent
// value seen or that any subsequent values are of a different value.
type singletonValue model.SampleValue

func (v singletonValue) get() model.SampleValue {
	return model.SampleValue(v)
}

type stream struct {
	sync.RWMutex

	metric model.Metric
	values model.Values
}

func (s *stream) add(timestamp time.Time, value model.SampleValue) {
	s.Lock()
	defer s.Unlock()

	// BUG(all): https://github.com/prometheus/prometheus/pull/265/files#r4336435.

	s.values = append(s.values, model.SamplePair{
		Timestamp: timestamp,
		Value:     value,
	})
}

func (s *stream) clone() model.Values {
	s.RLock()
	defer s.RUnlock()

	// BUG(all): Examine COW technique.

	clone := make(model.Values, len(s.values))
	copy(clone, s.values)

	return clone
}

func (s *stream) getValueAtTime(t time.Time) model.Values {
	s.RLock()
	defer s.RUnlock()

	// BUG(all): May be avenues for simplification.
	l := len(s.values)
	switch l {
	case 0:
		return model.Values{}
	case 1:
		return model.Values{s.values[0]}
	default:
		index := sort.Search(l, func(i int) bool {
			return !s.values[i].Timestamp.Before(t)
		})

		if index == 0 {
			return model.Values{s.values[0]}
		}
		if index == l {
			return model.Values{s.values[l-1]}
		}

		if s.values[index].Timestamp.Equal(t) {
			return model.Values{s.values[index]}
		}
		return model.Values{s.values[index-1], s.values[index]}
	}
}

func (s *stream) getBoundaryValues(i model.Interval) (model.Values, model.Values) {
	return s.getValueAtTime(i.OldestInclusive), s.getValueAtTime(i.NewestInclusive)
}

func (s *stream) getRangeValues(in model.Interval) model.Values {
	s.RLock()
	defer s.RUnlock()

	oldest := sort.Search(len(s.values), func(i int) bool {
		return !s.values[i].Timestamp.Before(in.OldestInclusive)
	})

	newest := sort.Search(len(s.values), func(i int) bool {
		return s.values[i].Timestamp.After(in.NewestInclusive)
	})

	result := make(model.Values, newest-oldest)
	copy(result, s.values[oldest:newest])

	return result
}

func newStream(metric model.Metric) *stream {
	return &stream{
		metric: metric,
		values: make(model.Values, 0, initialSeriesArenaSize),
	}
}

type memorySeriesStorage struct {
	sync.RWMutex

	fingerprintToSeries     map[model.Fingerprint]*stream
	labelPairToFingerprints map[model.LabelPair]model.Fingerprints
	labelNameToFingerprints map[model.LabelName]model.Fingerprints
}

func (s *memorySeriesStorage) AppendSamples(samples model.Samples) error {
	for _, sample := range samples {
		s.AppendSample(sample)
	}

	return nil
}

func (s *memorySeriesStorage) AppendSample(sample model.Sample) error {
	metric := sample.Metric
	fingerprint := model.NewFingerprintFromMetric(metric)
	s.RLock()
	series, ok := s.fingerprintToSeries[*fingerprint]
	s.RUnlock()

	if !ok {
		series = newStream(metric)
		s.Lock()
		s.fingerprintToSeries[*fingerprint] = series

		for k, v := range metric {
			labelPair := model.LabelPair{
				Name:  k,
				Value: v,
			}
			labelPairValues := s.labelPairToFingerprints[labelPair]
			labelPairValues = append(labelPairValues, fingerprint)
			s.labelPairToFingerprints[labelPair] = labelPairValues

			labelNameValues := s.labelNameToFingerprints[k]
			labelNameValues = append(labelNameValues, fingerprint)
			s.labelNameToFingerprints[k] = labelNameValues
		}

		s.Unlock()
	}

	series.add(sample.Timestamp, sample.Value)

	return nil
}

// Append raw sample, bypassing indexing. Only used to add data to views, which
// don't need to lookup by metric.
func (s *memorySeriesStorage) appendSampleWithoutIndexing(f *model.Fingerprint, timestamp time.Time, value model.SampleValue) {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()

	if !ok {
		series = newStream(model.Metric{})
		s.Lock()
		s.fingerprintToSeries[*f] = series
		s.Unlock()
	}

	series.add(timestamp, value)
}

func (s *memorySeriesStorage) GetFingerprintsForLabelSet(l model.LabelSet) (fingerprints model.Fingerprints, err error) {
	sets := []utility.Set{}

	s.RLock()
	for k, v := range l {
		values := s.labelPairToFingerprints[model.LabelPair{
			Name:  k,
			Value: v,
		}]
		set := utility.Set{}
		for _, fingerprint := range values {
			set.Add(*fingerprint)
		}
		sets = append(sets, set)
	}
	s.RUnlock()

	setCount := len(sets)
	if setCount == 0 {
		return fingerprints, nil
	}

	base := sets[0]
	for i := 1; i < setCount; i++ {
		base = base.Intersection(sets[i])
	}
	for _, e := range base.Elements() {
		fingerprint := e.(model.Fingerprint)
		fingerprints = append(fingerprints, &fingerprint)
	}

	return fingerprints, nil
}

func (s *memorySeriesStorage) GetFingerprintsForLabelName(l model.LabelName) (model.Fingerprints, error) {
	s.RLock()
	defer s.RUnlock()
	values, ok := s.labelNameToFingerprints[l]
	if !ok {
		return nil, nil
	}

	fingerprints := make(model.Fingerprints, len(values))
	copy(fingerprints, values)

	return fingerprints, nil
}

func (s *memorySeriesStorage) GetMetricForFingerprint(f *model.Fingerprint) (model.Metric, error) {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()
	if !ok {
		return nil, nil
	}

	metric := model.Metric{}
	for label, value := range series.metric {
		metric[label] = value
	}

	return metric, nil
}

func (s *memorySeriesStorage) CloneSamples(f *model.Fingerprint) model.Values {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()
	if !ok {
		return nil
	}

	return series.clone()
}

func (s *memorySeriesStorage) GetValueAtTime(f *model.Fingerprint, t time.Time) model.Values {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()
	if !ok {
		return nil
	}

	return series.getValueAtTime(t)
}

func (s *memorySeriesStorage) GetBoundaryValues(f *model.Fingerprint, i model.Interval) (model.Values, model.Values) {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()
	if !ok {
		return nil, nil
	}

	return series.getBoundaryValues(i)
}

func (s *memorySeriesStorage) GetRangeValues(f *model.Fingerprint, i model.Interval) model.Values {
	s.RLock()
	series, ok := s.fingerprintToSeries[*f]
	s.RUnlock()

	if !ok {
		return nil
	}

	return series.getRangeValues(i)
}

func (s *memorySeriesStorage) Close() {
	s.fingerprintToSeries = map[model.Fingerprint]*stream{}
	s.labelPairToFingerprints = map[model.LabelPair]model.Fingerprints{}
	s.labelNameToFingerprints = map[model.LabelName]model.Fingerprints{}
}

func (s *memorySeriesStorage) GetAllValuesForLabel(labelName model.LabelName) (values model.LabelValues, err error) {
	valueSet := map[model.LabelValue]bool{}
	for _, series := range s.fingerprintToSeries {
		if value, ok := series.metric[labelName]; ok {
			if !valueSet[value] {
				values = append(values, value)
				valueSet[value] = true
			}
		}
	}
	return
}

func NewMemorySeriesStorage() *memorySeriesStorage {
	return &memorySeriesStorage{
		fingerprintToSeries:     make(map[model.Fingerprint]*stream),
		labelPairToFingerprints: make(map[model.LabelPair]model.Fingerprints),
		labelNameToFingerprints: make(map[model.LabelName]model.Fingerprints),
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package common

import (
	"github.com/juju/errors"
	"github.com/juju/utils/series"

	apimetadata "github.com/juju/juju/api/imagemetadata"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/imagemetadata"
	"github.com/juju/juju/environs/simplestreams"
	"github.com/juju/juju/juju"
)

// FindImageMetadata looks for image metadata in state.
// If none are found, we fall back on original image search in simple streams.
func FindImageMetadata(env environs.Environ, imageConstraint *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	stateMetadata, stateInfo, err := imageMetadataFromState(env, imageConstraint, signedOnly)
	if err != nil && !errors.IsNotFound(err) {
		// look into simple stream if for some reason can't get from state server
		logger.Infof("could not get image metadata from state server: %v", err)
	}

	// No need to look in data sources if found in state?
	if len(stateMetadata) != 0 {
		return stateMetadata, stateInfo, nil
	}

	// If none are found, fall back to original simple stream impl.
	// Currently, an image metadata worker picks up this metadata periodically (daily),
	// and stores it in state. So potentially, this collection could be different
	// to what is in state.
	dsMetadata, dsInfo, err := imageMetadataFromDataSources(env, imageConstraint, signedOnly)
	if err != nil {
		if !errors.IsNotFound(err) {
			return nil, nil, errors.Trace(err)
		}
	}

	// If still none found, complain
	if len(dsMetadata) == 0 {
		return nil, nil, errors.NotFoundf("image metadata for series %v, arch %v", imageConstraint.Series, imageConstraint.Arches)
	}

	return dsMetadata, dsInfo, nil
}

var metadataAPI = func(env environs.Environ) (*apimetadata.Client, error) {
	api, err := juju.NewAPIFromName(env.Config().Name())
	if err != nil {
		return nil, errors.Annotate(err, "could not connect to api")
	}
	return apimetadata.NewClient(api), nil
}

func imageMetadataFromState(env environs.Environ, ic *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	metadataAPI, err := metadataAPI(env)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	stored, err := metadataAPI.List(ic.Stream, ic.Region, ic.Series, ic.Arches, "", "")
	if err != nil {
		return nil, nil, errors.Annotate(err, "could not list image metadata from state server")
	}

	// Convert to common format.
	images := make([]*imagemetadata.ImageMetadata, len(stored))
	for i, one := range stored {
		m := &imagemetadata.ImageMetadata{
			Storage:    one.RootStorageType,
			Id:         one.ImageId,
			VirtType:   one.VirtType,
			Arch:       one.Arch,
			RegionName: one.Region,
			Stream:     one.Stream,
		}
		m.Version, _ = series.SeriesVersion(one.Series)
		images[i] = m
	}

	info := &simplestreams.ResolveInfo{}
	info.Source = "state server"
	// This is currently ignored for image metadata that is stored in state
	// since when stored, both signed and unsigned metadata are written,
	// but whether it was signed or not is not.
	info.Signed = signedOnly
	return images, info, nil
}

// imageMetadataFromDataSources finds image metadata using
// existing data sources.
func imageMetadataFromDataSources(env environs.Environ, constraint *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	sources, err := environs.ImageMetadataSources(env)
	if err != nil {
		return nil, nil, err
	}
	return imagemetadata.Fetch(sources, constraint, signedOnly)
}
<commit_msg>My blutanly obvious log statements.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package common

import (
	"github.com/juju/errors"
	"github.com/juju/utils/series"

	apimetadata "github.com/juju/juju/api/imagemetadata"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/imagemetadata"
	"github.com/juju/juju/environs/simplestreams"
	"github.com/juju/juju/juju"
)

// FindImageMetadata looks for image metadata in state.
// If none are found, we fall back on original image search in simple streams.
func FindImageMetadata(env environs.Environ, imageConstraint *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	stateMetadata, stateInfo, err := imageMetadataFromState(env, imageConstraint, signedOnly)
	if err != nil && !errors.IsNotFound(err) {
		// look into simple stream if for some reason can't get from state server
		logger.Infof("could not get image metadata from state server: %v", err)
	}

	logger.Debugf("\n\n GOT FROM STATE %d METADATA \n\n", len(stateMetadata))
	// No need to look in data sources if found in state?
	if len(stateMetadata) != 0 {
		return stateMetadata, stateInfo, nil
	}

	// If none are found, fall back to original simple stream impl.
	// Currently, an image metadata worker picks up this metadata periodically (daily),
	// and stores it in state. So potentially, this collection could be different
	// to what is in state.
	dsMetadata, dsInfo, err := imageMetadataFromDataSources(env, imageConstraint, signedOnly)
	if err != nil {
		if !errors.IsNotFound(err) {
			return nil, nil, errors.Trace(err)
		}
	}

	logger.Debugf("\n\n GOT FROM DATA SOURCES %d METADATA \n\n", len(dsMetadata))
	// If still none found, complain
	if len(dsMetadata) == 0 {
		return nil, nil, errors.NotFoundf("image metadata for series %v, arch %v", imageConstraint.Series, imageConstraint.Arches)
	}

	return dsMetadata, dsInfo, nil
}

var metadataAPI = func(env environs.Environ) (*apimetadata.Client, error) {
	api, err := juju.NewAPIFromName(env.Config().Name())
	if err != nil {
		return nil, errors.Annotate(err, "could not connect to api")
	}
	return apimetadata.NewClient(api), nil
}

func imageMetadataFromState(env environs.Environ, ic *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	metadataAPI, err := metadataAPI(env)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	stored, err := metadataAPI.List(ic.Stream, ic.Region, ic.Series, ic.Arches, "", "")
	if err != nil {
		return nil, nil, errors.Annotate(err, "could not list image metadata from state server")
	}

	// Convert to common format.
	images := make([]*imagemetadata.ImageMetadata, len(stored))
	for i, one := range stored {
		m := &imagemetadata.ImageMetadata{
			Storage:    one.RootStorageType,
			Id:         one.ImageId,
			VirtType:   one.VirtType,
			Arch:       one.Arch,
			RegionName: one.Region,
			Stream:     one.Stream,
		}
		m.Version, _ = series.SeriesVersion(one.Series)
		images[i] = m
	}

	info := &simplestreams.ResolveInfo{}
	info.Source = "state server"
	// This is currently ignored for image metadata that is stored in state
	// since when stored, both signed and unsigned metadata are written,
	// but whether it was signed or not is not.
	info.Signed = signedOnly
	return images, info, nil
}

// imageMetadataFromDataSources finds image metadata using
// existing data sources.
func imageMetadataFromDataSources(env environs.Environ, constraint *imagemetadata.ImageConstraint, signedOnly bool) ([]*imagemetadata.ImageMetadata, *simplestreams.ResolveInfo, error) {
	sources, err := environs.ImageMetadataSources(env)
	if err != nil {
		return nil, nil, err
	}
	return imagemetadata.Fetch(sources, constraint, signedOnly)
}
<|endoftext|>


<commit_before>// This package implements a provisioner for Packer that executes
// shell scripts within the remote machine.
package shell

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"io/ioutil"
	"log"
	"os"
	"strings"
)

const DefaultRemotePath = "/tmp/script.sh"

type config struct {
	// An inline script to execute. Multiple strings are all executed
	// in the context of a single shell.
	Inline []string

	// The shebang value used when running inline scripts.
	InlineShebang string `mapstructure:"inline_shebang"`

	// The local path of the shell script to upload and execute.
	Script string

	// An array of multiple scripts to run.
	Scripts []string

	// An array of environment variables that will be injected before
	// your command(s) are executed.
	Vars []string `mapstructure:"environment_vars"`

	// The remote path where the local shell script will be uploaded to.
	// This should be set to a writable file that is in a pre-existing directory.
	RemotePath string `mapstructure:"remote_path"`

	// The command used to execute the script. The '{{ .Path }}' variable
	// should be used to specify where the script goes, {{ .Vars }}
	// can be used to inject the environment_vars into the environment.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Packer configurations, these come from Packer itself
	PackerBuildName   string `mapstructure:"packer_build_name"`
	PackerBuilderType string `mapstructure:"packer_builder_type"`

	tpl *common.Template
}

type Provisioner struct {
	config config
}

type ExecuteCommandTemplate struct {
	Vars string
	Path string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	md, err := common.DecodeConfig(&p.config, raws...)
	if err != nil {
		return err
	}

	p.config.tpl, err = common.NewTemplate()
	if err != nil {
		return err
	}

	// Accumulate any errors
	errs := common.CheckUnusedConfig(md)

	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "chmod +x {{.Path}}; {{.Vars}} {{.Path}}"
	}

	if p.config.Inline != nil && len(p.config.Inline) == 0 {
		p.config.Inline = nil
	}

	if p.config.InlineShebang == "" {
		p.config.InlineShebang = "/bin/sh"
	}

	if p.config.RemotePath == "" {
		p.config.RemotePath = DefaultRemotePath
	}

	if p.config.Scripts == nil {
		p.config.Scripts = make([]string, 0)
	}

	if p.config.Vars == nil {
		p.config.Vars = make([]string, 0)
	}

	if p.config.Script != "" && len(p.config.Scripts) > 0 {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only one of script or scripts can be specified."))
	}

	if p.config.Script != "" {
		p.config.Scripts = []string{p.config.Script}
	}

	templates := map[string]*string{
		"inline_shebang": &p.config.InlineShebang,
		"script":         &p.config.Script,
		"remote_path":    &p.config.RemotePath,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = p.config.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	sliceTemplates := map[string][]string{
		"inline":           p.config.Inline,
		"scripts":          p.config.Scripts,
		"environment_vars": p.config.Vars,
	}

	for n, slice := range sliceTemplates {
		for i, elem := range slice {
			var err error
			slice[i], err = p.config.tpl.Process(elem, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(
					errs, fmt.Errorf("Error processing %s[%d]: %s", n, i, err))
			}
		}
	}

	if len(p.config.Scripts) == 0 && p.config.Inline == nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Either a script file or inline script must be specified."))
	} else if len(p.config.Scripts) > 0 && p.config.Inline != nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only a script file or an inline script can be specified, not both."))
	}

	for _, path := range p.config.Scripts {
		if _, err := os.Stat(path); err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Bad script '%s': %s", path, err))
		}
	}

	// Do a check for bad environment variables, such as '=foo', 'foobar'
	for _, kv := range p.config.Vars {
		vs := strings.Split(kv, "=")
		if len(vs) != 2 || vs[0] == "" {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	scripts := make([]string, len(p.config.Scripts))
	copy(scripts, p.config.Scripts)

	// If we have an inline script, then turn that into a temporary
	// shell script and use that.
	if p.config.Inline != nil {
		tf, err := ioutil.TempFile("", "packer-shell")
		if err != nil {
			return fmt.Errorf("Error preparing shell script: %s", err)
		}
		defer os.Remove(tf.Name())

		// Set the path to the temporary file
		scripts = append(scripts, tf.Name())

		// Write our contents to it
		writer := bufio.NewWriter(tf)
		writer.WriteString(fmt.Sprintf("#!%s\n", p.config.InlineShebang))
		for _, command := range p.config.Inline {
			if _, err := writer.WriteString(command + "\n"); err != nil {
				return fmt.Errorf("Error preparing shell script: %s", err)
			}
		}

		if err := writer.Flush(); err != nil {
			return fmt.Errorf("Error preparing shell script: %s", err)
		}

		tf.Close()
	}

	// Build our variables up by adding in the build name and builder type
	envVars := make([]string, len(p.config.Vars)+2)
	envVars[0] = "PACKER_BUILD_NAME=" + p.config.PackerBuildName
	envVars[1] = "PACKER_BUILDER_TYPE=" + p.config.PackerBuilderType
	copy(envVars[2:], p.config.Vars)

	for _, path := range scripts {
		ui.Say(fmt.Sprintf("Provisioning with shell script: %s", path))

		log.Printf("Opening %s for reading", path)
		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Error opening shell script: %s", err)
		}
		defer f.Close()

		log.Printf("Uploading %s => %s", path, p.config.RemotePath)
		err = comm.Upload(p.config.RemotePath, f)
		if err != nil {
			return fmt.Errorf("Error uploading shell script: %s", err)
		}

		// Close the original file since we copied it
		f.Close()

		// Flatten the environment variables
		flattendVars := strings.Join(envVars, " ")

		// Compile the command
		command := p.config.tpl.Process(p.config.ExecuteCommand, &ExecuteCommandTemplate{
			Vars: flattendVars,
			Path: p.config.RemotePath,
		})

		cmd := &packer.RemoteCmd{Command: command}
		log.Printf("Executing command: %s", cmd.Command)
		if err := cmd.StartWithUi(comm, ui); err != nil {
			return fmt.Errorf("Failed executing command: %s", err)
		}

		if cmd.ExitStatus != 0 {
			return fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus)
		}
	}

	return nil
}
<commit_msg>provisioner/shell: tests passing and compiling<commit_after>// This package implements a provisioner for Packer that executes
// shell scripts within the remote machine.
package shell

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/packer"
	"io/ioutil"
	"log"
	"os"
	"strings"
)

const DefaultRemotePath = "/tmp/script.sh"

type config struct {
	// An inline script to execute. Multiple strings are all executed
	// in the context of a single shell.
	Inline []string

	// The shebang value used when running inline scripts.
	InlineShebang string `mapstructure:"inline_shebang"`

	// The local path of the shell script to upload and execute.
	Script string

	// An array of multiple scripts to run.
	Scripts []string

	// An array of environment variables that will be injected before
	// your command(s) are executed.
	Vars []string `mapstructure:"environment_vars"`

	// The remote path where the local shell script will be uploaded to.
	// This should be set to a writable file that is in a pre-existing directory.
	RemotePath string `mapstructure:"remote_path"`

	// The command used to execute the script. The '{{ .Path }}' variable
	// should be used to specify where the script goes, {{ .Vars }}
	// can be used to inject the environment_vars into the environment.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Packer configurations, these come from Packer itself
	PackerBuildName   string `mapstructure:"packer_build_name"`
	PackerBuilderType string `mapstructure:"packer_builder_type"`

	tpl *common.Template
}

type Provisioner struct {
	config config
}

type ExecuteCommandTemplate struct {
	Vars string
	Path string
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	md, err := common.DecodeConfig(&p.config, raws...)
	if err != nil {
		return err
	}

	p.config.tpl, err = common.NewTemplate()
	if err != nil {
		return err
	}

	// Accumulate any errors
	errs := common.CheckUnusedConfig(md)

	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "chmod +x {{.Path}}; {{.Vars}} {{.Path}}"
	}

	if p.config.Inline != nil && len(p.config.Inline) == 0 {
		p.config.Inline = nil
	}

	if p.config.InlineShebang == "" {
		p.config.InlineShebang = "/bin/sh"
	}

	if p.config.RemotePath == "" {
		p.config.RemotePath = DefaultRemotePath
	}

	if p.config.Scripts == nil {
		p.config.Scripts = make([]string, 0)
	}

	if p.config.Vars == nil {
		p.config.Vars = make([]string, 0)
	}

	if p.config.Script != "" && len(p.config.Scripts) > 0 {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only one of script or scripts can be specified."))
	}

	if p.config.Script != "" {
		p.config.Scripts = []string{p.config.Script}
	}

	templates := map[string]*string{
		"inline_shebang": &p.config.InlineShebang,
		"script":         &p.config.Script,
		"remote_path":    &p.config.RemotePath,
	}

	for n, ptr := range templates {
		var err error
		*ptr, err = p.config.tpl.Process(*ptr, nil)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", n, err))
		}
	}

	sliceTemplates := map[string][]string{
		"inline":           p.config.Inline,
		"scripts":          p.config.Scripts,
		"environment_vars": p.config.Vars,
	}

	for n, slice := range sliceTemplates {
		for i, elem := range slice {
			var err error
			slice[i], err = p.config.tpl.Process(elem, nil)
			if err != nil {
				errs = packer.MultiErrorAppend(
					errs, fmt.Errorf("Error processing %s[%d]: %s", n, i, err))
			}
		}
	}

	if len(p.config.Scripts) == 0 && p.config.Inline == nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Either a script file or inline script must be specified."))
	} else if len(p.config.Scripts) > 0 && p.config.Inline != nil {
		errs = packer.MultiErrorAppend(errs,
			errors.New("Only a script file or an inline script can be specified, not both."))
	}

	for _, path := range p.config.Scripts {
		if _, err := os.Stat(path); err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Bad script '%s': %s", path, err))
		}
	}

	// Do a check for bad environment variables, such as '=foo', 'foobar'
	for _, kv := range p.config.Vars {
		vs := strings.Split(kv, "=")
		if len(vs) != 2 || vs[0] == "" {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("Environment variable not in format 'key=value': %s", kv))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	scripts := make([]string, len(p.config.Scripts))
	copy(scripts, p.config.Scripts)

	// If we have an inline script, then turn that into a temporary
	// shell script and use that.
	if p.config.Inline != nil {
		tf, err := ioutil.TempFile("", "packer-shell")
		if err != nil {
			return fmt.Errorf("Error preparing shell script: %s", err)
		}
		defer os.Remove(tf.Name())

		// Set the path to the temporary file
		scripts = append(scripts, tf.Name())

		// Write our contents to it
		writer := bufio.NewWriter(tf)
		writer.WriteString(fmt.Sprintf("#!%s\n", p.config.InlineShebang))
		for _, command := range p.config.Inline {
			if _, err := writer.WriteString(command + "\n"); err != nil {
				return fmt.Errorf("Error preparing shell script: %s", err)
			}
		}

		if err := writer.Flush(); err != nil {
			return fmt.Errorf("Error preparing shell script: %s", err)
		}

		tf.Close()
	}

	// Build our variables up by adding in the build name and builder type
	envVars := make([]string, len(p.config.Vars)+2)
	envVars[0] = "PACKER_BUILD_NAME=" + p.config.PackerBuildName
	envVars[1] = "PACKER_BUILDER_TYPE=" + p.config.PackerBuilderType
	copy(envVars[2:], p.config.Vars)

	for _, path := range scripts {
		ui.Say(fmt.Sprintf("Provisioning with shell script: %s", path))

		log.Printf("Opening %s for reading", path)
		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Error opening shell script: %s", err)
		}
		defer f.Close()

		log.Printf("Uploading %s => %s", path, p.config.RemotePath)
		err = comm.Upload(p.config.RemotePath, f)
		if err != nil {
			return fmt.Errorf("Error uploading shell script: %s", err)
		}

		// Close the original file since we copied it
		f.Close()

		// Flatten the environment variables
		flattendVars := strings.Join(envVars, " ")

		// Compile the command
		command, err := p.config.tpl.Process(p.config.ExecuteCommand, &ExecuteCommandTemplate{
			Vars: flattendVars,
			Path: p.config.RemotePath,
		})
		if err != nil {
			return fmt.Errorf("Error processing command: %s", err)
		}

		cmd := &packer.RemoteCmd{Command: command}
		log.Printf("Executing command: %s", cmd.Command)
		if err := cmd.StartWithUi(comm, ui); err != nil {
			return fmt.Errorf("Failed executing command: %s", err)
		}

		if cmd.ExitStatus != 0 {
			return fmt.Errorf("Script exited with non-zero exit status: %d", cmd.ExitStatus)
		}
	}

	return nil
}
<|endoftext|>


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

package io

import (
	"bufio"
	"io"
)

// ResettableWriter is a resettable writer.
type ResettableWriter interface {
	io.Writer
	Flush() error
	Reset(w io.Writer)
}

// ResettableWriterOptions are options for a resettable writer.
type ResettableWriterOptions struct {
	WriteBufferSize int
}

// ResettableWriterFn creates a resettable writer.
type ResettableWriterFn func(r io.Writer, opts ResettableWriterOptions) ResettableWriter

// defaultResettableWriterFn creates a default resettable writer.
func defaultResettableWriterFn() ResettableWriterFn {
	return func(r io.Writer, opts ResettableWriterOptions) ResettableWriter {
		return bufio.NewWriterSize(r, opts.WriteBufferSize)
	}
}
<commit_msg>Use pass-through writer if no buffering is requested (#3178)<commit_after>// Copyright (c) 2020 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package io

import (
	"bufio"
	"io"
)

// ResettableWriter is a resettable writer.
type ResettableWriter interface {
	io.Writer
	Flush() error
	Reset(w io.Writer)
}

// ResettableWriterOptions are options for a resettable writer.
type ResettableWriterOptions struct {
	WriteBufferSize int
}

// ResettableWriterFn creates a resettable writer.
type ResettableWriterFn func(r io.Writer, opts ResettableWriterOptions) ResettableWriter

// defaultResettableWriterFn creates a default resettable writer.
func defaultResettableWriterFn() ResettableWriterFn {
	return func(w io.Writer, opts ResettableWriterOptions) ResettableWriter {
		if opts.WriteBufferSize <= 0 {
			return &passthroughResettableWriter{w: w}
		}
		return bufio.NewWriterSize(w, opts.WriteBufferSize)
	}
}

type passthroughResettableWriter struct {
	w io.Writer
}

func (p passthroughResettableWriter) Write(b []byte) (n int, err error) {
	return p.w.Write(b)
}

func (p passthroughResettableWriter) Flush() error {
	return nil
}

func (p *passthroughResettableWriter) Reset(w io.Writer) {
	p.w = w
}
<|endoftext|>


<commit_before>/*
This tool recursively searches $SRCSEARCHROOT for the directory queried and will return the path
of the most shallow result. Directories are searched in lexicographic order.
*/

package main

import (
	"container/list"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
)

const envRootVar = "SRCSEARCHROOT"

var (
	errNotFound = errors.New("could not find directory")

	ignoreHidden = flag.Bool("ignorehidden", true, "ignore hidden directories")
	maxDepth     = flag.Int("maxdepth", 5, "maximum search depth")
)

type location struct {
	path  string
	depth int
}

func search(dir, name string, startdepth int) (path string, err error) {
	q := list.New()
	q.PushBack(location{dir, startdepth})
	for q.Len() > 0 {
		front := q.Front()
		cloc := q.Remove(front).(location)
		entries, err := ioutil.ReadDir(cloc.path)
		if err != nil {
			return "", err
		}
		for _, e := range entries {
			if !e.IsDir() {
				continue
			}
			absPath := cloc.path + "/" + e.Name()
			if e.Name() == name {
				return absPath, nil
			}
			if cloc.depth < *maxDepth {
				q.PushBack(location{absPath, cloc.depth + 1})
			}
		}
	}
	return "", errNotFound
}

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s dirname\n", os.Args[0])
	flag.PrintDefaults()
}

// TODO: Add support for partial paths (i.e. dir1/dir2/.../dirN)
func main() {
	flag.Usage = usage
	flag.Parse()
	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(1)
	}
	name := flag.Arg(0)
	root := os.Getenv(envRootVar)
	if root == "" {
		fmt.Fprintln(os.Stderr, envRootVar+" must be set.")
		os.Exit(2)
	}
	p, err := search(root, name, 0)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	fmt.Println(p)
}
<commit_msg>Actually check if directory is hidden<commit_after>/*
This tool recursively searches $SRCSEARCHROOT for the directory queried and will return the path
of the most shallow result. Directories are searched in lexicographic order.
*/

package main

import (
	"container/list"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

const envRootVar = "SRCSEARCHROOT"

var (
	errNotFound = errors.New("could not find directory")

	ignoreHidden = flag.Bool("ignorehidden", true, "ignore hidden directories")
	maxDepth     = flag.Int("maxdepth", 5, "maximum search depth")
)

type location struct {
	path  string
	depth int
}

func search(dir, name string, startdepth int) (path string, err error) {
	q := list.New()
	q.PushBack(location{dir, startdepth})
	for q.Len() > 0 {
		front := q.Front()
		cloc := q.Remove(front).(location)
		entries, err := ioutil.ReadDir(cloc.path)
		if err != nil {
			return "", err
		}
		for _, e := range entries {
			if !e.IsDir() || (*ignoreHidden && strings.HasPrefix(e.Name(), ".")) {
				continue
			}
			absPath := cloc.path + "/" + e.Name()
			if e.Name() == name {
				return absPath, nil
			}
			if cloc.depth < *maxDepth {
				q.PushBack(location{absPath, cloc.depth + 1})
			}
		}
	}
	return "", errNotFound
}

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s dirname\n", os.Args[0])
	flag.PrintDefaults()
}

// TODO: Add support for partial paths (i.e. dir1/dir2/.../dirN)
func main() {
	flag.Usage = usage
	flag.Parse()
	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(1)
	}
	name := flag.Arg(0)
	root := os.Getenv(envRootVar)
	if root == "" {
		fmt.Fprintln(os.Stderr, envRootVar+" must be set.")
		os.Exit(2)
	}
	p, err := search(root, name, 0)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(3)
	}
	fmt.Println(p)
}
<|endoftext|>


<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 cli

import (
	"fmt"
	"time"

	"github.com/spf13/cobra"
	"mynewt.apache.org/newt/newtmgr/config"
	"mynewt.apache.org/newt/newtmgr/protocol"
	"mynewt.apache.org/newt/newtmgr/transport"
)

func crashRunCmd(cmd *cobra.Command, args []string) {
	cpm, err := config.NewConnProfileMgr()
	if err != nil {
		nmUsage(cmd, err)
	}

	if len(args) != 1 {
		nmUsage(cmd, nil)
	}
	crashType := args[0]

	profile, err := cpm.GetConnProfile(ConnProfileName)
	if err != nil {
		nmUsage(cmd, err)
	}

	conn, err := transport.NewConnWithTimeout(profile, time.Second*1)
	if err != nil {
		nmUsage(cmd, err)
	}
	defer conn.Close()

	runner, err := protocol.NewCmdRunner(conn)
	if err != nil {
		nmUsage(cmd, err)
	}

	crash, err := protocol.NewCrash(crashType)
	if err != nil {
		nmUsage(cmd, err)
	}

	nmr, err := crash.EncodeWriteRequest()
	if err != nil {
		nmUsage(cmd, err)
	}

	if err := runner.WriteReq(nmr); err != nil {
		nmUsage(cmd, err)
	}

	rsp, err := runner.ReadResp()
	if err == nil {
		cRsp, err := protocol.DecodeCrashResponse(rsp.Data)
		if err != nil {
			nmUsage(cmd, err)
		}
		if cRsp.Err != 0 {
			fmt.Printf("Failed, error:%d\n", cRsp.Err)
		}
	}
	fmt.Println("Done")
}

func crashCmd() *cobra.Command {
	crashEx := "   newtmgr -c olimex crash div0\n"

	crashCmd := &cobra.Command{
		Use:     "crash [div0|jump0|ref0|assert]",
		Short:   "Send crash command to remote endpoint using newtmgr",
		Example: crashEx,
		Run:     crashRunCmd,
	}

	return crashCmd
}
<commit_msg>newtmgr; advertise wdog as a crash option.<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 cli

import (
	"fmt"
	"time"

	"github.com/spf13/cobra"
	"mynewt.apache.org/newt/newtmgr/config"
	"mynewt.apache.org/newt/newtmgr/protocol"
	"mynewt.apache.org/newt/newtmgr/transport"
)

func crashRunCmd(cmd *cobra.Command, args []string) {
	cpm, err := config.NewConnProfileMgr()
	if err != nil {
		nmUsage(cmd, err)
	}

	if len(args) != 1 {
		nmUsage(cmd, nil)
	}
	crashType := args[0]

	profile, err := cpm.GetConnProfile(ConnProfileName)
	if err != nil {
		nmUsage(cmd, err)
	}

	conn, err := transport.NewConnWithTimeout(profile, time.Second*1)
	if err != nil {
		nmUsage(cmd, err)
	}
	defer conn.Close()

	runner, err := protocol.NewCmdRunner(conn)
	if err != nil {
		nmUsage(cmd, err)
	}

	crash, err := protocol.NewCrash(crashType)
	if err != nil {
		nmUsage(cmd, err)
	}

	nmr, err := crash.EncodeWriteRequest()
	if err != nil {
		nmUsage(cmd, err)
	}

	if err := runner.WriteReq(nmr); err != nil {
		nmUsage(cmd, err)
	}

	rsp, err := runner.ReadResp()
	if err == nil {
		cRsp, err := protocol.DecodeCrashResponse(rsp.Data)
		if err != nil {
			nmUsage(cmd, err)
		}
		if cRsp.Err != 0 {
			fmt.Printf("Failed, error:%d\n", cRsp.Err)
		}
	}
	fmt.Println("Done")
}

func crashCmd() *cobra.Command {
	crashEx := "   newtmgr -c olimex crash div0\n"

	crashCmd := &cobra.Command{
		Use:     "crash [div0|jump0|ref0|assert|wdog]",
		Short:   "Send crash command to remote endpoint using newtmgr",
		Example: crashEx,
		Run:     crashRunCmd,
	}

	return crashCmd
}
<|endoftext|>


<commit_before>package state

import (
	"fmt"
	"launchpad.net/gozk/zookeeper"
)

// RelationRole defines the role of a relation endpoint.
type RelationRole string

const (
	RoleProvider RelationRole = "provider"
	RoleRequirer RelationRole = "requirer"
	RolePeer     RelationRole = "peer"
)

// CounterpartRole returns the RelationRole that this RelationRole can
// relate to.
func (r RelationRole) CounterpartRole() RelationRole {
	switch r {
	case RoleProvider:
		return RoleRequirer
	case RoleRequirer:
		return RoleProvider
	case RolePeer:
		return RolePeer
	}
	panic(fmt.Errorf("unknown RelationRole: %q", r))
}

// RelationScope describes the scope of a relation endpoint.
type RelationScope string

const (
	ScopeGlobal    RelationScope = "global"
	ScopeContainer RelationScope = "container"
)

// RelationEndpoint represents one endpoint of a relation.
type RelationEndpoint struct {
	ServiceName   string
	Interface     string
	RelationName  string
	RelationRole  RelationRole
	RelationScope RelationScope
}

// CanRelateTo returns whether a relation may be established between e and other.
func (e *RelationEndpoint) CanRelateTo(other *RelationEndpoint) bool {
	if e.Interface != other.Interface {
		return false
	}
	if e.RelationRole == RolePeer {
		// Peer relations do not currently work with multiple endpoints.
		return false
	}
	return e.RelationRole.CounterpartRole() == other.RelationRole
}

// String returns the unique identifier of the relation endpoint.
func (e RelationEndpoint) String() string {
	return e.ServiceName + ":" + e.RelationName
}

// ServiceRelation represents an established relation from
// the viewpoint of a participant service.
type ServiceRelation struct {
	st            *State
	relationKey   string
	serviceKey    string
	relationScope RelationScope
	relationRole  RelationRole
	relationName  string
}

// RelationScope returns the scope of the relation.
func (r *ServiceRelation) RelationScope() RelationScope {
	return r.relationScope
}

// RelationRole returns the service role within the relation.
func (r *ServiceRelation) RelationRole() RelationRole {
	return r.relationRole
}

// RelationName returns the name this relation has within the service.
func (r *ServiceRelation) RelationName() string {
	return r.relationName
}

// unitScope represents a set of units that can (transitively) affect one
// another within the context of a particular relation. For a globally-scoped
// relation, the unitScope holds every unit of every service in the relation;
// for a container-scoped relation, the unitScope holds every unit of the
// relation that is located within a particular container.
type unitScope struct {
	zk   *zookeeper.Conn
	path string
}

// SettingsPath returns the path to the relation unit settings node for the
// unit identified by key, or to the relation group settings node if key is
// empty.
func (s *unitScope) SettingsPath(key string) string {
	return s.subpath("settings", key)
}

// PresencePath returns the path to the relation unit presence node for a
// unit (identified by key) of a service acting as role; or to the relation
// group role node if key is empty.
func (s *unitScope) PresencePath(role RelationRole, key string) string {
	return s.subpath(string(role), key)
}

// PrepareJoin ensures that ZooKeeper nodes exist such that a unit of a
// service with the supplied role will be able to join the relation.
func (s *unitScope) PrepareJoin(role RelationRole) error {
	paths := []string{
		s.path,
		s.SettingsPath(""),
		s.PresencePath(role, ""),
	}
	for _, path := range paths {
		if _, err := s.zk.Create(path, "", 0, zkPermAll); err != nil {
			if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
				continue
			}
			return err
		}
	}
	return nil
}

// subpath returns an absolute ZooKeeper path to the node whose path relative
// to the group node is composed of parts. Empty parts will be stripped.
func (s *unitScope) subpath(parts ...string) string {
	path := s.path
	for _, part := range parts {
		if part != "" {
			path = path + "/" + part
		}
	}
	return path
}
<commit_msg>expand docs slightly<commit_after>package state

import (
	"fmt"
	"launchpad.net/gozk/zookeeper"
)

// RelationRole defines the role of a relation endpoint.
type RelationRole string

const (
	RoleProvider RelationRole = "provider"
	RoleRequirer RelationRole = "requirer"
	RolePeer     RelationRole = "peer"
)

// CounterpartRole returns the RelationRole that this RelationRole can
// relate to.
func (r RelationRole) CounterpartRole() RelationRole {
	switch r {
	case RoleProvider:
		return RoleRequirer
	case RoleRequirer:
		return RoleProvider
	case RolePeer:
		return RolePeer
	}
	panic(fmt.Errorf("unknown RelationRole: %q", r))
}

// RelationScope describes the scope of a relation endpoint.
type RelationScope string

const (
	ScopeGlobal    RelationScope = "global"
	ScopeContainer RelationScope = "container"
)

// RelationEndpoint represents one endpoint of a relation.
type RelationEndpoint struct {
	ServiceName   string
	Interface     string
	RelationName  string
	RelationRole  RelationRole
	RelationScope RelationScope
}

// CanRelateTo returns whether a relation may be established between e and other.
func (e *RelationEndpoint) CanRelateTo(other *RelationEndpoint) bool {
	if e.Interface != other.Interface {
		return false
	}
	if e.RelationRole == RolePeer {
		// Peer relations do not currently work with multiple endpoints.
		return false
	}
	return e.RelationRole.CounterpartRole() == other.RelationRole
}

// String returns the unique identifier of the relation endpoint.
func (e RelationEndpoint) String() string {
	return e.ServiceName + ":" + e.RelationName
}

// ServiceRelation represents an established relation from
// the viewpoint of a participant service.
type ServiceRelation struct {
	st            *State
	relationKey   string
	serviceKey    string
	relationScope RelationScope
	relationRole  RelationRole
	relationName  string
}

// RelationScope returns the scope of the relation.
func (r *ServiceRelation) RelationScope() RelationScope {
	return r.relationScope
}

// RelationRole returns the service role within the relation.
func (r *ServiceRelation) RelationRole() RelationRole {
	return r.relationRole
}

// RelationName returns the name this relation has within the service.
func (r *ServiceRelation) RelationName() string {
	return r.relationName
}

// unitScope represents a set of units that can (transitively) affect one
// another within the context of a particular relation. For a globally-scoped
// relation, the unitScope holds every unit of every service in the relation;
// for a container-scoped relation, the unitScope holds every unit of the
// relation that is located within a particular container.
// Thus, unitScope paths will take one of the following forms:
//
//   /relations/<relation-id>
//   /relations/<relation-id>/<container-id>
type unitScope struct {
	zk   *zookeeper.Conn
	path string
}

// SettingsPath returns the path to the relation unit settings node for the
// unit identified by key, or to the relation group settings node if key is
// empty.
func (s *unitScope) SettingsPath(key string) string {
	return s.subpath("settings", key)
}

// PresencePath returns the path to the relation unit presence node for a
// unit (identified by key) of a service acting as role; or to the relation
// group role node if key is empty.
func (s *unitScope) PresencePath(role RelationRole, key string) string {
	return s.subpath(string(role), key)
}

// PrepareJoin ensures that ZooKeeper nodes exist such that a unit of a
// service with the supplied role will be able to join the relation.
func (s *unitScope) PrepareJoin(role RelationRole) error {
	paths := []string{
		s.path,
		s.SettingsPath(""),
		s.PresencePath(role, ""),
	}
	for _, path := range paths {
		if _, err := s.zk.Create(path, "", 0, zkPermAll); err != nil {
			if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
				continue
			}
			return err
		}
	}
	return nil
}

// subpath returns an absolute ZooKeeper path to the node whose path relative
// to the group node is composed of parts. Empty parts will be stripped.
func (s *unitScope) subpath(parts ...string) string {
	path := s.path
	for _, part := range parts {
		if part != "" {
			path = path + "/" + part
		}
	}
	return path
}
<|endoftext|>


<commit_before>package pixel

// Drawer glues all the fundamental interfaces (Target, Triangles, Picture) into a coherent and the
// only intended usage pattern.
//
// Drawer makes it possible to draw any combination of Triangles and Picture onto any Target
// efficiently.
//
// To create a Drawer, just assign it's Triangles and Picture fields:
//
//   d := pixel.Drawer{Triangles: t, Picture: p}
//
// If Triangles is nil, nothing will be drawn. If Picture is nil, Triangles will be drawn without a
// Picture.
//
// Whenever you change the Triangles, call Dirty to notify Drawer that Triangles changed. You don't
// need to notify Drawer about a change of the Picture.
//
// Note, that Drawer caches the results of MakePicture from Targets it's drawn to for each Picture
// it's set to. What it means is that using a Drawer with an unbounded number of Pictures leads to a
// memory leak, since Drawer caches them and never forgets. In such a situation, create a new Drawer
// for each Picture.
type Drawer struct {
	Triangles Triangles
	Picture   Picture

	tris   map[Target]TargetTriangles
	clean  map[Target]bool
	pics   map[targetPicturePair]TargetPicture
	dirty  bool
	inited bool
}

type targetPicturePair struct {
	Target  Target
	Picture Picture
}

func (d *Drawer) lazyInit() {
	if !d.inited {
		d.tris = make(map[Target]TargetTriangles)
		d.clean = make(map[Target]bool)
		d.pics = make(map[targetPicturePair]TargetPicture)
		d.inited = true
	}
}

// Dirty marks the Triangles of this Drawer as changed. If not called, changes will not be visible
// when drawing.
func (d *Drawer) Dirty() {
	d.lazyInit()

	d.dirty = true
}

// Draw efficiently draws Triangles with Picture onto the provided Target.
//
// If Triangles is nil, nothing will be drawn. If Picture is nil, Triangles will be drawn without a
// Picture.
func (d *Drawer) Draw(t Target) {
	d.lazyInit()

	if d.dirty {
		for t := range d.clean {
			d.clean[t] = false
		}
		d.dirty = false
	}

	if d.Triangles == nil {
		return
	}

	tri := d.tris[t]
	if tri == nil {
		tri = t.MakeTriangles(d.Triangles)
		d.tris[t] = tri
		d.clean[t] = true
	}

	if !d.clean[t] {
		tri.SetLen(d.Triangles.Len())
		tri.Update(d.Triangles)
		d.clean[t] = true
	}

	if d.Picture == nil {
		tri.Draw()
		return
	}

	pic := d.pics[targetPicturePair{t, d.Picture}]
	if pic == nil {
		pic = t.MakePicture(d.Picture)
		d.pics[targetPicturePair{t, d.Picture}] = pic
	}

	pic.Draw(tri)
}
<commit_msg>optimize Drawer (reduce map access)<commit_after>package pixel

// Drawer glues all the fundamental interfaces (Target, Triangles, Picture) into a coherent and the
// only intended usage pattern.
//
// Drawer makes it possible to draw any combination of Triangles and Picture onto any Target
// efficiently.
//
// To create a Drawer, just assign it's Triangles and Picture fields:
//
//   d := pixel.Drawer{Triangles: t, Picture: p}
//
// If Triangles is nil, nothing will be drawn. If Picture is nil, Triangles will be drawn without a
// Picture.
//
// Whenever you change the Triangles, call Dirty to notify Drawer that Triangles changed. You don't
// need to notify Drawer about a change of the Picture.
//
// Note, that Drawer caches the results of MakePicture from Targets it's drawn to for each Picture
// it's set to. What it means is that using a Drawer with an unbounded number of Pictures leads to a
// memory leak, since Drawer caches them and never forgets. In such a situation, create a new Drawer
// for each Picture.
type Drawer struct {
	Triangles Triangles
	Picture   Picture

	targets map[Target]*drawerTarget
	inited  bool
}

type drawerTarget struct {
	tris  TargetTriangles
	pics  map[Picture]TargetPicture
	clean bool
}

func (d *Drawer) lazyInit() {
	if !d.inited {
		d.targets = make(map[Target]*drawerTarget)
		d.inited = true
	}
}

// Dirty marks the Triangles of this Drawer as changed. If not called, changes will not be visible
// when drawing.
func (d *Drawer) Dirty() {
	d.lazyInit()

	for _, t := range d.targets {
		t.clean = false
	}
}

// Draw efficiently draws Triangles with Picture onto the provided Target.
//
// If Triangles is nil, nothing will be drawn. If Picture is nil, Triangles will be drawn without a
// Picture.
func (d *Drawer) Draw(t Target) {
	d.lazyInit()

	if d.Triangles == nil {
		return
	}

	dt := d.targets[t]
	if dt == nil {
		dt = &drawerTarget{
			pics: make(map[Picture]TargetPicture),
		}
		d.targets[t] = dt
	}

	if dt.tris == nil {
		dt.tris = t.MakeTriangles(d.Triangles)
		dt.clean = true
	}

	if !dt.clean {
		dt.tris.SetLen(d.Triangles.Len())
		dt.tris.Update(d.Triangles)
		dt.clean = true
	}

	if d.Picture == nil {
		dt.tris.Draw()
		return
	}

	pic := dt.pics[d.Picture]
	if pic == nil {
		pic = t.MakePicture(d.Picture)
		dt.pics[d.Picture] = pic
	}

	pic.Draw(dt.tris)
}
<|endoftext|>


<commit_before>package gorethink

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	p "gopkg.in/dancannon/gorethink.v1/ql2"
)

var (
	// ErrNoHosts is returned when no hosts to the Connect method.
	ErrNoHosts = errors.New("no hosts provided")
	// ErrNoConnectionsStarted is returned when the driver couldn't to any of
	// the provided hosts.
	ErrNoConnectionsStarted = errors.New("no connections were made when creating the session")
	// ErrInvalidNode is returned when attempting to connect to a node which
	// returns an invalid response.
	ErrInvalidNode = errors.New("invalid node")
	// ErrNoConnections is returned when there are no active connections in the
	// clusters connection pool.
	ErrNoConnections = errors.New("gorethink: no connections were available")
	// ErrConnectionClosed is returned when trying to send a query with a closed
	// connection.
	ErrConnectionClosed = errors.New("gorethink: the connection is closed")
)

func printCarrots(t Term, frames []*p.Frame) string {
	var frame *p.Frame
	if len(frames) > 1 {
		frame, frames = frames[0], frames[1:]
	} else if len(frames) == 1 {
		frame, frames = frames[0], []*p.Frame{}
	}

	for i, arg := range t.args {
		if frame.GetPos() == int64(i) {
			t.args[i] = Term{
				termType: p.Term_DATUM,
				data:     printCarrots(arg, frames),
			}
		}
	}

	for k, arg := range t.optArgs {
		if frame.GetOpt() == k {
			t.optArgs[k] = Term{
				termType: p.Term_DATUM,
				data:     printCarrots(arg, frames),
			}
		}
	}

	b := &bytes.Buffer{}
	for _, c := range t.String() {
		if c != '^' {
			b.WriteString(" ")
		} else {
			b.WriteString("^")
		}
	}

	return b.String()
}

// Error constants
var ErrEmptyResult = errors.New("The result does not contain any more rows")

// Connection/Response errors

// rqlResponseError is the base type for all errors, it formats both
// for the response and query if set.
type rqlServerError struct {
	response *Response
	term     *Term
}

func (e rqlServerError) Error() string {
	var err = "An error occurred"
	if e.response != nil {
		json.Unmarshal(e.response.Responses[0], &err)
	}

	if e.term == nil {
		return fmt.Sprintf("gorethink: %s", err)
	}

	return fmt.Sprintf("gorethink: %s in: \n%s", err, e.term.String())

}

func (e rqlServerError) String() string {
	return e.Error()
}

type rqlError string

func (e rqlError) Error() string {
	return fmt.Sprintf("gorethink: %s", string(e))
}

func (e rqlError) String() string {
	return e.Error()
}

// Exported Error "Implementations"

type RQLClientError struct{ rqlServerError }
type RQLCompileError struct{ rqlServerError }
type RQLDriverCompileError struct{ RQLCompileError }
type RQLServerCompileError struct{ RQLCompileError }
type RQLAuthError struct{ RQLDriverError }
type RQLRuntimeError struct{ rqlServerError }

type RQLQueryLogicError struct{ RQLRuntimeError }
type RQLNonExistenceError struct{ RQLQueryLogicError }
type RQLResourceLimitError struct{ RQLRuntimeError }
type RQLUserError struct{ RQLRuntimeError }
type RQLInternalError struct{ RQLRuntimeError }
type RQLTimeoutError struct{ rqlServerError }
type RQLAvailabilityError struct{ RQLRuntimeError }
type RQLOpFailedError struct{ RQLAvailabilityError }
type RQLOpIndeterminateError struct{ RQLAvailabilityError }

// RQLDriverError represents an unexpected error with the driver, if this error
// persists please create an issue.
type RQLDriverError struct {
	rqlError
}

// RQLConnectionError represents an error when communicating with the database
// server.
type RQLConnectionError struct {
	rqlError
}

func createRuntimeError(errorType p.Response_ErrorType, response *Response, term *Term) error {
	serverErr := rqlServerError{response, term}

	switch errorType {
	case p.Response_QUERY_LOGIC:
		return RQLQueryLogicError{RQLRuntimeError{serverErr}}
	case p.Response_NON_EXISTENCE:
		return RQLNonExistenceError{RQLQueryLogicError{RQLRuntimeError{serverErr}}}
	case p.Response_RESOURCE_LIMIT:
		return RQLResourceLimitError{RQLRuntimeError{serverErr}}
	case p.Response_USER:
		return RQLUserError{RQLRuntimeError{serverErr}}
	case p.Response_INTERNAL:
		return RQLInternalError{RQLRuntimeError{serverErr}}
	case p.Response_OP_FAILED:
		return RQLOpFailedError{RQLAvailabilityError{RQLRuntimeError{serverErr}}}
	case p.Response_OP_INDETERMINATE:
		return RQLOpIndeterminateError{RQLAvailabilityError{RQLRuntimeError{serverErr}}}
	default:
		return RQLRuntimeError{serverErr}
	}
}

// Error type helpers

// IsConflictErr returns true if the error is non-nil and the query failed
// due to a duplicate primary key.
func IsConflictErr(err error) bool {
	return strings.HasPrefix(err.Error(), "Duplicate primary key")
}

// IsTypeErr returns true if the error is non-nil and the query failed due
// to a type error.
func IsTypeErr(err error) bool {
	return strings.HasPrefix(err.Error(), "Expected type")
}
<commit_msg>Check for nil in IsConflictErr and IsTypeErr<commit_after>package gorethink

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	p "gopkg.in/dancannon/gorethink.v1/ql2"
)

var (
	// ErrNoHosts is returned when no hosts to the Connect method.
	ErrNoHosts = errors.New("no hosts provided")
	// ErrNoConnectionsStarted is returned when the driver couldn't to any of
	// the provided hosts.
	ErrNoConnectionsStarted = errors.New("no connections were made when creating the session")
	// ErrInvalidNode is returned when attempting to connect to a node which
	// returns an invalid response.
	ErrInvalidNode = errors.New("invalid node")
	// ErrNoConnections is returned when there are no active connections in the
	// clusters connection pool.
	ErrNoConnections = errors.New("gorethink: no connections were available")
	// ErrConnectionClosed is returned when trying to send a query with a closed
	// connection.
	ErrConnectionClosed = errors.New("gorethink: the connection is closed")
)

func printCarrots(t Term, frames []*p.Frame) string {
	var frame *p.Frame
	if len(frames) > 1 {
		frame, frames = frames[0], frames[1:]
	} else if len(frames) == 1 {
		frame, frames = frames[0], []*p.Frame{}
	}

	for i, arg := range t.args {
		if frame.GetPos() == int64(i) {
			t.args[i] = Term{
				termType: p.Term_DATUM,
				data:     printCarrots(arg, frames),
			}
		}
	}

	for k, arg := range t.optArgs {
		if frame.GetOpt() == k {
			t.optArgs[k] = Term{
				termType: p.Term_DATUM,
				data:     printCarrots(arg, frames),
			}
		}
	}

	b := &bytes.Buffer{}
	for _, c := range t.String() {
		if c != '^' {
			b.WriteString(" ")
		} else {
			b.WriteString("^")
		}
	}

	return b.String()
}

// Error constants
var ErrEmptyResult = errors.New("The result does not contain any more rows")

// Connection/Response errors

// rqlResponseError is the base type for all errors, it formats both
// for the response and query if set.
type rqlServerError struct {
	response *Response
	term     *Term
}

func (e rqlServerError) Error() string {
	var err = "An error occurred"
	if e.response != nil {
		json.Unmarshal(e.response.Responses[0], &err)
	}

	if e.term == nil {
		return fmt.Sprintf("gorethink: %s", err)
	}

	return fmt.Sprintf("gorethink: %s in: \n%s", err, e.term.String())

}

func (e rqlServerError) String() string {
	return e.Error()
}

type rqlError string

func (e rqlError) Error() string {
	return fmt.Sprintf("gorethink: %s", string(e))
}

func (e rqlError) String() string {
	return e.Error()
}

// Exported Error "Implementations"

type RQLClientError struct{ rqlServerError }
type RQLCompileError struct{ rqlServerError }
type RQLDriverCompileError struct{ RQLCompileError }
type RQLServerCompileError struct{ RQLCompileError }
type RQLAuthError struct{ RQLDriverError }
type RQLRuntimeError struct{ rqlServerError }

type RQLQueryLogicError struct{ RQLRuntimeError }
type RQLNonExistenceError struct{ RQLQueryLogicError }
type RQLResourceLimitError struct{ RQLRuntimeError }
type RQLUserError struct{ RQLRuntimeError }
type RQLInternalError struct{ RQLRuntimeError }
type RQLTimeoutError struct{ rqlServerError }
type RQLAvailabilityError struct{ RQLRuntimeError }
type RQLOpFailedError struct{ RQLAvailabilityError }
type RQLOpIndeterminateError struct{ RQLAvailabilityError }

// RQLDriverError represents an unexpected error with the driver, if this error
// persists please create an issue.
type RQLDriverError struct {
	rqlError
}

// RQLConnectionError represents an error when communicating with the database
// server.
type RQLConnectionError struct {
	rqlError
}

func createRuntimeError(errorType p.Response_ErrorType, response *Response, term *Term) error {
	serverErr := rqlServerError{response, term}

	switch errorType {
	case p.Response_QUERY_LOGIC:
		return RQLQueryLogicError{RQLRuntimeError{serverErr}}
	case p.Response_NON_EXISTENCE:
		return RQLNonExistenceError{RQLQueryLogicError{RQLRuntimeError{serverErr}}}
	case p.Response_RESOURCE_LIMIT:
		return RQLResourceLimitError{RQLRuntimeError{serverErr}}
	case p.Response_USER:
		return RQLUserError{RQLRuntimeError{serverErr}}
	case p.Response_INTERNAL:
		return RQLInternalError{RQLRuntimeError{serverErr}}
	case p.Response_OP_FAILED:
		return RQLOpFailedError{RQLAvailabilityError{RQLRuntimeError{serverErr}}}
	case p.Response_OP_INDETERMINATE:
		return RQLOpIndeterminateError{RQLAvailabilityError{RQLRuntimeError{serverErr}}}
	default:
		return RQLRuntimeError{serverErr}
	}
}

// Error type helpers

// IsConflictErr returns true if the error is non-nil and the query failed
// due to a duplicate primary key.
func IsConflictErr(err error) bool {
	if err == nil {
		return false
	}

	return strings.HasPrefix(err.Error(), "Duplicate primary key")
}

// IsTypeErr returns true if the error is non-nil and the query failed due
// to a type error.
func IsTypeErr(err error) bool {
	if err == nil {
		return false
	}

	return strings.HasPrefix(err.Error(), "Expected type")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"log"
	"net/http"
)

type response struct {
	Message string
}

func main() {
	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		log.Println("Req: ", r.URL)

		id, present := r.URL.Query()["id"]

		log.Println(present)

		if present {
			log.Println("Got id: ", id[0])

			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(200)
			json.NewEncoder(w).Encode(response{"OK"})
		} else {
			log.Println("No Id present, rejecting")

			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(404)
			json.NewEncoder(w).Encode(response{"Invalid ID"})
		}
	})

	log.Println("Listening on 8099")
	log.Fatal(http.ListenAndServe(":8099", nil))
}
<commit_msg>Delete leftover file<commit_after><|endoftext|>


<commit_before>// garble produces pseudo random bytes based on a phrase
// and uses it to garble and ungarble files
package main

// #include <stdint.h>
// #define BSIZE 65536
//
// void xor(int64_t *a, int64_t *b) {
//     int i = BSIZE / 8;
//     while(i--) {
//         a[i] ^= b[i];
//     }
// }
import "C"
import "unsafe"

import (
	"bytes"
	cryptorand "crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"runtime"
	"runtime/pprof"
	"time"
)

const (
	BSIZE    = C.BSIZE
	BSIZE7   = BSIZE - BSIZEMOD
	BSIZEMOD = BSIZE % 7
	MULTI    = 4
	SOURCES  = 8
	POOL     = SOURCES * (MULTI + 1)
)

var (
	narg   = int(0)
	phrase = ""
	pool   = make(chan []byte, POOL)
	loop   = make(chan []byte, POOL)
	data   = make([]chan []byte, SOURCES)
)

// randomSeed produces a int64 seed based on crypto/rand and time.
func randomSeed() int64 {
	var seed int64

	urandom := make([]byte, 8)
	cryptorand.Reader.Read(urandom)

	for key, value := range urandom {
		seed ^= (int64(value) ^ time.Now().UTC().UnixNano()) << (uint(key) * 8)
	}

	return seed
}

// randomBytes fills byte buffers with random data
func randomBytes(src rand.Source, out chan<- []byte) {
	var (
		r int64
		i = BSIZE
	)

	for buf, ok := <-pool; ok; buf, ok = <-pool {
		r = src.Int63()
		switch { // Go seems to eliminate impossible cases
		case BSIZEMOD == 6:
			buf[BSIZE-6] = byte(r >> 48)
			fallthrough
		case BSIZEMOD == 5:
			buf[BSIZE-5] = byte(r >> 32)
			fallthrough
		case BSIZEMOD == 4:
			buf[BSIZE-4] = byte(r >> 24)
			fallthrough
		case BSIZEMOD == 3:
			buf[BSIZE-3] = byte(r >> 16)
			fallthrough
		case BSIZEMOD == 2:
			buf[BSIZE-2] = byte(r >> 8)
			fallthrough
		case BSIZEMOD == 1:
			buf[BSIZE-1] = byte(r)
		}

		for i = 0; i < BSIZE7; i += 7 {
			r = src.Int63()
			buf[i] = byte(r)
			buf[i+1] = byte(r >> 8)
			buf[i+2] = byte(r >> 16)
			buf[i+3] = byte(r >> 24)
			buf[i+4] = byte(r >> 32)
			buf[i+5] = byte(r >> 40)
			buf[i+6] = byte(r >> 48)
		}

		out <- buf
	}
}

// xor a file with random data
func garble(f *os.File, in <-chan []byte, out chan<- bool) {
	var n, m int
	var err error
	data := make([]byte, BSIZE)
	var buf []byte
	pos := int64(0)
	for {
		// read
		n, err = f.Read(data)
		for n != BSIZE || err != nil {
			if err != nil && err != io.EOF {
				panic(err)
			}
			if n == 0 {
				close(out)
				for {
					<-in // sleep forever
				}
				return
			}
			m, err = f.Read(data[n:BSIZE])
			if m == 0 && err == io.EOF {
				// last partial block
				break
			}
			n += m
		}

		// xor with random data
		buf = <-in
		C.xor((*C.int64_t)(unsafe.Pointer(&data[0])), (*C.int64_t)(unsafe.Pointer(&buf[0])))
		out <- true // done with buf

		// write
		_, err = f.WriteAt(data[0:n], pos)
		if err != nil {
			panic(err)
		}
		pos += int64(n)
	}
}

// parse command line arguments
var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

func init() {
	flag.StringVar(&phrase, "phrase", "", "the Garble phrase, by default random")
	flag.Parse()

	narg = flag.NArg()

	if narg <= 0 {
		flag.Usage()
		os.Exit(1)
	}

	if phrase == "" {
		phrase = fmt.Sprintf("%016x", uint64(randomSeed()))
	}

	fmt.Println("Using phrase:", phrase)
}

// the main program...
func main() {
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Use available CPUs:
	if runtime.GOMAXPROCS(0) == 1 &&
		runtime.NumCPU() > 1 &&
		os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	// Open files:
	files := make([]*os.File, narg)
	writers := make([]chan []byte, narg)
	signals := make([]chan bool, narg)

	for i, arg := range flag.Args() {
		f, err := os.OpenFile(arg, os.O_RDWR, 0666)
		if err != nil {
			panic(err)
		}
		defer f.Close()
		files[i] = f
		writers[i] = make(chan []byte, MULTI)
		signals[i] = make(chan bool, MULTI)
		go garble(files[i], writers[i], signals[i])
	}

	// Allocate byte buffer pool:
	buffer := make([]byte, BSIZE*POOL)

	for i := 0; i < BSIZE*POOL; i += BSIZE {
		pool <- buffer[i : i+BSIZE]
	}

	// Initialize random sources:
	hash := sha512.New()
	sum := make([]byte, hash.Size())

	for i := 0; i < SOURCES; i++ {
		var seed, s int64
		var err error

		hash.Write([]byte(":garble:" + phrase))
		hash.Sum(sum[:0])

		buf := bytes.NewReader(sum)
		s = 0

		for err == nil {
			err = binary.Read(buf, binary.LittleEndian, &s)
			seed ^= s
		}

		src := rand.NewSource(seed)
		data[i] = make(chan []byte, MULTI)
		go randomBytes(src, data[i])
	}

	// Route data channels:
	go func(data []chan []byte, writers []chan []byte, signals []chan bool) {
		var buf []byte
		for {
			for _, r := range data {
				buf = <-r

				for i, w := range writers {
					if signals[i] != nil {
						w <- buf
					}
				}

				loop <- buf
			}
		}
	}(data, writers, signals)

	var buf []byte

	for narg > 0 {
		buf = <-loop

		for i, s := range signals {
			if s != nil {
				_, ok := <-s

				if !ok {
					signals[i] = nil
					narg--
				}
			}
		}

		pool <- buf
	}

	fmt.Println("All done!")
}
<commit_msg>prepare for stdin/stdout<commit_after>// garble produces pseudo random bytes based on a phrase
// and uses it to garble and ungarble files
package main

// #include <stdint.h>
// #define BSIZE 65536
//
// void xor(int64_t *a, int64_t *b) {
//     int i = BSIZE / 8;
//     while(i--) {
//         a[i] ^= b[i];
//     }
// }
import "C"
import "unsafe"

import (
	"bytes"
	cryptorand "crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"os"
	"runtime"
	"runtime/pprof"
	"time"
)

const (
	BSIZE    = C.BSIZE
	BSIZE7   = BSIZE - BSIZEMOD
	BSIZEMOD = BSIZE % 7
	MULTI    = 4
	SOURCES  = 8
	POOL     = SOURCES * (MULTI + 1)
)

var (
	narg   = int(0)
	phrase = ""
	pool   = make(chan []byte, POOL)
	loop   = make(chan []byte, POOL)
	data   = make([]chan []byte, SOURCES)
)

// randomSeed produces a int64 seed based on crypto/rand and time.
func randomSeed() int64 {
	var seed int64

	urandom := make([]byte, 8)
	cryptorand.Reader.Read(urandom)

	for key, value := range urandom {
		seed ^= (int64(value) ^ time.Now().UTC().UnixNano()) << (uint(key) * 8)
	}

	return seed
}

// randomBytes fills byte buffers with random data
func randomBytes(src rand.Source, out chan<- []byte) {
	var (
		r int64
		i = BSIZE
	)

	for buf, ok := <-pool; ok; buf, ok = <-pool {
		r = src.Int63()
		switch { // Go seems to eliminate impossible cases
		case BSIZEMOD == 6:
			buf[BSIZE-6] = byte(r >> 48)
			fallthrough
		case BSIZEMOD == 5:
			buf[BSIZE-5] = byte(r >> 32)
			fallthrough
		case BSIZEMOD == 4:
			buf[BSIZE-4] = byte(r >> 24)
			fallthrough
		case BSIZEMOD == 3:
			buf[BSIZE-3] = byte(r >> 16)
			fallthrough
		case BSIZEMOD == 2:
			buf[BSIZE-2] = byte(r >> 8)
			fallthrough
		case BSIZEMOD == 1:
			buf[BSIZE-1] = byte(r)
		}

		for i = 0; i < BSIZE7; i += 7 {
			r = src.Int63()
			buf[i] = byte(r)
			buf[i+1] = byte(r >> 8)
			buf[i+2] = byte(r >> 16)
			buf[i+3] = byte(r >> 24)
			buf[i+4] = byte(r >> 32)
			buf[i+5] = byte(r >> 40)
			buf[i+6] = byte(r >> 48)
		}

		out <- buf
	}
}

// xor a file with random data
func garble(fin *os.File, fout *os.File, in <-chan []byte, out chan<- bool) {
	var n, m int
	var err error
	data := make([]byte, BSIZE)
	var buf []byte

	for {
		// read
		n, err = fin.Read(data)
		for n != BSIZE || err != nil {
			if err != nil && err != io.EOF {
				panic(err)
			}
			if n == 0 {
				close(out)
				for {
					<-in // sleep forever
				}
				return
			}
			m, err = fin.Read(data[n:BSIZE])
			if m == 0 && err == io.EOF {
				// last partial block
				break
			}
			n += m
		}

		// xor with random data
		buf = <-in
		C.xor((*C.int64_t)(unsafe.Pointer(&data[0])), (*C.int64_t)(unsafe.Pointer(&buf[0])))
		out <- true // done with buf

		// write
		_, err = fout.Write(data[0:n])
		if err != nil {
			panic(err)
		}
	}
}

// parse command line arguments
var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

func init() {
	flag.StringVar(&phrase, "phrase", "", "the Garble phrase, by default random")
	flag.Parse()

	narg = flag.NArg()

	if narg <= 0 {
		flag.Usage()
		os.Exit(1)
	}

	if phrase == "" {
		phrase = fmt.Sprintf("%016x", uint64(randomSeed()))
	}

	fmt.Println("Using phrase:", phrase)
}

// the main program...
func main() {
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			panic(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Use available CPUs:
	if runtime.GOMAXPROCS(0) == 1 &&
		runtime.NumCPU() > 1 &&
		os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	// Open files:
	writers := make([]chan []byte, narg)
	signals := make([]chan bool, narg)

	for i, arg := range flag.Args() {
		var fd [2]*os.File
		for i, _ := range fd {
			f, err := os.OpenFile(arg, os.O_RDWR, 0666)
			if err != nil {
				panic(err)
			}
			defer f.Close()
			fd[i] = f
		}

		writers[i] = make(chan []byte, MULTI)
		signals[i] = make(chan bool, MULTI)
		go garble(fd[0], fd[1], writers[i], signals[i])
	}

	// Allocate byte buffer pool:
	buffer := make([]byte, BSIZE*POOL)

	for i := 0; i < BSIZE*POOL; i += BSIZE {
		pool <- buffer[i : i+BSIZE]
	}

	// Initialize random sources:
	hash := sha512.New()
	sum := make([]byte, hash.Size())

	for i := 0; i < SOURCES; i++ {
		var seed, s int64
		var err error

		hash.Write([]byte(":garble:" + phrase))
		hash.Sum(sum[:0])

		buf := bytes.NewReader(sum)
		s = 0

		for err == nil {
			err = binary.Read(buf, binary.LittleEndian, &s)
			seed ^= s
		}

		src := rand.NewSource(seed)
		data[i] = make(chan []byte, MULTI)
		go randomBytes(src, data[i])
	}

	// Route data channels:
	go func(data []chan []byte, writers []chan []byte, signals []chan bool) {
		var buf []byte
		for {
			for _, r := range data {
				buf = <-r

				for i, w := range writers {
					if signals[i] != nil {
						w <- buf
					}
				}

				loop <- buf
			}
		}
	}(data, writers, signals)

	var buf []byte

	for narg > 0 {
		buf = <-loop

		for i, s := range signals {
			if s != nil {
				_, ok := <-s

				if !ok {
					signals[i] = nil
					narg--
				}
			}
		}

		pool <- buf
	}

	fmt.Println("All done!")
}
<|endoftext|>


<commit_before>package rpcd

import (
	"encoding/gob"
	"errors"
	"flag"
	"fmt"
	"github.com/Symantec/Dominator/lib/format"
	"github.com/Symantec/Dominator/lib/fsutil"
	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/netspeed"
	"github.com/Symantec/Dominator/lib/objectcache"
	objectclient "github.com/Symantec/Dominator/lib/objectserver/client"
	"github.com/Symantec/Dominator/lib/rateio"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/sub"
	"io"
	"os"
	"path"
	"syscall"
	"time"
)

const filePerms = syscall.S_IRUSR | syscall.S_IWUSR | syscall.S_IRGRP

var (
	exitOnFetchFailure = flag.Bool("exitOnFetchFailure", false,
		"If true, exit if there are fetch failures. For debugging only")
)

func (t *rpcType) Fetch(conn *srpc.Conn) error {
	defer conn.Flush()
	var request sub.FetchRequest
	var response sub.FetchResponse
	decoder := gob.NewDecoder(conn)
	if err := decoder.Decode(&request); err != nil {
		_, err = conn.WriteString(err.Error() + "\n")
		return err
	}
	if err := t.fetch(request, &response); err != nil {
		_, err = conn.WriteString(err.Error() + "\n")
		return err
	}
	if _, err := conn.WriteString("\n"); err != nil {
		return err
	}
	return gob.NewEncoder(conn).Encode(response)
}

func (t *rpcType) fetch(request sub.FetchRequest,
	reply *sub.FetchResponse) error {
	if *readOnly {
		txt := "Fetch() rejected due to read-only mode"
		t.logger.Println(txt)
		return errors.New(txt)
	}
	t.rwLock.Lock()
	defer t.rwLock.Unlock()
	t.logger.Printf("Fetch(%s) %d objects\n",
		request.ServerAddress, len(request.Hashes))
	if t.fetchInProgress {
		t.logger.Println("Error: fetch already in progress")
		return errors.New("fetch already in progress")
	}
	if t.updateInProgress {
		t.logger.Println("Error: update in progress")
		return errors.New("update in progress")
	}
	t.fetchInProgress = true
	go func() {
		t.lastFetchError = t.doFetch(request)
		if t.lastFetchError != nil && *exitOnFetchFailure {
			os.Exit(1)
		}
	}()
	return nil
}

func (t *rpcType) doFetch(request sub.FetchRequest) error {
	defer t.clearFetchInProgress()
	objectServer := objectclient.NewObjectClient(request.ServerAddress)
	benchmark := false
	linkSpeed, haveLinkSpeed := netspeed.GetSpeedToAddress(
		request.ServerAddress)
	if !haveLinkSpeed && t.networkReaderContext.MaximumSpeed() < 1 {
		benchmark = enoughBytesForBenchmark(objectServer, request)
		if benchmark {
			objectServer.SetExclusiveGetObjects(true)
			t.logger.Println("Benchmarking network speed")
		}
	}
	objectsReader, err := objectServer.GetObjects(request.Hashes)
	if err != nil {
		t.logger.Printf("Error getting object reader:\t%s\n", err.Error())
		return err
	}
	defer objectsReader.Close()
	var totalLength uint64
	defer func() { t.rescanObjectCacheChannel <- true }()
	timeStart := time.Now()
	for _, hash := range request.Hashes {
		length, reader, err := objectsReader.NextObject()
		if err != nil {
			t.logger.Println(err)
			return err
		}
		r := io.Reader(reader)
		if haveLinkSpeed {
			if linkSpeed > 0 {
				r = rateio.NewReaderContext(linkSpeed,
					uint64(t.networkReaderContext.SpeedPercent()),
					&rateio.ReadMeasurer{}).NewReader(reader)
			}
		} else if !benchmark {
			r = t.networkReaderContext.NewReader(reader)
		}
		err = readOne(t.objectsDir, hash, length, r)
		reader.Close()
		if err != nil {
			t.logger.Println(err)
			return err
		}
		totalLength += length
	}
	duration := time.Since(timeStart)
	speed := uint64(float64(totalLength) / duration.Seconds())
	if benchmark {
		file, err := os.Create(t.netbenchFilename)
		if err == nil {
			fmt.Fprintf(file, "%d\n", speed)
			file.Close()
		}
		t.networkReaderContext.InitialiseMaximumSpeed(speed)
	}
	t.logger.Printf("Fetch() complete. Read: %s in %s (%s/s)\n",
		format.FormatBytes(totalLength), duration, format.FormatBytes(speed))
	return nil
}

func enoughBytesForBenchmark(objectServer *objectclient.ObjectClient,
	request sub.FetchRequest) bool {
	lengths, err := objectServer.CheckObjects(request.Hashes)
	if err != nil {
		return false
	}
	var totalLength uint64
	for _, length := range lengths {
		totalLength += length
	}
	if totalLength > 1024*1024*64 {
		return true
	}
	return false
}

func readOne(objectsDir string, hash hash.Hash, length uint64,
	reader io.Reader) error {
	filename := path.Join(objectsDir, objectcache.HashToFilename(hash))
	dirname := path.Dir(filename)
	if err := os.MkdirAll(dirname, syscall.S_IRWXU); err != nil {
		return err
	}
	return fsutil.CopyToFile(filename, filePerms, reader, length)
}

func (t *rpcType) clearFetchInProgress() {
	t.rwLock.Lock()
	defer t.rwLock.Unlock()
	t.fetchInProgress = false
}
<commit_msg>Log speed limit in Subd.Fetch() RPC method.<commit_after>package rpcd

import (
	"encoding/gob"
	"errors"
	"flag"
	"fmt"
	"github.com/Symantec/Dominator/lib/format"
	"github.com/Symantec/Dominator/lib/fsutil"
	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/netspeed"
	"github.com/Symantec/Dominator/lib/objectcache"
	objectclient "github.com/Symantec/Dominator/lib/objectserver/client"
	"github.com/Symantec/Dominator/lib/rateio"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/sub"
	"io"
	"os"
	"path"
	"syscall"
	"time"
)

const filePerms = syscall.S_IRUSR | syscall.S_IWUSR | syscall.S_IRGRP

var (
	exitOnFetchFailure = flag.Bool("exitOnFetchFailure", false,
		"If true, exit if there are fetch failures. For debugging only")
)

func (t *rpcType) Fetch(conn *srpc.Conn) error {
	defer conn.Flush()
	var request sub.FetchRequest
	var response sub.FetchResponse
	decoder := gob.NewDecoder(conn)
	if err := decoder.Decode(&request); err != nil {
		_, err = conn.WriteString(err.Error() + "\n")
		return err
	}
	if err := t.fetch(request, &response); err != nil {
		_, err = conn.WriteString(err.Error() + "\n")
		return err
	}
	if _, err := conn.WriteString("\n"); err != nil {
		return err
	}
	return gob.NewEncoder(conn).Encode(response)
}

func (t *rpcType) fetch(request sub.FetchRequest,
	reply *sub.FetchResponse) error {
	if *readOnly {
		txt := "Fetch() rejected due to read-only mode"
		t.logger.Println(txt)
		return errors.New(txt)
	}
	t.rwLock.Lock()
	defer t.rwLock.Unlock()
	if t.fetchInProgress {
		t.logger.Println("Error: fetch already in progress")
		return errors.New("fetch already in progress")
	}
	if t.updateInProgress {
		t.logger.Println("Error: update in progress")
		return errors.New("update in progress")
	}
	t.fetchInProgress = true
	go func() {
		t.lastFetchError = t.doFetch(request)
		if t.lastFetchError != nil && *exitOnFetchFailure {
			os.Exit(1)
		}
	}()
	return nil
}

func (t *rpcType) doFetch(request sub.FetchRequest) error {
	defer t.clearFetchInProgress()
	objectServer := objectclient.NewObjectClient(request.ServerAddress)
	benchmark := false
	linkSpeed, haveLinkSpeed := netspeed.GetSpeedToAddress(
		request.ServerAddress)
	if haveLinkSpeed {
		t.logFetch(request, linkSpeed)
	} else {
		if t.networkReaderContext.MaximumSpeed() < 1 {
			benchmark = enoughBytesForBenchmark(objectServer, request)
			if benchmark {
				objectServer.SetExclusiveGetObjects(true)
				t.logger.Printf("Fetch(%s) %d objects and benchmark speed\n",
					request.ServerAddress, len(request.Hashes))
			} else {
				t.logFetch(request, 0)
			}
		} else {
			t.logFetch(request, t.networkReaderContext.MaximumSpeed())
		}
	}
	objectsReader, err := objectServer.GetObjects(request.Hashes)
	if err != nil {
		t.logger.Printf("Error getting object reader:\t%s\n", err.Error())
		return err
	}
	defer objectsReader.Close()
	var totalLength uint64
	defer func() { t.rescanObjectCacheChannel <- true }()
	timeStart := time.Now()
	for _, hash := range request.Hashes {
		length, reader, err := objectsReader.NextObject()
		if err != nil {
			t.logger.Println(err)
			return err
		}
		r := io.Reader(reader)
		if haveLinkSpeed {
			if linkSpeed > 0 {
				r = rateio.NewReaderContext(linkSpeed,
					uint64(t.networkReaderContext.SpeedPercent()),
					&rateio.ReadMeasurer{}).NewReader(reader)
			}
		} else if !benchmark {
			r = t.networkReaderContext.NewReader(reader)
		}
		err = readOne(t.objectsDir, hash, length, r)
		reader.Close()
		if err != nil {
			t.logger.Println(err)
			return err
		}
		totalLength += length
	}
	duration := time.Since(timeStart)
	speed := uint64(float64(totalLength) / duration.Seconds())
	if benchmark {
		file, err := os.Create(t.netbenchFilename)
		if err == nil {
			fmt.Fprintf(file, "%d\n", speed)
			file.Close()
		}
		t.networkReaderContext.InitialiseMaximumSpeed(speed)
	}
	t.logger.Printf("Fetch() complete. Read: %s in %s (%s/s)\n",
		format.FormatBytes(totalLength), format.Duration(duration),
		format.FormatBytes(speed))
	return nil
}

func (t *rpcType) logFetch(request sub.FetchRequest, speed uint64) {
	speedString := "unlimited speed"
	if speed > 0 {
		speedString = format.FormatBytes(
			speed*uint64(t.networkReaderContext.SpeedPercent())/100) + "/s"
	}
	t.logger.Printf("Fetch(%s) %d objects at %s\n",
		request.ServerAddress, len(request.Hashes), speedString)
}

func enoughBytesForBenchmark(objectServer *objectclient.ObjectClient,
	request sub.FetchRequest) bool {
	lengths, err := objectServer.CheckObjects(request.Hashes)
	if err != nil {
		return false
	}
	var totalLength uint64
	for _, length := range lengths {
		totalLength += length
	}
	if totalLength > 1024*1024*64 {
		return true
	}
	return false
}

func readOne(objectsDir string, hash hash.Hash, length uint64,
	reader io.Reader) error {
	filename := path.Join(objectsDir, objectcache.HashToFilename(hash))
	dirname := path.Dir(filename)
	if err := os.MkdirAll(dirname, syscall.S_IRWXU); err != nil {
		return err
	}
	return fsutil.CopyToFile(filename, filePerms, reader, length)
}

func (t *rpcType) clearFetchInProgress() {
	t.rwLock.Lock()
	defer t.rwLock.Unlock()
	t.fetchInProgress = false
}
<|endoftext|>


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

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime/pprof"
	"strings"

	"github.com/blevesearch/bleve"
	_ "github.com/blevesearch/bleve/config"
)

var indexPath = flag.String("index", "", "index path")
var limit = flag.Int("limit", 10, "limit to first N results")
var skip = flag.Int("skip", 0, "skip the first N results")
var explain = flag.Bool("explain", false, "explain scores")
var includeHighlights = flag.Bool("highlight", true, "highlight matches")
var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
var repeat = flag.Int("repeat", 1, "repeat query n times")

func main() {

	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		err = pprof.StartCPUProfile(f)
		if err != nil {
			log.Fatal(err)
		}
		defer pprof.StopCPUProfile()
	}

	if *indexPath == "" {
		log.Fatal("Specify index to query")
	}

	if flag.NArg() < 1 {
		log.Fatal("Specify search query")
	}

	// open index
	index, err := bleve.Open(*indexPath)
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		cerr := index.Close()
		if cerr != nil {
			log.Fatalf("error closing index: %v", err)
		}
	}()

	for i := 0; i < *repeat; i++ {
		// build a search with the provided parameters
		queryString := strings.Join(flag.Args(), " ")
		query := bleve.NewQueryStringQuery(queryString)
		searchRequest := bleve.NewSearchRequestOptions(query, *limit, *skip, *explain)

		// enable highlights if requested
		if *includeHighlights {
			searchRequest.Highlight = bleve.NewHighlightWithStyle("ansi")
		}

		// execute the search
		searchResult, err := index.Search(searchRequest)
		if err != nil {
			log.Fatalf("search error: %v", err)
		}
		fmt.Println(searchResult)
	}
}
<commit_msg>support metrics through bleve query<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime/pprof"
	"strings"

	"github.com/blevesearch/bleve"
	_ "github.com/blevesearch/bleve/config"
	_ "github.com/blevesearch/bleve/index/store/metrics"
)

var indexPath = flag.String("index", "", "index path")
var limit = flag.Int("limit", 10, "limit to first N results")
var skip = flag.Int("skip", 0, "skip the first N results")
var explain = flag.Bool("explain", false, "explain scores")
var includeHighlights = flag.Bool("highlight", true, "highlight matches")
var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
var repeat = flag.Int("repeat", 1, "repeat query n times")

func main() {

	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		err = pprof.StartCPUProfile(f)
		if err != nil {
			log.Fatal(err)
		}
		defer pprof.StopCPUProfile()
	}

	if *indexPath == "" {
		log.Fatal("Specify index to query")
	}

	if flag.NArg() < 1 {
		log.Fatal("Specify search query")
	}

	// open index
	index, err := bleve.Open(*indexPath)
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		cerr := index.Close()
		if cerr != nil {
			log.Fatalf("error closing index: %v", err)
		}
	}()

	for i := 0; i < *repeat; i++ {
		// build a search with the provided parameters
		queryString := strings.Join(flag.Args(), " ")
		query := bleve.NewQueryStringQuery(queryString)
		searchRequest := bleve.NewSearchRequestOptions(query, *limit, *skip, *explain)

		// enable highlights if requested
		if *includeHighlights {
			searchRequest.Highlight = bleve.NewHighlightWithStyle("ansi")
		}

		// execute the search
		searchResult, err := index.Search(searchRequest)
		if err != nil {
			log.Fatalf("search error: %v", err)
		}
		fmt.Println(searchResult)
	}
}
<|endoftext|>