<commit_before>package dapperish

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"strconv"
	"strings"
	"sync"

	"github.com/opentracing/api-golang/opentracing"
	"github.com/opentracing/api-golang/opentracing/standardtracer"
)

// NewTracer returns a new dapperish Tracer instance.
func NewTracer(processName string) opentracing.Tracer {
	return standardtracer.New(
		NewTrivialRecorder(processName),
		NewTraceContextSource())
}

// TraceContext is an implementation of opentracing.TraceContext.
type TraceContext struct {
	// A probabilistically unique identifier for a [multi-span] trace.
	TraceID int64

	// A probabilistically unique identifier for a span.
	SpanID int64

	// Whether the trace is sampled.
	Sampled bool

	// `tagLock` protects the `traceAttrs` map, which in turn supports
	// `SetTraceAttribute` and `TraceAttribute`.
	tagLock    sync.RWMutex
	traceAttrs map[string]string
}

const (
	// Note that these strings are designed to be unchanged by the conversion
	// into standard HTTP headers (which messes with capitalization).
	fieldNameTraceID = "Traceid"
	fieldNameSpanID  = "Spanid"
	fieldNameSampled = "Sampled"
)

// NewChild complies with the opentracing.TraceContext interface.
func (d *TraceContext) NewChild() (opentracing.TraceContext, opentracing.Tags) {
	d.tagLock.RLock()
	newTags := make(map[string]string, len(d.traceTags))
	for k, v := range d.traceTags {
		newTags[k] = v
	}
	d.tagLock.RUnlock()

	return &TraceContext{
		TraceID:   d.TraceID,
		SpanID:    randomID(),
		Sampled:   d.Sampled,
		traceTags: newTags,
	}, opentracing.Tags{"parent_span_id": d.SpanID}
}

// SetTraceAttribute complies with the opentracing.TraceContext interface.
func (d *DapperishTraceContext) SetTraceAttribute(restrictedKey, val string) opentracing.TraceContext {
	canonicalKey, valid := opentracing.CanonicalizeTraceAttributeKey(restrictedKey)
	if !valid {
		panic(fmt.Errorf("Invalid key: %q", restrictedKey))
	}

	d.tagLock.Lock()
	defer d.tagLock.Unlock()

	d.traceAttrs[canonicalKey] = val
	return d
}

// TraceAttribute complies with the opentracing.TraceContext interface.
func (d *TraceContext) TraceAttribute(restrictedKey string) string {
	canonicalKey, valid := opentracing.CanonicalizeTraceAttributeKey(restrictedKey)
	if !valid {
		panic(fmt.Errorf("Invalid key: %q", restrictedKey))
	}

	d.tagLock.RLock()
	defer d.tagLock.RUnlock()

	return d.traceAttrs[canonicalKey]
}

// TraceContextSource is an implementation of
// opentracing.TraceContextSource.
type TraceContextSource struct{}

// NewTraceContextSource returns a dapperish opentracing.TraceContextSource
// implementation.
func NewTraceContextSource() *TraceContextSource {
	return &TraceContextSource{}
}

// NewRootTraceContext complies with the opentracing.TraceContextSource interface.
func (d *TraceContextSource) NewRootTraceContext() opentracing.TraceContext {
	return &DapperishTraceContext{
		TraceID:    randomID(),
		SpanID:     randomID(),
		Sampled:    randomID()%1024 == 0,
		traceAttrs: make(map[string]string),
	}
}

// NewChildTraceContext complies with the opentracing.TraceContextSource interface.
func (d *DapperishTraceContextSource) NewChildTraceContext(parent opentracing.TraceContext) (opentracing.TraceContext, opentracing.Tags) {
	dParent := parent.(*DapperishTraceContext)
	dParent.tagLock.RLock()
	newTags := make(map[string]string, len(dParent.traceAttrs))
	for k, v := range dParent.traceAttrs {
		newTags[k] = v
	}
	dParent.tagLock.RUnlock()

	return &DapperishTraceContext{
		TraceID:    dParent.TraceID,
		SpanID:     randomID(),
		Sampled:    dParent.Sampled,
		traceAttrs: newTags,
	}, opentracing.Tags{"parent_span_id": dParent.SpanID}
}

// MarshalTraceContextStringMap complies with the
// opentracing.TraceContextSource interface.
func (d *TraceContextSource) MarshalTraceContextStringMap(
	ctx opentracing.TraceContext,
) (contextIDMap map[string]string, tagsMap map[string]string) {
	dctx := ctx.(*TraceContext)
	contextIDMap = map[string]string{
		fieldNameTraceID: strconv.FormatInt(dctx.TraceID, 10),
		fieldNameSpanID:  strconv.FormatInt(dctx.SpanID, 10),
		fieldNameSampled: strconv.FormatBool(dctx.Sampled),
	}
	dctx.tagLock.RLock()
	tagsMap = make(map[string]string, len(dctx.traceAttrs))
	for k, v := range dctx.traceAttrs {
		tagsMap[k] = v
	}
	dctx.tagLock.RUnlock()
	return contextIDMap, tagsMap
}

// UnmarshalTraceContextStringMap complies with the
// opentracing.TraceContextSource interface.
func (d *TraceContextSource) UnmarshalTraceContextStringMap(
	contextIDMap map[string]string,
	tagsMap map[string]string,
) (opentracing.TraceContext, error) {
	requiredFieldCount := 0
	var traceID, spanID int64
	var sampled bool
	var err error
	for k, v := range contextIDMap {
		switch k {
		case fieldNameTraceID:
			traceID, err = strconv.ParseInt(v, 10, 64)
			if err != nil {
				return nil, err
			}
			requiredFieldCount++
		case fieldNameSpanID:
			spanID, err = strconv.ParseInt(v, 10, 64)
			if err != nil {
				return nil, err
			}
			requiredFieldCount++
		case fieldNameSampled:
			sampled, err = strconv.ParseBool(v)
			if err != nil {
				return nil, err
			}
			requiredFieldCount++
		default:
			return nil, fmt.Errorf("Unknown contextIDMap field: %v", k)
		}
	}
	if requiredFieldCount < 3 {
		return nil, fmt.Errorf("Only found %v of 3 required fields", requiredFieldCount)
	}

	lowercaseTagsMap := make(map[string]string, len(tagsMap))
	for k, v := range tagsMap {
		lowercaseTagsMap[strings.ToLower(k)] = v
	}

	return &TraceContext{
		TraceID:    traceID,
		SpanID:     spanID,
		Sampled:    sampled,
		traceAttrs: lowercaseTagsMap,
	}, nil
}

// MarshalTraceContextBinary complies with the opentracing.TraceContextSource
// interface.
func (d *TraceContextSource) MarshalTraceContextBinary(ctx opentracing.TraceContext) (contextID []byte, traceAttrs []byte) {
	dtc := ctx.(*DapperishTraceContext)
	var err error
	buf := new(bytes.Buffer)
	err = binary.Write(buf, binary.BigEndian, dtc.TraceID)
	if err != nil {
		panic(err)
	}
	err = binary.Write(buf, binary.BigEndian, dtc.SpanID)
	if err != nil {
		panic(err)
	}
	var sampledByte byte
	if dtc.Sampled {
		sampledByte = 1
	}
	err = binary.Write(buf, binary.BigEndian, sampledByte)
	if err != nil {
		panic(err)
	}
	// XXX: support tags
	return buf.Bytes(), []byte{}
}

// UnmarshalTraceContextBinary complies with the opentracing.TraceContextSource
// interface.
func (d *TraceContextSource) UnmarshalTraceContextBinary(
	contextID []byte,
	traceAttrs []byte,
) (opentracing.TraceContext, error) {
	var err error
	reader := bytes.NewReader(contextID)
	var traceID, spanID int64
	var sampledByte byte

	err = binary.Read(reader, binary.BigEndian, &traceID)
	if err != nil {
		return nil, err
	}
	err = binary.Read(reader, binary.BigEndian, &spanID)
	if err != nil {
		return nil, err
	}
	err = binary.Read(reader, binary.BigEndian, &sampledByte)
	if err != nil {
		return nil, err
	}
	// XXX: support tags
	return &TraceContext{
		TraceID:    traceID,
		SpanID:     spanID,
		Sampled:    sampledByte != 0,
		traceAttrs: make(map[string]string),
	}, nil
}
<commit_msg>fix build after rebase<commit_after>package dapperish

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"strconv"
	"strings"
	"sync"

	"github.com/opentracing/api-golang/opentracing"
	"github.com/opentracing/api-golang/opentracing/standardtracer"
)

// NewTracer returns a new dapperish Tracer instance.
func NewTracer(processName string) opentracing.Tracer {
	return standardtracer.New(
		NewTrivialRecorder(processName),
		NewTraceContextSource())
}

// TraceContext is an implementation of opentracing.TraceContext.
type TraceContext struct {
	// A probabilistically unique identifier for a [multi-span] trace.
	TraceID int64

	// A probabilistically unique identifier for a span.
	SpanID int64

	// Whether the trace is sampled.
	Sampled bool

	// `tagLock` protects the `traceAttrs` map, which in turn supports
	// `SetTraceAttribute` and `TraceAttribute`.
	tagLock    sync.RWMutex
	traceAttrs map[string]string
}

const (
	// Note that these strings are designed to be unchanged by the conversion
	// into standard HTTP headers (which messes with capitalization).
	fieldNameTraceID = "Traceid"
	fieldNameSpanID  = "Spanid"
	fieldNameSampled = "Sampled"
)

// SetTraceAttribute complies with the opentracing.TraceContext interface.
func (d *TraceContext) SetTraceAttribute(restrictedKey, val string) opentracing.TraceContext {
	canonicalKey, valid := opentracing.CanonicalizeTraceAttributeKey(restrictedKey)
	if !valid {
		panic(fmt.Errorf("Invalid key: %q", restrictedKey))
	}

	d.tagLock.Lock()
	defer d.tagLock.Unlock()

	d.traceAttrs[canonicalKey] = val
	return d
}

// TraceAttribute complies with the opentracing.TraceContext interface.
func (d *TraceContext) TraceAttribute(restrictedKey string) string {
	canonicalKey, valid := opentracing.CanonicalizeTraceAttributeKey(restrictedKey)
	if !valid {
		panic(fmt.Errorf("Invalid key: %q", restrictedKey))
	}

	d.tagLock.RLock()
	defer d.tagLock.RUnlock()

	return d.traceAttrs[canonicalKey]
}

// TraceContextSource is an implementation of
// opentracing.TraceContextSource.
type TraceContextSource struct{}

// NewTraceContextSource returns a dapperish opentracing.TraceContextSource
// implementation.
func NewTraceContextSource() *TraceContextSource {
	return &TraceContextSource{}
}

// NewRootTraceContext complies with the opentracing.TraceContextSource interface.
func (d *TraceContextSource) NewRootTraceContext() opentracing.TraceContext {
	return &TraceContext{
		TraceID:    randomID(),
		SpanID:     randomID(),
		Sampled:    randomID()%1024 == 0,
		traceAttrs: make(map[string]string),
	}
}

// NewChildTraceContext complies with the opentracing.TraceContextSource interface.
func (d *TraceContextSource) NewChildTraceContext(
	parent opentracing.TraceContext,
) (opentracing.TraceContext, opentracing.Tags) {
	dParent := parent.(*TraceContext)
	dParent.tagLock.RLock()
	newTags := make(map[string]string, len(dParent.traceAttrs))
	for k, v := range dParent.traceAttrs {
		newTags[k] = v
	}
	dParent.tagLock.RUnlock()

	return &TraceContext{
		TraceID:    dParent.TraceID,
		SpanID:     randomID(),
		Sampled:    dParent.Sampled,
		traceAttrs: newTags,
	}, opentracing.Tags{"parent_span_id": dParent.SpanID}
}

// MarshalTraceContextStringMap complies with the
// opentracing.TraceContextSource interface.
func (d *TraceContextSource) MarshalTraceContextStringMap(
	ctx opentracing.TraceContext,
) (contextIDMap map[string]string, tagsMap map[string]string) {
	dctx := ctx.(*TraceContext)
	contextIDMap = map[string]string{
		fieldNameTraceID: strconv.FormatInt(dctx.TraceID, 10),
		fieldNameSpanID:  strconv.FormatInt(dctx.SpanID, 10),
		fieldNameSampled: strconv.FormatBool(dctx.Sampled),
	}
	dctx.tagLock.RLock()
	tagsMap = make(map[string]string, len(dctx.traceAttrs))
	for k, v := range dctx.traceAttrs {
		tagsMap[k] = v
	}
	dctx.tagLock.RUnlock()
	return contextIDMap, tagsMap
}

// UnmarshalTraceContextStringMap complies with the
// opentracing.TraceContextSource interface.
func (d *TraceContextSource) UnmarshalTraceContextStringMap(
	contextIDMap map[string]string,
	tagsMap map[string]string,
) (opentracing.TraceContext, error) {
	requiredFieldCount := 0
	var traceID, spanID int64
	var sampled bool
	var err error
	for k, v := range contextIDMap {
		switch k {
		case fieldNameTraceID:
			traceID, err = strconv.ParseInt(v, 10, 64)
			if err != nil {
				return nil, err
			}
			requiredFieldCount++
		case fieldNameSpanID:
			spanID, err = strconv.ParseInt(v, 10, 64)
			if err != nil {
				return nil, err
			}
			requiredFieldCount++
		case fieldNameSampled:
			sampled, err = strconv.ParseBool(v)
			if err != nil {
				return nil, err
			}
			requiredFieldCount++
		default:
			return nil, fmt.Errorf("Unknown contextIDMap field: %v", k)
		}
	}
	if requiredFieldCount < 3 {
		return nil, fmt.Errorf("Only found %v of 3 required fields", requiredFieldCount)
	}

	lowercaseTagsMap := make(map[string]string, len(tagsMap))
	for k, v := range tagsMap {
		lowercaseTagsMap[strings.ToLower(k)] = v
	}

	return &TraceContext{
		TraceID:    traceID,
		SpanID:     spanID,
		Sampled:    sampled,
		traceAttrs: lowercaseTagsMap,
	}, nil
}

// MarshalTraceContextBinary complies with the opentracing.TraceContextSource
// interface.
func (d *TraceContextSource) MarshalTraceContextBinary(ctx opentracing.TraceContext) (contextID []byte, traceAttrs []byte) {
	dtc := ctx.(*TraceContext)
	var err error
	buf := new(bytes.Buffer)
	err = binary.Write(buf, binary.BigEndian, dtc.TraceID)
	if err != nil {
		panic(err)
	}
	err = binary.Write(buf, binary.BigEndian, dtc.SpanID)
	if err != nil {
		panic(err)
	}
	var sampledByte byte
	if dtc.Sampled {
		sampledByte = 1
	}
	err = binary.Write(buf, binary.BigEndian, sampledByte)
	if err != nil {
		panic(err)
	}
	// XXX: support tags
	return buf.Bytes(), []byte{}
}

// UnmarshalTraceContextBinary complies with the opentracing.TraceContextSource
// interface.
func (d *TraceContextSource) UnmarshalTraceContextBinary(
	contextID []byte,
	traceAttrs []byte,
) (opentracing.TraceContext, error) {
	var err error
	reader := bytes.NewReader(contextID)
	var traceID, spanID int64
	var sampledByte byte

	err = binary.Read(reader, binary.BigEndian, &traceID)
	if err != nil {
		return nil, err
	}
	err = binary.Read(reader, binary.BigEndian, &spanID)
	if err != nil {
		return nil, err
	}
	err = binary.Read(reader, binary.BigEndian, &sampledByte)
	if err != nil {
		return nil, err
	}
	// XXX: support tags
	return &TraceContext{
		TraceID:    traceID,
		SpanID:     spanID,
		Sampled:    sampledByte != 0,
		traceAttrs: make(map[string]string),
	}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"os"
	"path/filepath"
	"github.com/songtianyi/wechat-go/wxweb"
	"github.com/songtianyi/rrframework/logs"
)

func main() {
	// get web server root path
	cur_dir := filepath.Dir(os.Args[0])
	public_dir := filepath.Join(cur_dir, "public")

	// create session and put qrcode image to webroot
	session, err := wxweb.CreateWebSessionWithPath(nil, nil, public_dir)
	if err != nil {
		logs.Error(err)
		return
	}

	// serve and wait for wechat msg
	go session.LoginAndServe(true)

	// serve http
	http.ListenAndServe(":8080", http.FileServer(http.Dir(public_dir)))

	// then visit http://target:8080/public/  + session.QrcodePath
}
<commit_msg>gofmt code<commit_after>package main

import (
	"github.com/songtianyi/rrframework/logs"
	"github.com/songtianyi/wechat-go/wxweb"
	"net/http"
	"os"
	"path/filepath"
)

func main() {
	// get web server root path
	cur_dir := filepath.Dir(os.Args[0])
	public_dir := filepath.Join(cur_dir, "public")

	// create session and put qrcode image to webroot
	session, err := wxweb.CreateWebSessionWithPath(nil, nil, public_dir)
	if err != nil {
		logs.Error(err)
		return
	}

	// serve and wait for wechat msg
	go session.LoginAndServe(true)

	// serve http
	http.ListenAndServe(":8080", http.FileServer(http.Dir(public_dir)))

	// then visit http://target:8080/public/  + session.QrcodePath
}
<|endoftext|>


<commit_before>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

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 smart

import (
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core"

	"github.com/intelsdi-x/snap-plugin-utilities/config"
)

const (
	PluginName = "smart-disk"
	version    = 7
	pluginType = plugin.CollectorPluginType

	nsVendor = "intel"
	nsClass  = "disk"
	nsType   = "smart"
)

var (
	//procPath source of data for metrics
	procPath = "/proc"
	//devPath source of data for metrics
	devPath = "/dev"

	namespace_prefix = []string{nsVendor, nsClass}
	namespace_suffix = []string{nsType}

	sysUtilProvider SysutilProvider
)

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(
		PluginName,
		version,
		pluginType,
		[]string{plugin.SnapGOBContentType},
		[]string{plugin.SnapGOBContentType},
		plugin.ConcurrencyCount(1),
	)
}

func NewSmartCollector() *SmartCollector {
	logger := log.New()
	imutex := new(sync.Mutex)
	return &SmartCollector{
		logger:           logger,
		initializedMutex: imutex,
		proc_path:        procPath,
		dev_path:         devPath,
	}
}

func makeName(device, metric string) []string {
	splited := strings.Split(metric, "/")

	name := []string{}
	name = append(name, namespace_prefix...)
	name = append(name, device)
	name = append(name, namespace_suffix...)
	name = append(name, splited...)

	return name
}

func parseName(namespace []string) (disk, attribute string) {
	disk = namespace[len(namespace_prefix)]
	smart_namespace := namespace[len(namespace_prefix)+len(namespace_suffix)+1:]
	attribute = strings.Join(smart_namespace, "/")
	return
}

func validateName(namespace []string) bool {
	for i, v := range namespace_prefix {
		if namespace[i] != v {
			return false
		}
	}

	offset := len(namespace_prefix) + 1
	for i, v := range namespace_suffix {
		if namespace[offset+i] != v {
			return false
		}
	}

	return true
}

// Function to check properness of configuration parameters
// and set plugin attribute accordingly
func (sc *SmartCollector) setProcDevPath(cfg interface{}) error {
	sc.initializedMutex.Lock()
	defer sc.initializedMutex.Unlock()
	procPath, err := config.GetConfigItem(cfg, "proc_path")
	if err == nil && len(procPath.(string)) > 0 {
		procPathStats, err := os.Stat(procPath.(string))
		if err != nil {
			return err
		}
		if !procPathStats.IsDir() {
			return errors.New(fmt.Sprintf("%s is not a directory", procPath.(string)))
		}
		sc.proc_path = procPath.(string)
	}
	devPath, err := config.GetConfigItem(cfg, "dev_path")
	if err == nil && len(devPath.(string)) > 0 {
		devPathStats, err := os.Stat(devPath.(string))
		if err != nil {
			return err
		}
		if !devPathStats.IsDir() {
			return errors.New(fmt.Sprintf("%s is not a directory", devPath.(string)))
		}
		sc.dev_path = devPath.(string)
	}
	if sysUtilProvider == nil {
		sysUtilProvider = NewSysutilProvider(sc.proc_path, sc.dev_path)
	}
	sc.initialized = true
	return nil
}

type SmartCollector struct {
	initialized      bool
	initializedMutex *sync.Mutex
	logger           *log.Logger
	proc_path        string
	dev_path         string
}

type smartResults map[string]interface{}

// CollectMetrics returns metrics from smart
func (sc *SmartCollector) CollectMetrics(mts []plugin.MetricType) ([]plugin.MetricType, error) {
	if !sc.initialized {
		if err := sc.setProcDevPath(mts[0]); err != nil {
			return nil, err
		}
	}

	buffered_results := map[string]smartResults{}

	results := make([]plugin.MetricType, len(mts))
	errs := make([]string, 0)

	collected := false

	t := time.Now()

	for i, mt := range mts {
		namespace := mt.Namespace().Strings()
		results[i] = plugin.MetricType{
			Namespace_: mt.Namespace(),
			Timestamp_: t,
		}

		if !validateName(namespace) {
			errs = append(errs, fmt.Sprintf("%s is not valid metric", mt.Namespace().String()))
			continue
		}
		disk, attribute_path := parseName(namespace)
		buffered, ok := buffered_results[disk]
		if !ok {
			values, err := ReadSmartData(disk, sysUtilProvider)
			if err != nil {
				return nil, err
			}
			buffered = values.GetAttributes()
			buffered_results[disk] = buffered
		}

		attribute, ok := buffered[attribute_path]

		if !ok {
			errs = append(errs, "Unknown attribute "+attribute_path)
		} else {
			collected = true
			results[i].Data_ = attribute
		}
	}

	errsStr := strings.Join(errs, "; ")
	if collected {
		if len(errs) > 0 {
			log.Printf("Data collected but error(s) occured: %v", errsStr)
		}
		return results, nil
	} else {
		return nil, errors.New(errsStr)
	}
}

// GetMetricTypes returns the metric types exposed by smart
func (sc *SmartCollector) GetMetricTypes(cfg plugin.ConfigType) ([]plugin.MetricType, error) {
	if !sc.initialized {
		if err := sc.setProcDevPath(cfg); err != nil {
			return nil, err
		}
	}
	smart_metrics := ListAllKeys()
	devices, err := sysUtilProvider.ListDevices()
	if err != nil {
		return nil, err
	}
	mts := make([]plugin.MetricType, 0, len(smart_metrics)*len(devices))

	for _, device := range devices {
		for _, metric := range smart_metrics {
			path := makeName(device, metric)
			mts = append(mts, plugin.MetricType{Namespace_: core.NewNamespace(path...)})
		}
	}

	return mts, nil
}

//GetConfigPolicy returns a ConfigPolicy
func (p *SmartCollector) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	cp := cpolicy.New()
	rule, _ := cpolicy.NewStringRule("proc_path", false, "/proc")
	node := cpolicy.NewPolicyNode()
	node.Add(rule)
	cp.Add([]string{nsVendor, nsClass, nsType}, node)
	rule, _ = cpolicy.NewStringRule("dev_path", false, "/dev")
	node.Add(rule)
	return cp, nil
}
<commit_msg>Proper implementation of locked calls<commit_after>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

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 smart

import (
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core"

	"github.com/intelsdi-x/snap-plugin-utilities/config"
)

const (
	PluginName = "smart-disk"
	version    = 7
	pluginType = plugin.CollectorPluginType

	nsVendor = "intel"
	nsClass  = "disk"
	nsType   = "smart"
)

var (
	//procPath source of data for metrics
	procPath = "/proc"
	//devPath source of data for metrics
	devPath = "/dev"

	namespace_prefix = []string{nsVendor, nsClass}
	namespace_suffix = []string{nsType}

	sysUtilProvider SysutilProvider
)

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(
		PluginName,
		version,
		pluginType,
		[]string{plugin.SnapGOBContentType},
		[]string{plugin.SnapGOBContentType},
		plugin.ConcurrencyCount(1),
	)
}

func NewSmartCollector() *SmartCollector {
	logger := log.New()
	imutex := new(sync.Mutex)
	return &SmartCollector{
		logger:           logger,
		initializedMutex: imutex,
		proc_path:        procPath,
		dev_path:         devPath,
	}
}

func makeName(device, metric string) []string {
	splited := strings.Split(metric, "/")

	name := []string{}
	name = append(name, namespace_prefix...)
	name = append(name, device)
	name = append(name, namespace_suffix...)
	name = append(name, splited...)

	return name
}

func parseName(namespace []string) (disk, attribute string) {
	disk = namespace[len(namespace_prefix)]
	smart_namespace := namespace[len(namespace_prefix)+len(namespace_suffix)+1:]
	attribute = strings.Join(smart_namespace, "/")
	return
}

func validateName(namespace []string) bool {
	for i, v := range namespace_prefix {
		if namespace[i] != v {
			return false
		}
	}

	offset := len(namespace_prefix) + 1
	for i, v := range namespace_suffix {
		if namespace[offset+i] != v {
			return false
		}
	}

	return true
}

// Function to check properness of configuration parameters
// and set plugin attribute accordingly
func (sc *SmartCollector) setProcDevPath(cfg interface{}) error {
	sc.initializedMutex.Lock()
	defer sc.initializedMutex.Unlock()
	if sc.initialized {
		return nil
	}
	procPath, err := config.GetConfigItem(cfg, "proc_path")
	if err == nil && len(procPath.(string)) > 0 {
		procPathStats, err := os.Stat(procPath.(string))
		if err != nil {
			return err
		}
		if !procPathStats.IsDir() {
			return errors.New(fmt.Sprintf("%s is not a directory", procPath.(string)))
		}
		sc.proc_path = procPath.(string)
	}
	devPath, err := config.GetConfigItem(cfg, "dev_path")
	if err == nil && len(devPath.(string)) > 0 {
		devPathStats, err := os.Stat(devPath.(string))
		if err != nil {
			return err
		}
		if !devPathStats.IsDir() {
			return errors.New(fmt.Sprintf("%s is not a directory", devPath.(string)))
		}
		sc.dev_path = devPath.(string)
	}
	if sysUtilProvider == nil {
		sysUtilProvider = NewSysutilProvider(sc.proc_path, sc.dev_path)
	}
	sc.initialized = true
	return nil
}

type SmartCollector struct {
	initialized      bool
	initializedMutex *sync.Mutex
	logger           *log.Logger
	proc_path        string
	dev_path         string
}

type smartResults map[string]interface{}

// CollectMetrics returns metrics from smart
func (sc *SmartCollector) CollectMetrics(mts []plugin.MetricType) ([]plugin.MetricType, error) {
	if err := sc.setProcDevPath(mts[0]); err != nil {
		return nil, err
	}

	buffered_results := map[string]smartResults{}

	results := make([]plugin.MetricType, len(mts))
	errs := make([]string, 0)

	collected := false

	t := time.Now()

	for i, mt := range mts {
		namespace := mt.Namespace().Strings()
		results[i] = plugin.MetricType{
			Namespace_: mt.Namespace(),
			Timestamp_: t,
		}

		if !validateName(namespace) {
			errs = append(errs, fmt.Sprintf("%s is not valid metric", mt.Namespace().String()))
			continue
		}
		disk, attribute_path := parseName(namespace)
		buffered, ok := buffered_results[disk]
		if !ok {
			values, err := ReadSmartData(disk, sysUtilProvider)
			if err != nil {
				return nil, err
			}
			buffered = values.GetAttributes()
			buffered_results[disk] = buffered
		}

		attribute, ok := buffered[attribute_path]

		if !ok {
			errs = append(errs, "Unknown attribute "+attribute_path)
		} else {
			collected = true
			results[i].Data_ = attribute
		}
	}

	errsStr := strings.Join(errs, "; ")
	if collected {
		if len(errs) > 0 {
			log.Printf("Data collected but error(s) occured: %v", errsStr)
		}
		return results, nil
	} else {
		return nil, errors.New(errsStr)
	}
}

// GetMetricTypes returns the metric types exposed by smart
func (sc *SmartCollector) GetMetricTypes(cfg plugin.ConfigType) ([]plugin.MetricType, error) {
	if err := sc.setProcDevPath(cfg); err != nil {
		return nil, err
	}
	smart_metrics := ListAllKeys()
	devices, err := sysUtilProvider.ListDevices()
	if err != nil {
		return nil, err
	}
	mts := make([]plugin.MetricType, 0, len(smart_metrics)*len(devices))

	for _, device := range devices {
		for _, metric := range smart_metrics {
			path := makeName(device, metric)
			mts = append(mts, plugin.MetricType{Namespace_: core.NewNamespace(path...)})
		}
	}

	return mts, nil
}

//GetConfigPolicy returns a ConfigPolicy
func (p *SmartCollector) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	cp := cpolicy.New()
	rule, _ := cpolicy.NewStringRule("proc_path", false, "/proc")
	node := cpolicy.NewPolicyNode()
	node.Add(rule)
	cp.Add([]string{nsVendor, nsClass, nsType}, node)
	rule, _ = cpolicy.NewStringRule("dev_path", false, "/dev")
	node.Add(rule)
	return cp, nil
}
<|endoftext|>


<commit_before>package commandline

import "strings"

func Parse(s string) *Command {
	ss := strings.Split(s, " ")
	if ss[0] == "" {
		return nil
	}
	return &Command{
		Name: ss[0],
		Args: ss[1:],
	}
}

type Command struct {
	Name string
	Args []string
}

type scanner struct {
	src []byte
}
<commit_msg>Add scanner.scan<commit_after>package commandline

import "strings"

func Parse(s string) *Command {
	ss := strings.Split(s, " ")
	if ss[0] == "" {
		return nil
	}
	return &Command{
		Name: ss[0],
		Args: ss[1:],
	}
}

type Command struct {
	Name string
	Args []string
}

type scanner struct {
	src []byte
}

func (s scanner) scan() {
}
<|endoftext|>


<commit_before>package main

import "fmt"

// Data is a generic struct with two generic types
type Data[K comparable, V any] struct {
	Key   K
	Value V
}

// getSliceData returns a slice of Data[K, V]
func getSliceData[K comparable, V any](m map[K]V) []Data[K, V] {
	// define a slice with Data type passing K, V type parameters
	d := make([]Data[K, V], len(m))
	i := 0
	for k, v := range m {
		// creating value using new keyword
		newData := Data[K, V]{}
		newData.Key = k
		newData.Value = v
		d[i] = newData
		i++
	}
	return d
}

func main() {
	data := getSliceData(map[string]string{
		"go":  "Go programming language",
		"rs":  "Rust programming language",
		"zig": "Zig programming language",
	})
	for _, v := range data {
		fmt.Printf("Key: %s Value:%s\n", v.Key, v.Value)
	}

	data1 := getSliceData(map[int]string{
		1: "Go programming language",
		2: "Rust programming language",
		3: "Zig programming language",
	})
	for _, v := range data1 {
		fmt.Printf("Key: %d Value:%s\n", v.Key, v.Value)
	}
}
<commit_msg>Added generic type example<commit_after>package main

import "fmt"

// Data is a generic struct with two generic types
type Data[K comparable, V any] struct {
	Key   K
	Value V
}

// getSliceData returns a slice of Data[K, V]
func getSliceData[K comparable, V any](m map[K]V) []Data[K, V] {
	// define a slice with Data type passing K, V type parameters
	d := make([]Data[K, V], len(m))
	i := 0
	for k, v := range m {
		// creating value of generic type of Data
		newData := Data[K, V]{}
		newData.Key = k
		newData.Value = v
		d[i] = newData
		i++
	}
	return d
}

func main() {
	data := getSliceData(map[string]string{
		"go":  "Go programming language",
		"rs":  "Rust programming language",
		"zig": "Zig programming language",
	})
	for _, v := range data {
		fmt.Printf("Key: %s Value:%s\n", v.Key, v.Value)
	}

	data1 := getSliceData(map[int]string{
		1: "Go programming language",
		2: "Rust programming language",
		3: "Zig programming language",
	})
	for _, v := range data1 {
		fmt.Printf("Key: %d Value:%s\n", v.Key, v.Value)
	}
}
<|endoftext|>


<commit_before>// Copyright 2020-2022 Buf Technologies, 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 httpserver

import (
	"context"
	"crypto/tls"
	"net"
	"net/http"
	"time"

	"github.com/go-chi/chi/v5"
	"go.uber.org/zap"
)

const (
	// DefaultShutdownTimeout is the default shutdown timeout.
	DefaultShutdownTimeout = 10 * time.Second
	// DefaultReadHeaderTimeout is the default read header timeout.
	DefaultReadHeaderTimeout = 30 * time.Second
	// DefaultIdleTimeout is the amount of time an HTTP/2 connection can be idle.
	DefaultIdleTimeout = 3 * time.Minute
)

// Mapper initializes a Router.
type Mapper interface {
	Map(chi.Router) error
}

// MapperFunc is a function for a Mapper
type MapperFunc func(chi.Router) error

// Map implements Mapper.
func (m MapperFunc) Map(router chi.Router) error {
	return m(router)
}

// NewHTTPHandlerMapper returns a new Mapper for the http.Handler.
func NewHTTPHandlerMapper(
	handler http.Handler,
	options ...HTTPHandlerMapperOption,
) Mapper {
	return newHTTPHandlerMapper(handler, options...)
}

// HTTPHandlerMapperOption is an option for a new HTTPHandlerMapper.
type HTTPHandlerMapperOption func(*httpHandlerMapper)

// HTTPHandlerMapperWithPrefix returns a new HTTPHandlerMapperOption that uses the prefix.
func HTTPHandlerMapperWithPrefix(prefix string) HTTPHandlerMapperOption {
	return func(httpHandlerMapper *httpHandlerMapper) {
		httpHandlerMapper.prefix = prefix
	}
}

// Runner is a runner.
type Runner interface {
	// Run runs the router.
	//
	// Blocking.
	// The runner is cancelled when the input context is cancelled.
	// The listener is closed upon return.
	//
	// Response write errors are logged. Response write errors can be ignored.
	//
	// Can be called multiple times, resulting in different runs.
	Run(ctx context.Context, listener net.Listener, mappers ...Mapper) error
}

// NewRunner returns a new Runner.
func NewRunner(logger *zap.Logger, options ...RunnerOption) Runner {
	return newRunner(logger, options...)
}

// RunnerOption is an option for a new Runner.
type RunnerOption func(*runner)

// RunnerWithShutdownTimeout returns a new RunnerOption that uses the given shutdown timeout.
//
// The default is to use DefaultShutdownTimeout.
// If shutdownTimeout is 0, no graceful shutdown will be performed.
func RunnerWithShutdownTimeout(shutdownTimeout time.Duration) RunnerOption {
	return func(runner *runner) {
		runner.shutdownTimeout = shutdownTimeout
	}
}

// RunnerWithReadHeaderTimeout returns a new RunnerOption that uses the given read header timeout.
//
// The default is to use DefaultReadHeaderTimeout.
// If readHeaderTimeout is 0, no read header timeout will be used.
func RunnerWithReadHeaderTimeout(readHeaderTimeout time.Duration) RunnerOption {
	return func(runner *runner) {
		runner.readHeaderTimeout = readHeaderTimeout
	}
}

// RunnerWithTLSConfig returns a new RunnerOption that uses the given tls.Config.
//
// The default is to use no TLS.
func RunnerWithTLSConfig(tlsConfig *tls.Config) RunnerOption {
	return func(runner *runner) {
		runner.tlsConfig = tlsConfig
	}
}

// RunnerWithMiddlewares returns a new RunnerOption that use middlewares when the Runner Run.
func RunnerWithMiddlewares(middlewares ...func(http.Handler) http.Handler) RunnerOption {
	return func(runner *runner) {
		runner.middlewares = append(runner.middlewares, middlewares...)
	}
}

// RunnerWithWalkFunc returns a new RunnerOption that runs chi.Walk to walk the router
// after all middlewares and routes have been mounted, but before the server is started.
func RunnerWithWalkFunc(walkFunc chi.WalkFunc) RunnerOption {
	return func(runner *runner) {
		runner.walkFunc = walkFunc
	}
}

// RunnerWithMaxBodySize returns a new RunnerOption that sets the max size of
// incoming request body.
//
// The default is to not limit body size.
func RunnerWithMaxBodySize(maxBodySize int64) RunnerOption {
	return func(runner *runner) {
		runner.maxBodySize = maxBodySize
	}
}

// RunnerWithHealth returns a new RunnerOption that turns a health check endpoint on at /health.
//
// The default is to not turn on health.
func RunnerWithHealth() RunnerOption {
	return func(runner *runner) {
		runner.health = true
	}
}
<commit_msg>add MapperWithMiddlewares helper (#1330)<commit_after>// Copyright 2020-2022 Buf Technologies, 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 httpserver

import (
	"context"
	"crypto/tls"
	"net"
	"net/http"
	"time"

	"github.com/go-chi/chi/v5"
	"go.uber.org/zap"
)

const (
	// DefaultShutdownTimeout is the default shutdown timeout.
	DefaultShutdownTimeout = 10 * time.Second
	// DefaultReadHeaderTimeout is the default read header timeout.
	DefaultReadHeaderTimeout = 30 * time.Second
	// DefaultIdleTimeout is the amount of time an HTTP/2 connection can be idle.
	DefaultIdleTimeout = 3 * time.Minute
)

// Mapper initializes a Router.
type Mapper interface {
	Map(chi.Router) error
}

// MapperFunc is a function for a Mapper
type MapperFunc func(chi.Router) error

// Map implements Mapper.
func (m MapperFunc) Map(router chi.Router) error {
	return m(router)
}

// NewHTTPHandlerMapper returns a new Mapper for the http.Handler.
func NewHTTPHandlerMapper(
	handler http.Handler,
	options ...HTTPHandlerMapperOption,
) Mapper {
	return newHTTPHandlerMapper(handler, options...)
}

// HTTPHandlerMapperOption is an option for a new HTTPHandlerMapper.
type HTTPHandlerMapperOption func(*httpHandlerMapper)

// HTTPHandlerMapperWithPrefix returns a new HTTPHandlerMapperOption that uses the prefix.
func HTTPHandlerMapperWithPrefix(prefix string) HTTPHandlerMapperOption {
	return func(httpHandlerMapper *httpHandlerMapper) {
		httpHandlerMapper.prefix = prefix
	}
}

// MapperWithMiddlewares rewrites the Map function of the Mapper to call the provided middlewares
// inside a chi Group before mapping
func MapperWithMiddlewares(mapper Mapper, middlewares chi.Middlewares) Mapper {
	return MapperFunc(func(parent chi.Router) error {
		var err error
		parent.Group(func(r chi.Router) {
			r.Use(middlewares...)
			err = mapper.Map(r)
		})
		return err
	})
}

// Runner is a runner.
type Runner interface {
	// Run runs the router.
	//
	// Blocking.
	// The runner is cancelled when the input context is cancelled.
	// The listener is closed upon return.
	//
	// Response write errors are logged. Response write errors can be ignored.
	//
	// Can be called multiple times, resulting in different runs.
	Run(ctx context.Context, listener net.Listener, mappers ...Mapper) error
}

// NewRunner returns a new Runner.
func NewRunner(logger *zap.Logger, options ...RunnerOption) Runner {
	return newRunner(logger, options...)
}

// RunnerOption is an option for a new Runner.
type RunnerOption func(*runner)

// RunnerWithShutdownTimeout returns a new RunnerOption that uses the given shutdown timeout.
//
// The default is to use DefaultShutdownTimeout.
// If shutdownTimeout is 0, no graceful shutdown will be performed.
func RunnerWithShutdownTimeout(shutdownTimeout time.Duration) RunnerOption {
	return func(runner *runner) {
		runner.shutdownTimeout = shutdownTimeout
	}
}

// RunnerWithReadHeaderTimeout returns a new RunnerOption that uses the given read header timeout.
//
// The default is to use DefaultReadHeaderTimeout.
// If readHeaderTimeout is 0, no read header timeout will be used.
func RunnerWithReadHeaderTimeout(readHeaderTimeout time.Duration) RunnerOption {
	return func(runner *runner) {
		runner.readHeaderTimeout = readHeaderTimeout
	}
}

// RunnerWithTLSConfig returns a new RunnerOption that uses the given tls.Config.
//
// The default is to use no TLS.
func RunnerWithTLSConfig(tlsConfig *tls.Config) RunnerOption {
	return func(runner *runner) {
		runner.tlsConfig = tlsConfig
	}
}

// RunnerWithMiddlewares returns a new RunnerOption that use middlewares when the Runner Run.
func RunnerWithMiddlewares(middlewares ...func(http.Handler) http.Handler) RunnerOption {
	return func(runner *runner) {
		runner.middlewares = append(runner.middlewares, middlewares...)
	}
}

// RunnerWithWalkFunc returns a new RunnerOption that runs chi.Walk to walk the router
// after all middlewares and routes have been mounted, but before the server is started.
func RunnerWithWalkFunc(walkFunc chi.WalkFunc) RunnerOption {
	return func(runner *runner) {
		runner.walkFunc = walkFunc
	}
}

// RunnerWithMaxBodySize returns a new RunnerOption that sets the max size of
// incoming request body.
//
// The default is to not limit body size.
func RunnerWithMaxBodySize(maxBodySize int64) RunnerOption {
	return func(runner *runner) {
		runner.maxBodySize = maxBodySize
	}
}

// RunnerWithHealth returns a new RunnerOption that turns a health check endpoint on at /health.
//
// The default is to not turn on health.
func RunnerWithHealth() RunnerOption {
	return func(runner *runner) {
		runner.health = true
	}
}
<|endoftext|>


<commit_before>package storage

import (
	"database/sql"
	"encoding/base64"
	"errors"
	"fmt"
	log "github.com/golang/glog"
	"github.com/jrupac/goliath/models"
	// PostgreSQL driver support
	_ "github.com/lib/pq"
	"strconv"
	"strings"
	"time"
)

const (
	DIALECT = "postgres"
	FOLDER_TABLE = "Folder"
	FOLDER_CHILDREN_TABLE = "FolderChildren"
	FEED_TABLE = "Feed"
	ARTICLE_TABLE = "Article"
	USER_TABLE = "UserTable"
	MAX_FETCHED_ROWS = 10000
)

type Database struct {
	db *sql.DB
}

func Open(dbPath string) (*Database, error) {
	db := new(Database)

	d, err := sql.Open(DIALECT, dbPath)
	if err != nil {
		return nil, err
	}
	if err = d.Ping(); err != nil {
		return nil, err
	}

	db.db = d
	return db, nil
}

func (d *Database) Close() error {
	return d.db.Close()
}

/*******************************************************************************
 * Insertion/deletion methods
 ******************************************************************************/

func (d *Database) InsertArticle(a models.Article) error {
	var articleId int64
	var count int
	err := d.db.QueryRow(
		`SELECT COUNT(*) FROM ` + ARTICLE_TABLE + ` WHERE hash = $1`, a.Hash()).Scan(&count)
	if err != nil {
		return err
	}
	if count > 0 {
		log.V(2).Infof("Duplicate article entry, skipping (hash): %s", a.Hash())
		return nil
	}

	err = d.db.QueryRow(
		`INSERT INTO ` + ARTICLE_TABLE + `
		(feed, folder, hash, title, summary, content, link, read, date, retrieved)
		VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING id`,
		a.FeedId, a.FolderId, a.Hash(), a.Title, a.Summary, a.Content, a.Link, a.Read, a.Date, a.Retrieved).Scan(&articleId)
	if err != nil {
		return err
	}
	a.Id = articleId
	return nil
}

func (d *Database) InsertFavicon(feedId int64, mime string, img []byte) error {
	var count int
	err := d.db.QueryRow(
		`SELECT COUNT(*) FROM ` + FEED_TABLE + ` WHERE id = $1`, feedId).Scan(&count)
	if err != nil {
		return err
	}
	if count == 0 {
		return fmt.Errorf("Original feed not in table: %d", feedId)
	}

	// Convert to a base64 encoded string before inserting
	h := base64.StdEncoding.EncodeToString(img)

	_, err = d.db.Query(
		`UPDATE ` + FEED_TABLE + ` SET favicon = $1, mime = $2 WHERE id = $3`,
		h, mime, feedId)
	return err
}

func (d *Database) InsertUser(u models.User) error {
	_, err := d.db.Query(
		`INSERT INTO ` + USER_TABLE + `(username, key) VALUES($1, $2)`, u.Username, u.Key)
	return err
}

func (d *Database) DeleteArticles(minTimestamp time.Time) (int, error) {
	count := 0
	rows, err := d.db.Query(
		`DELETE FROM ` + ARTICLE_TABLE + ` WHERE read AND (retrieved IS NULL OR retrieved < $1) RETURNING id`, minTimestamp)
	if err != nil {
		return count, err
	}
	defer rows.Close()

	var id int64
	for rows.Next() {
		rows.Scan(&id)
		count += 1
	}

	return count, err
}

/*******************************************************************************
 * Modification methods
 ******************************************************************************/

func (d *Database) MarkArticle(id int64, status string) error {
	state, err := parseState(status)
	if err != nil {
		return err
	}

	_, err = d.db.Query(`UPDATE ` + ARTICLE_TABLE + ` SET read = $1 WHERE id = $2`, state, id)
	return err
}

func (d *Database) MarkFeed(id int64, status string) error {
	state, err := parseState(status)
	if err != nil {
		return err
	}

	_, err = d.db.Query(`UPDATE ` + ARTICLE_TABLE + ` SET read = $1 WHERE feed = $2`, state, id)
	return err
}

func (d *Database) MarkFolder(id int64, status string) error {
	state, err := parseState(status)
	if err != nil {
		return err
	}

	_, err = d.db.Query(`UPDATE ` + ARTICLE_TABLE + ` SET read = $1 WHERE folder = $2`, state, id)
	children, err := d.GetFolderChildren(id)
	if err != nil {
		return err
	}
	for _, c := range children {
		if err := d.MarkFolder(c, status); err != nil {
			return err
		}
	}
	return err
}

func (d *Database) GetFolderChildren(id int64) ([]int64, error) {
	children := []int64{}
	rows, err := d.db.Query(`SELECT child FROM ` + FOLDER_CHILDREN_TABLE + ` WHERE parent = $1`, id)
	if err != nil {
		return children, err
	}
	defer rows.Close()

	var childId int64
	for rows.Next() {
		if err = rows.Scan(&childId); err != nil {
			return children, err
		}
		children = append(children, childId)
	}
	return children, err
}

func (d *Database) GetAllFolders() ([]models.Folder, error) {
	folders := []models.Folder{}
	rows, err := d.db.Query(`SELECT id, name FROM ` + FOLDER_TABLE)
	if err != nil {
		return folders, err
	}
	defer rows.Close()

	for rows.Next() {
		f := models.Folder{}
		if err = rows.Scan(&f.Id, &f.Name); err != nil {
			return folders, err
		}
		folders = append(folders, f)
	}

	return folders, err
}

func (d *Database) GetAllFeeds() ([]models.Feed, error) {
	feeds := []models.Feed{}
	rows, err := d.db.Query(`SELECT id, folder, title, description, url FROM ` + FEED_TABLE)
	if err != nil {
		return feeds, err
	}
	defer rows.Close()

	for rows.Next() {
		f := models.Feed{}
		if err = rows.Scan(&f.Id, &f.FolderId, &f.Title, &f.Description, &f.Url); err != nil {
			return feeds, err
		}
		feeds = append(feeds, f)
	}

	return feeds, err
}

func (d *Database) GetFeedsPerFolder() (map[int64]string, error) {
	agg := map[int64][]string{}
	resp := map[int64]string{}

	// CockroachDB doesn't have a concat-with-separator aggregation function
	rows, err := d.db.Query(`SELECT folder, id FROM ` + FEED_TABLE)
	if err != nil {
		return resp, err
	}
	defer rows.Close()

	var folderId, feedId int64
	for rows.Next() {
		if err = rows.Scan(&folderId, &feedId); err != nil {
			return resp, err
		}
		agg[folderId] = append(agg[folderId], strconv.FormatInt(feedId, 10))
	}

	for k, v := range agg {
		resp[k] = strings.Join(v, ",")
	}
	return resp, err
}

func (d *Database) GetAllFavicons() (map[int64]string, error) {
	favicons := map[int64]string{}
	rows, err := d.db.Query(
		`SELECT id, mime, favicon FROM ` + FEED_TABLE + ` WHERE favicon IS NOT NULL`)
	if err != nil {
		return favicons, err
	}
	defer rows.Close()

	var id int64
	var mime string
	var favicon string
	for rows.Next() {
		if err = rows.Scan(&id, &mime, &favicon); err != nil {
			return favicons, err
		}
		favicons[id] = fmt.Sprintf("%s;base64,%s", mime, favicon)
	}
	return favicons, err
}

func (d *Database) GetUnreadArticles(limit int, sinceId int64) ([]models.Article, error) {
	articles := []models.Article{}
	var rows *sql.Rows
	var err error

	if limit == -1 {
		limit = MAX_FETCHED_ROWS
	}
	if sinceId == -1 {
		sinceId = 0
	}

	rows, err = d.db.Query(
		`SELECT id, feed, folder, title, summary, content, link, date FROM ` + ARTICLE_TABLE + `
		WHERE NOT read AND id > $1 ORDER BY id LIMIT $2`, sinceId, limit)

	if err != nil {
		return articles, err
	}
	defer rows.Close()

	for rows.Next() {
		a := models.Article{}
		if err = rows.Scan(
			&a.Id, &a.FeedId, &a.FolderId, &a.Title, &a.Summary, &a.Content, &a.Link, &a.Date); err != nil {
			return articles, err
		}
		articles = append(articles, a)
	}
	return articles, err
}

func (d *Database) GetUserByKey(key string) (models.User, error) {
	var u models.User
	err := d.db.QueryRow(
		`SELECT username, key FROM ` + USER_TABLE + ` WHERE key = $1`, key).Scan(
		&u.Username, &u.Key)
	if !u.Valid() {
		return models.User{}, errors.New("Could not find user.")
	}
	return u, err
}

/*******************************************************************************
 * Import methods
 ******************************************************************************/

func (d *Database) ImportOpml(opml *models.Opml) error {
	root := opml.Folders
	// TODO: Remove extra read after https://github.com/cockroachdb/cockroach/issues/6637 is closed.
	_, err := d.db.Query(
		`INSERT INTO ` + FOLDER_TABLE + `(name) VALUES($1) ON CONFLICT(name) DO NOTHING`, root.Name)
	if err != nil {
		return err
	}
	err = d.db.QueryRow(
		`SELECT id FROM ` + FOLDER_TABLE + ` WHERE name = $1`, root.Name).Scan(&root.Id)
	if err != nil {
		return err
	}
	return d.importChildren(root)
}

func (d *Database) importChildren(parent models.Folder) error {
	var err error
	for _, f := range parent.Feed {
		// TODO: Remove extra read after https://github.com/cockroachdb/cockroach/issues/6637 is closed.
		_, err = d.db.Query(
			`INSERT INTO ` + FEED_TABLE + `(folder, hash, title, description, url)
			VALUES($1, $2, $3, $4, $5) ON CONFLICT(hash) DO NOTHING`,
			parent.Id, f.Hash(), f.Title, f.Description, f.Url)
		if err != nil {
			return err
		}
		err = d.db.QueryRow(
			`SELECT id FROM ` + FEED_TABLE + ` WHERE hash = $1`, f.Hash()).Scan(&f.Id)
		if err != nil {
			return err
		}
		f.FolderId = parent.Id
	}

	for _, child := range parent.Folders {
		// TODO: Remove extra read after https://github.com/cockroachdb/cockroach/issues/6637 is closed.
		_, err := d.db.Query(
			`INSERT INTO ` + FOLDER_TABLE + `(name) VALUES($1) ON CONFLICT(name) DO NOTHING`, child.Name)
		if err != nil {
			return err
		}
		err = d.db.QueryRow(
			`SELECT id FROM ` + FOLDER_TABLE + ` WHERE name = $1`, child.Name).Scan(&child.Id)
		if err != nil {
			return err
		}

		_, err = d.db.Query(
			`UPSERT INTO ` + FOLDER_CHILDREN_TABLE + `(parent, child) VALUES($1, $2)`, parent.Id, child.Id)
		if err != nil {
			return err
		}
		if err = d.importChildren(child); err != nil {
			return err
		}
	}
	return err
}

func parseState(status string) (bool, error) {
	var state bool
	switch status {
	case "saved", "unsaved":
		// Perhaps add support for saving articles in the future.
		return false, errors.New("Unsupported status: %s" + status)
	case "read":
		state = true
	case "unread":
		state = false
	}
	return state, nil
}
<commit_msg>core/db: Don't set limit on SQL queries.<commit_after>package storage

import (
	"database/sql"
	"encoding/base64"
	"errors"
	"fmt"
	log "github.com/golang/glog"
	"github.com/jrupac/goliath/models"
	// PostgreSQL driver support
	_ "github.com/lib/pq"
	"strconv"
	"strings"
	"time"
)

const (
	DIALECT = "postgres"
	FOLDER_TABLE = "Folder"
	FOLDER_CHILDREN_TABLE = "FolderChildren"
	FEED_TABLE = "Feed"
	ARTICLE_TABLE = "Article"
	USER_TABLE = "UserTable"
)

type Database struct {
	db *sql.DB
}

func Open(dbPath string) (*Database, error) {
	db := new(Database)

	d, err := sql.Open(DIALECT, dbPath)
	if err != nil {
		return nil, err
	}
	if err = d.Ping(); err != nil {
		return nil, err
	}

	db.db = d
	return db, nil
}

func (d *Database) Close() error {
	return d.db.Close()
}

/*******************************************************************************
 * Insertion/deletion methods
 ******************************************************************************/

func (d *Database) InsertArticle(a models.Article) error {
	var articleId int64
	var count int
	err := d.db.QueryRow(
		`SELECT COUNT(*) FROM ` + ARTICLE_TABLE + ` WHERE hash = $1`, a.Hash()).Scan(&count)
	if err != nil {
		return err
	}
	if count > 0 {
		log.V(2).Infof("Duplicate article entry, skipping (hash): %s", a.Hash())
		return nil
	}

	err = d.db.QueryRow(
		`INSERT INTO ` + ARTICLE_TABLE + `
		(feed, folder, hash, title, summary, content, link, read, date, retrieved)
		VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING id`,
		a.FeedId, a.FolderId, a.Hash(), a.Title, a.Summary, a.Content, a.Link, a.Read, a.Date, a.Retrieved).Scan(&articleId)
	if err != nil {
		return err
	}
	a.Id = articleId
	return nil
}

func (d *Database) InsertFavicon(feedId int64, mime string, img []byte) error {
	var count int
	err := d.db.QueryRow(
		`SELECT COUNT(*) FROM ` + FEED_TABLE + ` WHERE id = $1`, feedId).Scan(&count)
	if err != nil {
		return err
	}
	if count == 0 {
		return fmt.Errorf("Original feed not in table: %d", feedId)
	}

	// Convert to a base64 encoded string before inserting
	h := base64.StdEncoding.EncodeToString(img)

	_, err = d.db.Query(
		`UPDATE ` + FEED_TABLE + ` SET favicon = $1, mime = $2 WHERE id = $3`,
		h, mime, feedId)
	return err
}

func (d *Database) InsertUser(u models.User) error {
	_, err := d.db.Query(
		`INSERT INTO ` + USER_TABLE + `(username, key) VALUES($1, $2)`, u.Username, u.Key)
	return err
}

func (d *Database) DeleteArticles(minTimestamp time.Time) (int, error) {
	count := 0
	rows, err := d.db.Query(
		`DELETE FROM ` + ARTICLE_TABLE + ` WHERE read AND (retrieved IS NULL OR retrieved < $1) RETURNING id`, minTimestamp)
	if err != nil {
		return count, err
	}
	defer rows.Close()

	var id int64
	for rows.Next() {
		rows.Scan(&id)
		count += 1
	}

	return count, err
}

/*******************************************************************************
 * Modification methods
 ******************************************************************************/

func (d *Database) MarkArticle(id int64, status string) error {
	state, err := parseState(status)
	if err != nil {
		return err
	}

	_, err = d.db.Query(`UPDATE ` + ARTICLE_TABLE + ` SET read = $1 WHERE id = $2`, state, id)
	return err
}

func (d *Database) MarkFeed(id int64, status string) error {
	state, err := parseState(status)
	if err != nil {
		return err
	}

	_, err = d.db.Query(`UPDATE ` + ARTICLE_TABLE + ` SET read = $1 WHERE feed = $2`, state, id)
	return err
}

func (d *Database) MarkFolder(id int64, status string) error {
	state, err := parseState(status)
	if err != nil {
		return err
	}

	_, err = d.db.Query(`UPDATE ` + ARTICLE_TABLE + ` SET read = $1 WHERE folder = $2`, state, id)
	children, err := d.GetFolderChildren(id)
	if err != nil {
		return err
	}
	for _, c := range children {
		if err := d.MarkFolder(c, status); err != nil {
			return err
		}
	}
	return err
}

func (d *Database) GetFolderChildren(id int64) ([]int64, error) {
	children := []int64{}
	rows, err := d.db.Query(`SELECT child FROM ` + FOLDER_CHILDREN_TABLE + ` WHERE parent = $1`, id)
	if err != nil {
		return children, err
	}
	defer rows.Close()

	var childId int64
	for rows.Next() {
		if err = rows.Scan(&childId); err != nil {
			return children, err
		}
		children = append(children, childId)
	}
	return children, err
}

func (d *Database) GetAllFolders() ([]models.Folder, error) {
	folders := []models.Folder{}
	rows, err := d.db.Query(`SELECT id, name FROM ` + FOLDER_TABLE)
	if err != nil {
		return folders, err
	}
	defer rows.Close()

	for rows.Next() {
		f := models.Folder{}
		if err = rows.Scan(&f.Id, &f.Name); err != nil {
			return folders, err
		}
		folders = append(folders, f)
	}

	return folders, err
}

func (d *Database) GetAllFeeds() ([]models.Feed, error) {
	feeds := []models.Feed{}
	rows, err := d.db.Query(`SELECT id, folder, title, description, url FROM ` + FEED_TABLE)
	if err != nil {
		return feeds, err
	}
	defer rows.Close()

	for rows.Next() {
		f := models.Feed{}
		if err = rows.Scan(&f.Id, &f.FolderId, &f.Title, &f.Description, &f.Url); err != nil {
			return feeds, err
		}
		feeds = append(feeds, f)
	}

	return feeds, err
}

func (d *Database) GetFeedsPerFolder() (map[int64]string, error) {
	agg := map[int64][]string{}
	resp := map[int64]string{}

	// CockroachDB doesn't have a concat-with-separator aggregation function
	rows, err := d.db.Query(`SELECT folder, id FROM ` + FEED_TABLE)
	if err != nil {
		return resp, err
	}
	defer rows.Close()

	var folderId, feedId int64
	for rows.Next() {
		if err = rows.Scan(&folderId, &feedId); err != nil {
			return resp, err
		}
		agg[folderId] = append(agg[folderId], strconv.FormatInt(feedId, 10))
	}

	for k, v := range agg {
		resp[k] = strings.Join(v, ",")
	}
	return resp, err
}

func (d *Database) GetAllFavicons() (map[int64]string, error) {
	favicons := map[int64]string{}
	rows, err := d.db.Query(
		`SELECT id, mime, favicon FROM ` + FEED_TABLE + ` WHERE favicon IS NOT NULL`)
	if err != nil {
		return favicons, err
	}
	defer rows.Close()

	var id int64
	var mime string
	var favicon string
	for rows.Next() {
		if err = rows.Scan(&id, &mime, &favicon); err != nil {
			return favicons, err
		}
		favicons[id] = fmt.Sprintf("%s;base64,%s", mime, favicon)
	}
	return favicons, err
}

func (d *Database) GetUnreadArticles(limit int, sinceId int64) ([]models.Article, error) {
	articles := []models.Article{}
	var rows *sql.Rows
	var err error

	if sinceId == -1 {
		sinceId = 0
	}

	rows, err = d.db.Query(
		`SELECT id, feed, folder, title, summary, content, link, date FROM ` + ARTICLE_TABLE + `
		WHERE NOT read AND id > $1 ORDER BY id`, sinceId)

	if err != nil {
		return articles, err
	}
	defer rows.Close()

	count := 0
	for rows.Next() {
		if count == limit {
			break
		}

		a := models.Article{}
		if err = rows.Scan(
			&a.Id, &a.FeedId, &a.FolderId, &a.Title, &a.Summary, &a.Content, &a.Link, &a.Date); err != nil {
			return articles, err
		}
		articles = append(articles, a)
		count += 1
	}
	return articles, err
}

func (d *Database) GetUserByKey(key string) (models.User, error) {
	var u models.User
	err := d.db.QueryRow(
		`SELECT username, key FROM ` + USER_TABLE + ` WHERE key = $1`, key).Scan(
		&u.Username, &u.Key)
	if !u.Valid() {
		return models.User{}, errors.New("Could not find user.")
	}
	return u, err
}

/*******************************************************************************
 * Import methods
 ******************************************************************************/

func (d *Database) ImportOpml(opml *models.Opml) error {
	root := opml.Folders
	// TODO: Remove extra read after https://github.com/cockroachdb/cockroach/issues/6637 is closed.
	_, err := d.db.Query(
		`INSERT INTO ` + FOLDER_TABLE + `(name) VALUES($1) ON CONFLICT(name) DO NOTHING`, root.Name)
	if err != nil {
		return err
	}
	err = d.db.QueryRow(
		`SELECT id FROM ` + FOLDER_TABLE + ` WHERE name = $1`, root.Name).Scan(&root.Id)
	if err != nil {
		return err
	}
	return d.importChildren(root)
}

func (d *Database) importChildren(parent models.Folder) error {
	var err error
	for _, f := range parent.Feed {
		// TODO: Remove extra read after https://github.com/cockroachdb/cockroach/issues/6637 is closed.
		_, err = d.db.Query(
			`INSERT INTO ` + FEED_TABLE + `(folder, hash, title, description, url)
			VALUES($1, $2, $3, $4, $5) ON CONFLICT(hash) DO NOTHING`,
			parent.Id, f.Hash(), f.Title, f.Description, f.Url)
		if err != nil {
			return err
		}
		err = d.db.QueryRow(
			`SELECT id FROM ` + FEED_TABLE + ` WHERE hash = $1`, f.Hash()).Scan(&f.Id)
		if err != nil {
			return err
		}
		f.FolderId = parent.Id
	}

	for _, child := range parent.Folders {
		// TODO: Remove extra read after https://github.com/cockroachdb/cockroach/issues/6637 is closed.
		_, err := d.db.Query(
			`INSERT INTO ` + FOLDER_TABLE + `(name) VALUES($1) ON CONFLICT(name) DO NOTHING`, child.Name)
		if err != nil {
			return err
		}
		err = d.db.QueryRow(
			`SELECT id FROM ` + FOLDER_TABLE + ` WHERE name = $1`, child.Name).Scan(&child.Id)
		if err != nil {
			return err
		}

		_, err = d.db.Query(
			`UPSERT INTO ` + FOLDER_CHILDREN_TABLE + `(parent, child) VALUES($1, $2)`, parent.Id, child.Id)
		if err != nil {
			return err
		}
		if err = d.importChildren(child); err != nil {
			return err
		}
	}
	return err
}

func parseState(status string) (bool, error) {
	var state bool
	switch status {
	case "saved", "unsaved":
		// Perhaps add support for saving articles in the future.
		return false, errors.New("Unsupported status: %s" + status)
	case "read":
		state = true
	case "unread":
		state = false
	}
	return state, nil
}
<|endoftext|>


<commit_before>package horsefeather

import (
	"github.com/SeanDolphin/horsefeather/test"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
)

var encodedKeys = []string{
	"aghnbGlicmFyeXIMCxIGUGVyc29uGAEM",
	"aghnbGlicmFyeXIdCxIFR3JhcGgiEmdyYXBoOjctZGF5LWFjdGl2ZQw",
	"aghnbGlicmFyeXIhCxIJV29yZEluZGV4GIChPgwLEglXb3JkSW5kZXgYiQgM",
}

var _ = Describe("Store", func() {
	var ctx context.Context
	var cache = test.NewCache()
	var store = test.NewStore()

	var keys []*datastore.Key
	for _, encodedKey := range encodedKeys {
		key, _ := datastore.DecodeKey(encodedKey)
		keys = append(keys, key)
	}

	It("should panic when memcache is not present", func() {
		Expect(func() { mc(context.Background()) }).To(Panic())
	})

	It("should panic when datastore is not present", func() {
		Expect(func() { ds(context.Background()) }).To(Panic())
	})

	Context("when caching context are present", func() {
		BeforeEach(func() {
			ctx = AddMemcache(context.Background(), cache)
			ctx = AddDatastore(ctx, store)
		})

		AfterEach(func() {
			store.Clear()
			cache.Clear()
		})

		Context("when dealing with single keys", func() {
			var data = "bob"
			Context("when deleteing", func() {
				It("should delete the key", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")

						err = Delete(ctx, key)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})

				It("should error when there is nothing to delete", func() {
					for _, key := range keys {
						Expect(Delete(ctx, key)).To(HaveOccurred())
					}
				})

				It("should delete only to the memcache when memcache only is allowed.", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}

					ctx = OnlyMemcache(ctx, true)
					for _, key := range keys {
						err := Delete(ctx, key)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})

				It("should delete only to the datastore when datastore only is allowed.", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}

					ctx = OnlyDatastore(ctx, true)
					for _, key := range keys {
						err := Delete(ctx, key)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})
			})

			Context("when getting", func() {
				It("should get the key", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})

				It("should get the key when it is not present in memcache", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						cache.Delete(ctx, key)

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})

				It("should get the key when it is not present in the datastore", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						store.Delete(ctx, key)

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})

				It("should error when trying to get something that does not exist", func() {
					for _, key := range keys {
						var result string
						Expect(Get(ctx, key, &result)).To(HaveOccurred())
					}
				})

				It("should work when only memcache is allowed", func() {
					ctx = OnlyMemcache(ctx, true)
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})

				It("should work when only datastore is allowed", func() {
					ctx = OnlyDatastore(ctx, true)
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})
			})

			Context("whenning putting a key", func() {
				It("should put the data", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})

				It("should error when trying to save nil", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, nil)
						Expect(err).To(HaveOccurred())
					}
				})

				It("should put only to the memcache when memcache only is allowed.", func() {
					ctx = OnlyMemcache(ctx, true)
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})

				It("should put only to the datastore when datastore only is allowed.", func() {
					ctx = OnlyDatastore(ctx, true)
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})
			})
		})

		Context("when dealing with multiple keys", func() {
			var data = []string{"t1", "t2", "t3"}
			Context("when deleting", func() {
				It("should delete all the keys", func() {
					_, err := PutMulti(ctx, keys, data)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "store")
					}

					err = DeleteMulti(ctx, keys)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {
						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})

				It("should return an error deleting an empty", func() {
					Expect(DeleteMulti(ctx, keys)).To(HaveOccurred())
				})

				It("should delete only to the memcache when memcache only is allowed.", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}

					ctx = OnlyMemcache(ctx, true)
					err := DeleteMulti(ctx, keys)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {

						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})

				It("should delete only to the datastore when datastore only is allowed.", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}

					ctx = OnlyDatastore(ctx, true)
					err := DeleteMulti(ctx, keys)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {

						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})
			})

			Context("when prefetching results", func() {
				It("should load data on multiple items", func() {
					PutMulti(ctx, keys, data)
					ctx = Prefetch(ctx, keys, data)
					store.Clear()
					cache.Clear()
					var result = make([]string, len(keys))
					err := GetMulti(ctx, keys, result)
					Expect(err).ToNot(HaveOccurred())
					for i := 0; i < len(keys); i++ {
						Expect(result[i]).To(Equal(data[i]))
					}
				})

				It("should load data for a single Item", func() {
					PutMulti(ctx, keys, data)
					ctx = Prefetch(ctx, keys, data)
					store.Clear()
					cache.Clear()

					for i, key := range keys {
						var result string
						err := Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data[i]))
					}
				})
			})

			Context("when getting", func() {
				It("should get all the keys", func() {
					_, err := PutMulti(ctx, keys, &data)
					Expect(err).ToNot(HaveOccurred())

					var results = make([]string, len(keys))
					err = GetMulti(ctx, keys, &results)
					Expect(err).ToNot(HaveOccurred())
					Expect(results).To(HaveLen(len(data)))
					for i, result := range results {
						Expect(result).To(Equal(data[i]))
					}
				})

				It("should get all the data even if the memcache is empty", func() {
					_, err := PutMulti(ctx, keys, &data)
					Expect(err).ToNot(HaveOccurred())

					for _, key := range keys {
						cache.Delete(ctx, key)
					}

					var results = make([]string, len(keys))
					err = GetMulti(ctx, keys, &results)
					Expect(err).ToNot(HaveOccurred())
					Expect(results).To(HaveLen(len(data)))
					for i, result := range results {
						Expect(result).To(Equal(data[i]))
					}
				})

				It("should get all that even if it is split between the store and memcache", func() {
					_, err := PutMulti(ctx, keys, &data)
					Expect(err).ToNot(HaveOccurred())

					for i, key := range keys {
						if i%2 == 0 {
							cache.Delete(ctx, key)
						} else {
							store.Delete(ctx, key)
						}
					}

					var results = make([]string, len(keys))
					err = GetMulti(ctx, keys, &results)
					Expect(err).ToNot(HaveOccurred())
					Expect(results).To(HaveLen(len(data)))
					for i, result := range results {
						Expect(result).To(Equal(data[i]))
					}
				})

				It("should work with arrays to pointers", func() {
					pts := []*string{}
					for _, d := range data {
						pts = append(pts, &d)
					}
					_, err := PutMulti(ctx, keys, pts)
					Expect(err).ToNot(HaveOccurred())

					var results = make([]*string, len(keys))
					err = GetMulti(ctx, keys, &results)
					Expect(err).ToNot(HaveOccurred())
					Expect(results).To(HaveLen(len(data)))

					// for i, result := range results {
					// 	Expect(*result).To(Equal(data[i]))
					// }
				})

				It("should error on things that cannot be gotton", func() {
					var result []string
					Expect(GetMulti(ctx, keys, &result)).To(HaveOccurred())
				})
			})

			Context("when putting", func() {
				It("should put all data to the keys", func() {
					_, err := PutMulti(ctx, keys, data)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})

				It("It should error on things that can be stored", func() {
					ts := []string{}
					_, err := PutMulti(ctx, keys, ts)
					Expect(err).To(HaveOccurred())
				})
			})
		})
	})
})
<commit_msg>fixed test import<commit_after>package horsefeather

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
	"gopkg.in/SeanDolphin/horsefeather.v1/test"
)

var encodedKeys = []string{
	"aghnbGlicmFyeXIMCxIGUGVyc29uGAEM",
	"aghnbGlicmFyeXIdCxIFR3JhcGgiEmdyYXBoOjctZGF5LWFjdGl2ZQw",
	"aghnbGlicmFyeXIhCxIJV29yZEluZGV4GIChPgwLEglXb3JkSW5kZXgYiQgM",
}

var _ = Describe("Store", func() {
	var ctx context.Context
	var cache = test.NewCache()
	var store = test.NewStore()

	var keys []*datastore.Key
	for _, encodedKey := range encodedKeys {
		key, _ := datastore.DecodeKey(encodedKey)
		keys = append(keys, key)
	}

	It("should panic when memcache is not present", func() {
		Expect(func() { mc(context.Background()) }).To(Panic())
	})

	It("should panic when datastore is not present", func() {
		Expect(func() { ds(context.Background()) }).To(Panic())
	})

	Context("when caching context are present", func() {
		BeforeEach(func() {
			ctx = AddMemcache(context.Background(), cache)
			ctx = AddDatastore(ctx, store)
		})

		AfterEach(func() {
			store.Clear()
			cache.Clear()
		})

		Context("when dealing with single keys", func() {
			var data = "bob"
			Context("when deleteing", func() {
				It("should delete the key", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")

						err = Delete(ctx, key)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})

				It("should error when there is nothing to delete", func() {
					for _, key := range keys {
						Expect(Delete(ctx, key)).To(HaveOccurred())
					}
				})

				It("should delete only to the memcache when memcache only is allowed.", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}

					ctx = OnlyMemcache(ctx, true)
					for _, key := range keys {
						err := Delete(ctx, key)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})

				It("should delete only to the datastore when datastore only is allowed.", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}

					ctx = OnlyDatastore(ctx, true)
					for _, key := range keys {
						err := Delete(ctx, key)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})
			})

			Context("when getting", func() {
				It("should get the key", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})

				It("should get the key when it is not present in memcache", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						cache.Delete(ctx, key)

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})

				It("should get the key when it is not present in the datastore", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						store.Delete(ctx, key)

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})

				It("should error when trying to get something that does not exist", func() {
					for _, key := range keys {
						var result string
						Expect(Get(ctx, key, &result)).To(HaveOccurred())
					}
				})

				It("should work when only memcache is allowed", func() {
					ctx = OnlyMemcache(ctx, true)
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})

				It("should work when only datastore is allowed", func() {
					ctx = OnlyDatastore(ctx, true)
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())

						var result string
						err = Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data))
					}
				})
			})

			Context("whenning putting a key", func() {
				It("should put the data", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})

				It("should error when trying to save nil", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, nil)
						Expect(err).To(HaveOccurred())
					}
				})

				It("should put only to the memcache when memcache only is allowed.", func() {
					ctx = OnlyMemcache(ctx, true)
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})

				It("should put only to the datastore when datastore only is allowed.", func() {
					ctx = OnlyDatastore(ctx, true)
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})
			})
		})

		Context("when dealing with multiple keys", func() {
			var data = []string{"t1", "t2", "t3"}
			Context("when deleting", func() {
				It("should delete all the keys", func() {
					_, err := PutMulti(ctx, keys, data)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "store")
					}

					err = DeleteMulti(ctx, keys)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {
						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})

				It("should return an error deleting an empty", func() {
					Expect(DeleteMulti(ctx, keys)).To(HaveOccurred())
				})

				It("should delete only to the memcache when memcache only is allowed.", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}

					ctx = OnlyMemcache(ctx, true)
					err := DeleteMulti(ctx, keys)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {

						Expect(cache.Contains(key)).To(BeFalse(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})

				It("should delete only to the datastore when datastore only is allowed.", func() {
					for _, key := range keys {
						_, err := Put(ctx, key, &data)
						Expect(err).ToNot(HaveOccurred())
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}

					ctx = OnlyDatastore(ctx, true)
					err := DeleteMulti(ctx, keys)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {

						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeFalse(), "datastore")
					}
				})
			})

			Context("when prefetching results", func() {
				It("should load data on multiple items", func() {
					PutMulti(ctx, keys, data)
					ctx = Prefetch(ctx, keys, data)
					store.Clear()
					cache.Clear()
					var result = make([]string, len(keys))
					err := GetMulti(ctx, keys, result)
					Expect(err).ToNot(HaveOccurred())
					for i := 0; i < len(keys); i++ {
						Expect(result[i]).To(Equal(data[i]))
					}
				})

				It("should load data for a single Item", func() {
					PutMulti(ctx, keys, data)
					ctx = Prefetch(ctx, keys, data)
					store.Clear()
					cache.Clear()

					for i, key := range keys {
						var result string
						err := Get(ctx, key, &result)
						Expect(err).ToNot(HaveOccurred())
						Expect(result).To(Equal(data[i]))
					}
				})
			})

			Context("when getting", func() {
				It("should get all the keys", func() {
					_, err := PutMulti(ctx, keys, &data)
					Expect(err).ToNot(HaveOccurred())

					var results = make([]string, len(keys))
					err = GetMulti(ctx, keys, &results)
					Expect(err).ToNot(HaveOccurred())
					Expect(results).To(HaveLen(len(data)))
					for i, result := range results {
						Expect(result).To(Equal(data[i]))
					}
				})

				It("should get all the data even if the memcache is empty", func() {
					_, err := PutMulti(ctx, keys, &data)
					Expect(err).ToNot(HaveOccurred())

					for _, key := range keys {
						cache.Delete(ctx, key)
					}

					var results = make([]string, len(keys))
					err = GetMulti(ctx, keys, &results)
					Expect(err).ToNot(HaveOccurred())
					Expect(results).To(HaveLen(len(data)))
					for i, result := range results {
						Expect(result).To(Equal(data[i]))
					}
				})

				It("should get all that even if it is split between the store and memcache", func() {
					_, err := PutMulti(ctx, keys, &data)
					Expect(err).ToNot(HaveOccurred())

					for i, key := range keys {
						if i%2 == 0 {
							cache.Delete(ctx, key)
						} else {
							store.Delete(ctx, key)
						}
					}

					var results = make([]string, len(keys))
					err = GetMulti(ctx, keys, &results)
					Expect(err).ToNot(HaveOccurred())
					Expect(results).To(HaveLen(len(data)))
					for i, result := range results {
						Expect(result).To(Equal(data[i]))
					}
				})

				It("should work with arrays to pointers", func() {
					pts := []*string{}
					for _, d := range data {
						pts = append(pts, &d)
					}
					_, err := PutMulti(ctx, keys, pts)
					Expect(err).ToNot(HaveOccurred())

					var results = make([]*string, len(keys))
					err = GetMulti(ctx, keys, &results)
					Expect(err).ToNot(HaveOccurred())
					Expect(results).To(HaveLen(len(data)))

					// for i, result := range results {
					// 	Expect(*result).To(Equal(data[i]))
					// }
				})

				It("should error on things that cannot be gotton", func() {
					var result []string
					Expect(GetMulti(ctx, keys, &result)).To(HaveOccurred())
				})
			})

			Context("when putting", func() {
				It("should put all data to the keys", func() {
					_, err := PutMulti(ctx, keys, data)
					Expect(err).ToNot(HaveOccurred())
					for _, key := range keys {
						Expect(cache.Contains(key)).To(BeTrue(), "memcache")
						Expect(store.Contains(key)).To(BeTrue(), "datastore")
					}
				})

				It("It should error on things that can be stored", func() {
					ts := []string{}
					_, err := PutMulti(ctx, keys, ts)
					Expect(err).To(HaveOccurred())
				})
			})
		})
	})
})
<|endoftext|>


<commit_before>package configloader

import (
	"bytes"
	"os"
	"strings"

	"github.com/Sirupsen/logrus"
	titancommon "github.com/iron-io/titan/common"
	"github.com/iron-io/titan/runner/agent"
	drivercommon "github.com/iron-io/titan/runner/drivers/common"
	"github.com/pivotal-golang/bytefmt"
	"github.com/spf13/viper"
)

// This package provides a function to read a complete runner configuration
// from various sources (currently JSON file and environment variables). Just
// importing the configloader package will cause it to read in configuration
// from the sources. Callers may then use RunnerConfiguration() to get an
// initialized runner configuration. Other methods provide access to some common functionality.
// A basic configuration looks like:
//
//     {
//       "runner": {
//         ...
//       },
//       "stats": {
//         ...
//       },
//       "log_level": "..."
//     }
//
// configloader will also set the runner's log package's log level to the value obtained here.
// StatsConfig() returns a stats configuration which can be passed to titan
// stats.New(). The runner's stats are not initialized using this
// configuration. Use the environment passed to agent.NewRunner() to handle
// this for now.
//
// For more configuration sub-trees related to the app specific behaviour that
// wraps a runner, use ExtendedConfiguration(key) to obtain a configuration
// sub-tree.
//
// This package is currently shared by the Titan runner and Swapi runner. Be
// careful when making changes to preserve compatibility.
//
// The following top-level keys are reserved for runner configuration, since
// they reflect the names of the environment variables captured in the process:
// - log_level
// - driver
// - concurrency
// - memory_per_job
// - instance_id
//
// configloader will terminate if any configuration does not match valid values.
// For this reason, it should only be used from a binary's main().

var runnerConfig agent.Config
var apiURL string

func init() {
	host, _ := os.Hostname() // if it's blank then at least we tried

	// Set up defaults.
	viper.SetConfigName("config")
	viper.AddConfigPath(".")
	replacer := strings.NewReplacer(".", "_")
	viper.SetEnvKeyReplacer(replacer)
	viper.AutomaticEnv()
	viper.SetDefault("log_to", "stderr")
	viper.SetDefault("log_prefix", host)
	viper.SetDefault("log_level", "info")
	viper.SetDefault("driver", "docker")
	viper.SetDefault("api_url", "http://localhost:8080")
	viper.SetDefault("memory_per_job", "256M")
	// If we do not set it here, viper will not load the relevant environment var. Set it to int's default value.
	viper.SetDefault("concurrency", "0")

	// Initialize viper from config.json.
	err := viper.ReadInConfig()
	if err != nil {
		if _, ok := err.(viper.UnsupportedConfigError); ok {
			logrus.WithError(err).Error("Couldn't read config file, this is fine, it's not required")
			// ignore
		} else {
			logrus.WithError(err).Fatal("Error reading config file")
		}
	}

	// Initialize viper from CONFIG environment variable.
	if os.Getenv("CONFIG") != "" {
		viper.SetConfigType("json")
		err = viper.ReadConfig(bytes.NewBufferString(os.Getenv("CONFIG")))
		if err != nil {
			logrus.WithError(err).Fatal("Error reading CONFIG from env")
		}
	}

	// Set up defaults.
	dconfig := drivercommon.DefaultConfig()
	memPerJob, err := bytefmt.ToBytes(viper.GetString("memory_per_job"))
	if err != nil {
		logrus.WithError(err).WithFields(logrus.Fields{"memory_per_job": viper.GetString("memory_per_job")}).Fatal("Invalid MEMORY_PER_JOB variable")
	}
	// todo: we should pass the entire config to the driver so they can use the same variable. Or make a config interface and pass it along.
	dconfig.Memory = memPerJob
	runnerConfig.DriverConfig = dconfig

	// Let loaded configuration overwrite defaults.
	if viper.InConfig("runner") {
		err = viper.Sub("runner").Unmarshal(&runnerConfig)
		if err != nil {
			logrus.WithError(err).Fatal("could not unmarshal runner configuration from config")
		}
	}

	// Finally, let top-level reserved keys overwrite any existing settings.
	runnerConfig.Driver = viper.GetString("driver")

	if concurrency := viper.GetInt("concurrency"); concurrency > 0 {
		runnerConfig.Concurrency = concurrency
	}

	runnerConfig.InstanceId = viper.GetString("instance_id")

	apiURL = viper.GetString("api_url")
	titancommon.SetLogLevel(viper.GetString("log_level"))
	titancommon.SetLogDest(viper.GetString("log_to"), viper.GetString("log_prefix"))
	logrus.WithFields(logrus.Fields{"runner_configuration": runnerConfig}).Info("configloader runner configuration")
}

func RunnerConfiguration() *agent.Config {
	return &runnerConfig
}

func ApiURL() string {
	return apiURL
}

func ExtendedConfiguration(key string) *viper.Viper {
	if !viper.InConfig(key) {
		return nil
	}

	return viper.Sub(key)
}
<commit_msg>fix capitalization; setting id for swapi_tests<commit_after>package configloader

import (
	"bytes"
	"os"
	"strings"

	"github.com/Sirupsen/logrus"
	titancommon "github.com/iron-io/titan/common"
	"github.com/iron-io/titan/runner/agent"
	drivercommon "github.com/iron-io/titan/runner/drivers/common"
	"github.com/pivotal-golang/bytefmt"
	"github.com/spf13/viper"
)

// This package provides a function to read a complete runner configuration
// from various sources (currently JSON file and environment variables). Just
// importing the configloader package will cause it to read in configuration
// from the sources. Callers may then use RunnerConfiguration() to get an
// initialized runner configuration. Other methods provide access to some common functionality.
// A basic configuration looks like:
//
//     {
//       "runner": {
//         ...
//       },
//       "stats": {
//         ...
//       },
//       "log_level": "..."
//     }
//
// configloader will also set the runner's log package's log level to the value obtained here.
// StatsConfig() returns a stats configuration which can be passed to titan
// stats.New(). The runner's stats are not initialized using this
// configuration. Use the environment passed to agent.NewRunner() to handle
// this for now.
//
// For more configuration sub-trees related to the app specific behaviour that
// wraps a runner, use ExtendedConfiguration(key) to obtain a configuration
// sub-tree.
//
// This package is currently shared by the Titan runner and Swapi runner. Be
// careful when making changes to preserve compatibility.
//
// The following top-level keys are reserved for runner configuration, since
// they reflect the names of the environment variables captured in the process:
// - log_level
// - driver
// - concurrency
// - memory_per_job
// - instance_id
//
// configloader will terminate if any configuration does not match valid values.
// For this reason, it should only be used from a binary's main().

var runnerConfig agent.Config
var apiURL string

func init() {
	host, _ := os.Hostname() // if it's blank then at least we tried

	// Set up defaults.
	viper.SetConfigName("config")
	viper.AddConfigPath(".")
	replacer := strings.NewReplacer(".", "_")
	viper.SetEnvKeyReplacer(replacer)
	viper.AutomaticEnv()
	viper.SetDefault("log_to", "stderr")
	viper.SetDefault("log_prefix", host)
	viper.SetDefault("log_level", "info")
	viper.SetDefault("driver", "docker")
	viper.SetDefault("api_url", "http://localhost:8080")
	viper.SetDefault("memory_per_job", "256M")
	// If we do not set it here, viper will not load the relevant environment var. Set it to int's default value.
	viper.SetDefault("concurrency", "0")

	// Initialize viper from config.json.
	err := viper.ReadInConfig()
	if err != nil {
		if _, ok := err.(viper.UnsupportedConfigError); ok {
			logrus.WithError(err).Error("Couldn't read config file, this is fine, it's not required")
			// ignore
		} else {
			logrus.WithError(err).Fatal("Error reading config file")
		}
	}

	// Initialize viper from CONFIG environment variable.
	if os.Getenv("CONFIG") != "" {
		viper.SetConfigType("json")
		err = viper.ReadConfig(bytes.NewBufferString(os.Getenv("CONFIG")))
		if err != nil {
			logrus.WithError(err).Fatal("Error reading CONFIG from env")
		}
	}

	// Set up defaults.
	dconfig := drivercommon.DefaultConfig()
	memPerJob, err := bytefmt.ToBytes(viper.GetString("memory_per_job"))
	if err != nil {
		logrus.WithError(err).WithFields(logrus.Fields{"memory_per_job": viper.GetString("memory_per_job")}).Fatal("Invalid MEMORY_PER_JOB variable")
	}
	// todo: we should pass the entire config to the driver so they can use the same variable. Or make a config interface and pass it along.
	dconfig.Memory = memPerJob
	runnerConfig.DriverConfig = dconfig

	// Let loaded configuration overwrite defaults.
	if viper.InConfig("runner") {
		err = viper.Sub("runner").Unmarshal(&runnerConfig)
		if err != nil {
			logrus.WithError(err).Fatal("could not unmarshal runner configuration from config")
		}
	}

	// Finally, let top-level reserved keys overwrite any existing settings.
	runnerConfig.Driver = viper.GetString("driver")

	if concurrency := viper.GetInt("concurrency"); concurrency > 0 {
		runnerConfig.Concurrency = concurrency
	}

	runnerConfig.InstanceID = viper.GetString("instance_id")

	apiURL = viper.GetString("api_url")
	titancommon.SetLogLevel(viper.GetString("log_level"))
	titancommon.SetLogDest(viper.GetString("log_to"), viper.GetString("log_prefix"))
	logrus.WithFields(logrus.Fields{"runner_configuration": runnerConfig}).Info("configloader runner configuration")
}

func RunnerConfiguration() *agent.Config {
	return &runnerConfig
}

func ApiURL() string {
	return apiURL
}

func ExtendedConfiguration(key string) *viper.Viper {
	if !viper.InConfig(key) {
		return nil
	}

	return viper.Sub(key)
}
<|endoftext|>


<commit_before>package main

/* Copyright 2014, Timothy Bogdala <tdb@animal-machine.com>
See the LICENSE file for more details. */

import (
	"bufio"
	"fmt"
	"math/rand"

	"github.com/tbogdala/noisey"
	"os"
)

func main() {
	fmt.Println("Doing a noisey test.")

	const imageSize int = 128

	// make a test generator seeded to 1
	r := rand.New(rand.NewSource(int64(1)))

	// create a new perlin noise generator with a 'standard' 256 int
	// table of random permutations
	perlin := noisey.NewPerlinGenerator2D(r, 256)

	// make an ascii pixel image by calculating random noise
	pixels := make([]float64, imageSize*imageSize)
	for i := 0; i < 100; i++ {
		for y := 0; y < imageSize; y++ {
			for x := 0; x < imageSize; x++ {
				v := perlin.Get(float64(x)*0.1, float64(y)*0.1, noisey.FastQuality)
				v = v*0.5 + 0.5
				pixels[y*imageSize+x] = v
			}
		}
	}

	// print the image out to the terminal using the symbols below
	symbols := []string{" ", "░", "▒", "▓", "█", "█"}
	out := bufio.NewWriter(os.Stdout)
	for y := 0; y < imageSize; y++ {
		for x := 0; x < imageSize; x++ {
			fmt.Fprint(out, symbols[int(pixels[y*imageSize+x]/0.2)])
		}
		fmt.Fprintln(out)
	}
	out.Flush()
}
<commit_msg>didnt need 100 runs of the loop.<commit_after>package main

/* Copyright 2014, Timothy Bogdala <tdb@animal-machine.com>
See the LICENSE file for more details. */

import (
	"bufio"
	"fmt"
	"math/rand"

	"github.com/tbogdala/noisey"
	"os"
)

func main() {
	fmt.Println("Doing a noisey test.")

	const imageSize int = 128

	// make a test generator seeded to 1
	r := rand.New(rand.NewSource(int64(1)))

	// create a new perlin noise generator with a 'standard' 256 int
	// table of random permutations
	perlin := noisey.NewPerlinGenerator2D(r, 256)

	// make an ascii pixel image by calculating random noise
	pixels := make([]float64, imageSize*imageSize)
	for y := 0; y < imageSize; y++ {
		for x := 0; x < imageSize; x++ {
			v := perlin.Get(float64(x)*0.1, float64(y)*0.1, noisey.FastQuality)
			v = v*0.5 + 0.5
			pixels[y*imageSize+x] = v
		}
	}

	// print the image out to the terminal using the symbols below
	symbols := []string{" ", "░", "▒", "▓", "█", "█"}
	out := bufio.NewWriter(os.Stdout)
	for y := 0; y < imageSize; y++ {
		for x := 0; x < imageSize; x++ {
			fmt.Fprint(out, symbols[int(pixels[y*imageSize+x]/0.2)])
		}
		fmt.Fprintln(out)
	}
	out.Flush()
}
<|endoftext|>