<commit_before>package ssh_agent

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"

	uuid "github.com/satori/go.uuid"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"

	"github.com/flant/logboek"
	"github.com/flant/werf/pkg/util"
	"github.com/flant/werf/pkg/werf"
)

var (
	SSHAuthSock string
	tmpSockPath string
)

func Init(keys []string) error {
	for _, key := range keys {
		if keyExists, err := util.FileExists(key); !keyExists {
			return fmt.Errorf("specified ssh key %s does not exist", key)
		} else if err != nil {
			return fmt.Errorf("specified ssh key %s does not exist: %v", key, err)
		}
	}

	if len(keys) > 0 {
		agentSock, err := runSSHAgentWithKeys(keys)
		if err != nil {
			return err
		}
		SSHAuthSock = agentSock

		return nil
	}

	systemAgentSock := os.Getenv("SSH_AUTH_SOCK")
	systemAgentSockExists, _ := util.FileExists(systemAgentSock)
	if systemAgentSock != "" && systemAgentSockExists {
		SSHAuthSock = systemAgentSock
		logboek.LogF("Using system ssh-agent: %s\n", systemAgentSock)
		return nil
	}

	var defaultKeys []string
	for _, defaultFileName := range []string{"id_rsa", "id_dsa"} {
		path := filepath.Join(os.Getenv("HOME"), ".ssh", defaultFileName)
		if keyExists, _ := util.FileExists(path); keyExists {
			defaultKeys = append(defaultKeys, path)
		}
	}

	if len(defaultKeys) > 0 {
		var validKeys []string

		for _, key := range defaultKeys {
			keyData, err := ioutil.ReadFile(key)
			if err != nil {
				continue
			}
			_, err = ssh.ParseRawPrivateKeyWithPassphrase(keyData, []byte{})
			if err != nil {
				continue
			}

			validKeys = append(validKeys, key)
		}

		if len(validKeys) > 0 {
			agentSock, err := runSSHAgentWithKeys(validKeys)
			if err != nil {
				return err
			}
			SSHAuthSock = agentSock
		}
	}

	return nil
}

func Terminate() error {
	if tmpSockPath != "" {
		err := os.RemoveAll(tmpSockPath)
		if err != nil {
			return fmt.Errorf("unable to remove tmp ssh agent sock %s: %s", tmpSockPath, err)
		}
	}

	return nil
}

func runSSHAgentWithKeys(keys []string) (string, error) {
	agentSock, err := runSSHAgent()
	if err != nil {
		return "", fmt.Errorf("error running ssh agent: %s", err)
	}

	for _, key := range keys {
		err := addSSHKey(agentSock, key)
		if err != nil {
			return "", fmt.Errorf("error adding ssh key %s: %s", key, err)
		}
	}

	return agentSock, nil
}

func runSSHAgent() (string, error) {
	sockPath := filepath.Join(werf.GetTmpDir(), "werf-ssh-agent", uuid.NewV4().String())
	tmpSockPath = sockPath

	err := os.MkdirAll(filepath.Dir(sockPath), os.ModePerm)
	if err != nil {
		return "", err
	}

	ln, err := net.Listen("unix", sockPath)
	if err != nil {
		return "", fmt.Errorf("error listen unix sock %s: %s", sockPath, err)
	}

	logboek.LogF("Running ssh agent on unix sock: %s\n", sockPath)

	go func() {
		agnt := agent.NewKeyring()

		for {
			conn, err := ln.Accept()
			if err != nil {
				logboek.LogErrorF("WARNING: failed to accept ssh-agent connection: %s\n", err)
				continue
			}

			go func() {
				var err error

				err = agent.ServeAgent(agnt, conn)
				if err != nil && err != io.EOF {
					logboek.LogErrorF("WARNING: ssh-agent server error: %s\n", err)
					return
				}

				err = conn.Close()
				if err != nil {
					logboek.LogErrorF("WARNING: ssh-agent server connection close error: %s\n", err)
					return
				}
			}()
		}
	}()

	return sockPath, nil
}

func addSSHKey(authSock string, key string) error {
	conn, err := net.Dial("unix", authSock)
	if err != nil {
		return fmt.Errorf("error dialing with ssh agent %s: %s", authSock, err)
	}
	defer conn.Close()

	agentClient := agent.NewClient(conn)

	keyData, err := ioutil.ReadFile(key)
	if err != nil {
		return fmt.Errorf("error reading key file %s: %s", key, err)
	}

	privateKey, err := ssh.ParseRawPrivateKeyWithPassphrase(keyData, []byte{})
	if err != nil {
		return fmt.Errorf("error parsing private key %s: %s", key, err)
	}

	err = agentClient.Add(agent.AddedKey{PrivateKey: privateKey})
	if err != nil {
		return err
	}

	logboek.LogF("Added private key %s to ssh agent %s\n", key, authSock)

	return nil
}
<commit_msg>Fix ssh-agent not adding default keys and keys passed by the `--ssh-key` option<commit_after>package ssh_agent

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"

	uuid "github.com/satori/go.uuid"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"

	"github.com/flant/logboek"
	"github.com/flant/werf/pkg/util"
	"github.com/flant/werf/pkg/werf"
)

var (
	SSHAuthSock string
	tmpSockPath string
)

func Init(keys []string) error {
	for _, key := range keys {
		if keyExists, err := util.FileExists(key); !keyExists {
			return fmt.Errorf("specified ssh key %s does not exist", key)
		} else if err != nil {
			return fmt.Errorf("specified ssh key %s does not exist: %v", key, err)
		}
	}

	if len(keys) > 0 {
		agentSock, err := runSSHAgentWithKeys(keys)
		if err != nil {
			return err
		}
		SSHAuthSock = agentSock

		return nil
	}

	systemAgentSock := os.Getenv("SSH_AUTH_SOCK")
	systemAgentSockExists, _ := util.FileExists(systemAgentSock)
	if systemAgentSock != "" && systemAgentSockExists {
		SSHAuthSock = systemAgentSock
		logboek.LogF("Using system ssh-agent: %s\n", systemAgentSock)
		return nil
	}

	var defaultKeys []string
	for _, defaultFileName := range []string{"id_rsa", "id_dsa"} {
		path := filepath.Join(os.Getenv("HOME"), ".ssh", defaultFileName)

		if keyExists, _ := util.FileExists(path); keyExists {
			defaultKeys = append(defaultKeys, path)
		}
	}

	if len(defaultKeys) > 0 {
		var validKeys []string

		for _, key := range defaultKeys {
			keyData, err := ioutil.ReadFile(key)
			if err != nil {
				logboek.LogErrorF("WARNING: cannot read default key %s: %s\n", key, err)
				continue
			}
			_, err = ssh.ParseRawPrivateKey(keyData)
			if err != nil {
				logboek.LogErrorF("WARNING: default key %s validation error: %s\n", key, err)
				continue
			}

			validKeys = append(validKeys, key)
		}

		if len(validKeys) > 0 {
			agentSock, err := runSSHAgentWithKeys(validKeys)
			if err != nil {
				return err
			}
			SSHAuthSock = agentSock
		}
	}

	return nil
}

func Terminate() error {
	if tmpSockPath != "" {
		err := os.RemoveAll(tmpSockPath)
		if err != nil {
			return fmt.Errorf("unable to remove tmp ssh agent sock %s: %s", tmpSockPath, err)
		}
	}

	return nil
}

func runSSHAgentWithKeys(keys []string) (string, error) {
	agentSock, err := runSSHAgent()
	if err != nil {
		return "", fmt.Errorf("error running ssh agent: %s", err)
	}

	for _, key := range keys {
		err := addSSHKey(agentSock, key)
		if err != nil {
			return "", fmt.Errorf("error adding ssh key %s: %s", key, err)
		}
	}

	return agentSock, nil
}

func runSSHAgent() (string, error) {
	sockPath := filepath.Join(werf.GetTmpDir(), "werf-ssh-agent", uuid.NewV4().String())
	tmpSockPath = sockPath

	err := os.MkdirAll(filepath.Dir(sockPath), os.ModePerm)
	if err != nil {
		return "", err
	}

	ln, err := net.Listen("unix", sockPath)
	if err != nil {
		return "", fmt.Errorf("error listen unix sock %s: %s", sockPath, err)
	}

	logboek.LogF("Running ssh agent on unix sock: %s\n", sockPath)

	go func() {
		agnt := agent.NewKeyring()

		for {
			conn, err := ln.Accept()
			if err != nil {
				logboek.LogErrorF("WARNING: failed to accept ssh-agent connection: %s\n", err)
				continue
			}

			go func() {
				var err error

				err = agent.ServeAgent(agnt, conn)
				if err != nil && err != io.EOF {
					logboek.LogErrorF("WARNING: ssh-agent server error: %s\n", err)
					return
				}

				err = conn.Close()
				if err != nil {
					logboek.LogErrorF("WARNING: ssh-agent server connection close error: %s\n", err)
					return
				}
			}()
		}
	}()

	return sockPath, nil
}

func addSSHKey(authSock string, key string) error {
	conn, err := net.Dial("unix", authSock)
	if err != nil {
		return fmt.Errorf("error dialing with ssh agent %s: %s", authSock, err)
	}
	defer conn.Close()

	agentClient := agent.NewClient(conn)

	keyData, err := ioutil.ReadFile(key)
	if err != nil {
		return fmt.Errorf("error reading key file %s: %s", key, err)
	}

	privateKey, err := ssh.ParseRawPrivateKey(keyData)
	if err != nil {
		return fmt.Errorf("error parsing private key %s: %s", key, err)
	}

	err = agentClient.Add(agent.AddedKey{PrivateKey: privateKey})
	if err != nil {
		return err
	}

	logboek.LogF("Added private key %s to ssh agent %s\n", key, authSock)

	return nil
}
<|endoftext|>


<commit_before>//
// Last.Backend LLC CONFIDENTIAL
// __________________
//
// [2014] - [2017] Last.Backend LLC
// All Rights Reserved.
//
// NOTICE:  All information contained herein is, and remains
// the property of Last.Backend LLC and its suppliers,
// if any.  The intellectual and technical concepts contained
// herein are proprietary to Last.Backend LLC
// and its suppliers and may be covered by Russian Federation and Foreign Patents,
// patents in process, and are protected by trade secret or copyright law.
// Dissemination of this information or reproduction of this material
// is strictly forbidden unless prior written permission is obtained
// from Last.Backend LLC.
//

package etcd3

import (
	"errors"
	"fmt"
	"github.com/coreos/etcd/clientv3"
	"github.com/lastbackend/lastbackend/pkg/util/serializer"
	st "github.com/lastbackend/lastbackend/pkg/storage/store"
	"github.com/lastbackend/lastbackend/pkg/util/converter"
	"github.com/lastbackend/lastbackend/pkg/util/validator"
	"golang.org/x/net/context"
	"path"
	"reflect"
	"regexp"
	"strings"
)

type store struct {
	client     *clientv3.Client
	opts       []clientv3.OpOption
	codec      serializer.Codec
	pathPrefix string
}

// Need for decode array bytes
type buffer []byte

func (s *store) Create(ctx context.Context, key string, obj, outPtr interface{}, ttl uint64) error {
	data, err := serializer.Encode(s.codec, obj)
	if err != nil {
		return err
	}
	key = path.Join(s.pathPrefix, key)
	opts, err := s.ttlOpts(ctx, int64(ttl))
	if err != nil {
		return err
	}
	fmt.Println("Create:", key, string(data))
	txnResp, err := s.client.KV.Txn(ctx).
		If(clientv3.Compare(clientv3.ModRevision(key), "=", 0)).
		Then(clientv3.OpPut(key, string(data), opts...)).Commit()
	if err != nil {
		return err
	}
	if !txnResp.Succeeded {
		return errors.New(st.ErrKeyExists)
	}
	if validator.IsNil(outPtr) {
		return nil
	}
	if outPtr != nil {
		return decode(s.codec, data, outPtr)
	}
	return nil
}

func (s *store) Get(ctx context.Context, key string, outPtr interface{}) error {
	key = path.Join(s.pathPrefix, key)
	fmt.Println("Get:", key)
	res, err := s.client.KV.Get(ctx, key, s.opts...)
	if err != nil {
		return err
	}
	if len(res.Kvs) == 0 {
		return errors.New(st.ErrKeyNotFound)
	}
	return decode(s.codec, res.Kvs[0].Value, outPtr)
}

func (s *store) List(ctx context.Context, key, keyRegexFilter string, listOutPtr interface{}) error {
	key = path.Join(s.pathPrefix, key)
	if !strings.HasSuffix(key, "/") {
		key += "/"
	}
	fmt.Println("List:", key)
	getResp, err := s.client.KV.Get(ctx, key, clientv3.WithPrefix())
	if err != nil {
		return err
	}
	r, _ := regexp.Compile(keyRegexFilter)
	items := make([]buffer, 0, len(getResp.Kvs))
	for _, kv := range getResp.Kvs {
		if r.MatchString(string(kv.Key)) {
			items = append(items, buffer(kv.Value))
		}
	}
	return decodeList(s.codec, items, listOutPtr)
}

func (s *store) Update(ctx context.Context, key string, obj, outPtr interface{}, ttl uint64) error {
	data, err := serializer.Encode(s.codec, obj)
	if err != nil {
		return err
	}
	key = path.Join(s.pathPrefix, key)
	fmt.Println("Update:", key, string(data))
	opts, err := s.ttlOpts(ctx, int64(ttl))
	if err != nil {
		return err
	}
	txnResp, err := s.client.KV.Txn(ctx).
		If(clientv3.Compare(clientv3.ModRevision(key), "!=", 0)).
		Then(clientv3.OpPut(key, string(data), opts...)).
		Commit()
	if err != nil {
		return err
	}
	if !txnResp.Succeeded {
		return errors.New(st.ErrKeyNotFound)
	}
	if validator.IsNil(outPtr) {
		return nil
	}
	if outPtr != nil {
		return decode(s.codec, data, outPtr)
	}
	return nil
}

func (s *store) Delete(ctx context.Context, key string, outPtr interface{}) error {
	key = path.Join(s.pathPrefix, key)
	fmt.Println("Del:", key)
	res, err := s.client.KV.Txn(ctx).Then(clientv3.OpGet(key), clientv3.OpDelete(key)).Commit()
	if err != nil {
		return err
	}
	if validator.IsNil(outPtr) {
		return nil
	}
	getResp := res.Responses[0].GetResponseRange()
	if len(getResp.Kvs) == 0 {
		return errors.New(st.ErrKeyNotFound)
	}
	return decode(s.codec, getResp.Kvs[0].Value, outPtr)
}

func (s *store) Begin(ctx context.Context) st.ITx {
	return &tx{
		store:   s,
		context: ctx,
		txn:     s.client.KV.Txn(ctx),
	}
}

func decode(s serializer.Codec, value []byte, outPtr interface{}) error {
	if _, err := converter.EnforcePtr(outPtr); err != nil {
		panic("Error: unable to convert output object to pointer")
	}
	return serializer.Decode(s, value, outPtr)
}

func decodeList(codec serializer.Codec, items []buffer, ListOutPtr interface{}) error {
	v, err := converter.EnforcePtr(ListOutPtr)
	if err != nil || v.Kind() != reflect.Slice {
		panic("Error: need ptr to slice")
	}
	for _, item := range items {
		var obj = reflect.New(v.Type().Elem()).Interface().(interface{})
		err := serializer.Decode(codec, item, obj)
		if err != nil {
			return err
		}
		v.Set(reflect.Append(v, reflect.ValueOf(obj).Elem()))
	}
	return nil
}

func (s *store) ttlOpts(ctx context.Context, ttl int64) ([]clientv3.OpOption, error) {
	if ttl == 0 {
		return nil, nil
	}
	lcr, err := s.client.Lease.Grant(ctx, ttl)
	if err != nil {
		return nil, err
	}
	return []clientv3.OpOption{clientv3.WithLease(clientv3.LeaseID(lcr.ID))}, nil
}
<commit_msg>update list method storage<commit_after>//
// Last.Backend LLC CONFIDENTIAL
// __________________
//
// [2014] - [2017] Last.Backend LLC
// All Rights Reserved.
//
// NOTICE:  All information contained herein is, and remains
// the property of Last.Backend LLC and its suppliers,
// if any.  The intellectual and technical concepts contained
// herein are proprietary to Last.Backend LLC
// and its suppliers and may be covered by Russian Federation and Foreign Patents,
// patents in process, and are protected by trade secret or copyright law.
// Dissemination of this information or reproduction of this material
// is strictly forbidden unless prior written permission is obtained
// from Last.Backend LLC.
//

package etcd3

import (
	"errors"
	"fmt"
	"github.com/coreos/etcd/clientv3"
	st "github.com/lastbackend/lastbackend/pkg/storage/store"
	"github.com/lastbackend/lastbackend/pkg/util/converter"
	"github.com/lastbackend/lastbackend/pkg/util/serializer"
	"github.com/lastbackend/lastbackend/pkg/util/validator"
	"golang.org/x/net/context"
	"path"
	"reflect"
	"regexp"
	"strings"
)

type store struct {
	client     *clientv3.Client
	opts       []clientv3.OpOption
	codec      serializer.Codec
	pathPrefix string
}

// Need for decode array bytes
type buffer []byte

func (s *store) Create(ctx context.Context, key string, obj, outPtr interface{}, ttl uint64) error {
	data, err := serializer.Encode(s.codec, obj)
	if err != nil {
		return err
	}
	key = path.Join(s.pathPrefix, key)
	opts, err := s.ttlOpts(ctx, int64(ttl))
	if err != nil {
		return err
	}
	fmt.Println("Create:", key, string(data))
	txnResp, err := s.client.KV.Txn(ctx).
		If(clientv3.Compare(clientv3.ModRevision(key), "=", 0)).
		Then(clientv3.OpPut(key, string(data), opts...)).Commit()
	if err != nil {
		return err
	}
	if !txnResp.Succeeded {
		return errors.New(st.ErrKeyExists)
	}
	if validator.IsNil(outPtr) {
		return nil
	}
	if outPtr != nil {
		return decode(s.codec, data, outPtr)
	}
	return nil
}

func (s *store) Get(ctx context.Context, key string, outPtr interface{}) error {
	key = path.Join(s.pathPrefix, key)
	fmt.Println("Get:", key)
	res, err := s.client.KV.Get(ctx, key, s.opts...)
	if err != nil {
		return err
	}
	if len(res.Kvs) == 0 {
		return errors.New(st.ErrKeyNotFound)
	}
	return decode(s.codec, res.Kvs[0].Value, outPtr)
}

func (s *store) List(ctx context.Context, key, keyRegexFilter string, listOutPtr interface{}) error {
	key = path.Join(s.pathPrefix, key)
	if !strings.HasSuffix(key, "/") {
		key += "/"
	}
	fmt.Println("List:", key)
	getResp, err := s.client.KV.Get(ctx, key, clientv3.WithPrefix())
	if err != nil {
		return err
	}
	r, _ := regexp.Compile(keyRegexFilter)
	items := make([]buffer, 0, len(getResp.Kvs))
	for _, kv := range getResp.Kvs {
		if r.MatchString(string(kv.Key)) {
			items = append(items, buffer(kv.Value))
		}
	}
	return decodeList(s.codec, items, listOutPtr)
}

func (s *store) Update(ctx context.Context, key string, obj, outPtr interface{}, ttl uint64) error {
	data, err := serializer.Encode(s.codec, obj)
	if err != nil {
		return err
	}
	key = path.Join(s.pathPrefix, key)
	fmt.Println("Update:", key, string(data))
	opts, err := s.ttlOpts(ctx, int64(ttl))
	if err != nil {
		return err
	}
	txnResp, err := s.client.KV.Txn(ctx).
		If(clientv3.Compare(clientv3.ModRevision(key), "!=", 0)).
		Then(clientv3.OpPut(key, string(data), opts...)).
		Commit()
	if err != nil {
		return err
	}
	if !txnResp.Succeeded {
		return errors.New(st.ErrKeyNotFound)
	}
	if validator.IsNil(outPtr) {
		return nil
	}
	if outPtr != nil {
		return decode(s.codec, data, outPtr)
	}
	return nil
}

func (s *store) Delete(ctx context.Context, key string, outPtr interface{}) error {
	key = path.Join(s.pathPrefix, key)
	fmt.Println("Del:", key)
	res, err := s.client.KV.Txn(ctx).Then(clientv3.OpGet(key), clientv3.OpDelete(key)).Commit()
	if err != nil {
		return err
	}
	if validator.IsNil(outPtr) {
		return nil
	}
	getResp := res.Responses[0].GetResponseRange()
	if len(getResp.Kvs) == 0 {
		return errors.New(st.ErrKeyNotFound)
	}
	return decode(s.codec, getResp.Kvs[0].Value, outPtr)
}

func (s *store) Begin(ctx context.Context) st.ITx {
	return &tx{
		store:   s,
		context: ctx,
		txn:     s.client.KV.Txn(ctx),
	}
}

func decode(s serializer.Codec, value []byte, outPtr interface{}) error {
	if _, err := converter.EnforcePtr(outPtr); err != nil {
		panic("Error: unable to convert output object to pointer")
	}
	return serializer.Decode(s, value, outPtr)
}

func decodeList(codec serializer.Codec, items []buffer, ListOutPtr interface{}) error {
	v, err := converter.EnforcePtr(ListOutPtr)
	if err != nil || (v.Kind() != reflect.Slice && v.Kind() != reflect.Map) {
		panic("Error: need ptr to slice")
	}
	for _, item := range items {
		var obj = reflect.New(v.Type().Elem()).Interface().(interface{})
		err := serializer.Decode(codec, item, obj)
		if err != nil {
			return err
		}
		v.Set(reflect.Append(v, reflect.ValueOf(obj).Elem()))
	}
	return nil
}

func (s *store) ttlOpts(ctx context.Context, ttl int64) ([]clientv3.OpOption, error) {
	if ttl == 0 {
		return nil, nil
	}
	lcr, err := s.client.Lease.Grant(ctx, ttl)
	if err != nil {
		return nil, err
	}
	return []clientv3.OpOption{clientv3.WithLease(clientv3.LeaseID(lcr.ID))}, nil
}
<|endoftext|>


<commit_before>package services

import (
	"jbrodriguez/mediagui/server/src/dto"
	"jbrodriguez/mediagui/server/src/lib"
	"jbrodriguez/mediagui/server/src/model"
	"jbrodriguez/mediagui/server/src/proto"

	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/jbrodriguez/actor"
	"github.com/jbrodriguez/mlog"
	"github.com/jbrodriguez/pubsub"
	"github.com/micro/go-micro/client"

	"golang.org/x/net/context"
)

// Core -
type Core struct {
	bus      *pubsub.PubSub
	settings *lib.Settings

	actor *actor.Actor
	re    *regexp.Regexp
	maps  map[string]bool

	wg sync.WaitGroup
}

// NewCore -
func NewCore(bus *pubsub.PubSub, settings *lib.Settings) *Core {
	core := &Core{
		bus:      bus,
		settings: settings,
		actor:    actor.NewActor(bus),
	}
	return core
}

// Start -
func (c *Core) Start() {
	mlog.Info("Starting service Core ...")

	c.actor.Register("/get/config", c.getConfig)
	c.actor.Register("/post/import", c.importMovies)
	c.actor.Register("/post/add", c.addMovie)
	c.actor.Register("/post/prune", c.pruneMovies)
	c.actor.Register("/put/config/folder", c.addMediaFolder)
	c.actor.Register("/put/movies/fix", c.fixMovie)

	c.actor.Register("/event/movie/found", c.doMovieFound)
	c.actor.Register("/event/movie/tmdbnotfound", c.doMovieTmdbNotFound)
	c.actor.Register("/event/movie/scraped", c.doMovieScraped)
	c.actor.Register("/event/movie/rescraped", c.doMovieReScraped)
	// c.actor.Register("/event/movie/updated", c.doMovieUpdated)
	// c.actor.Register("/event/movie/cached/forced", c.doMovieCachedForced)

	c.actor.Register("/event/workunit/done", c.doWorkUnitDone)

	c.re = regexp.MustCompile(`(?i)/Volumes/(.*?)/.*`)
	c.maps = make(map[string]bool)

	// maps["/Volumes/wopr-films"] = true // for example
	for _, folder := range c.settings.MediaFolders {
		c.maps[folder] = false
		if _, err := os.Stat(folder); err == nil {
			c.maps[folder] = true
		}
	}

	go c.actor.React()
}

// Stop -
func (c *Core) Stop() {
	mlog.Info("Stopped service Core ...")
}

func (c *Core) getConfig(msg *pubsub.Message) {
	msg.Reply <- &c.settings.Config
	mlog.Info("Sent config")
}

func (c *Core) importMovies(msg *pubsub.Message) {
	t0 := time.Now()
	// mlog.Info("Begin movie scanning ...")
	lib.Notify(c.bus, "import:begin", "Import process started")

	c.wg.Add(1)

	c.bus.Pub(nil, "/command/movie/scan")
	//	msg.Reply <- &c.settings.Config
	// mlog.Info("Import finished")

	go func() {
		c.wg.Wait()
		// 𝛥t := float64(time.Since(t0)) / 1e9
		lib.Notify(c.bus, "import:end", fmt.Sprintf("Import process finished (%s elapsed)", time.Since(t0).String()))
	}()

}

func (c *Core) addMovie(msg *pubsub.Message) {
	// I'm expecting movie to have the following fields filled
	// Title
	movie := msg.Payload.(*model.Movie)

	// Since it's coming in through this endpoint, we know it's a stub
	movie.File_Title = movie.Title
	movie.Resolution = "n/a"
	movie.FileType = "n/a"
	movie.Location = "n/a"
	movie.Stub = 1

	// 3 operations, rescrape, update and cache
	c.wg.Add(1)

	// rescrape
	scrape := &pubsub.Message{Payload: movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(scrape, "/command/movie/scrape")

	go func() {
		c.wg.Wait()
		msg.Reply <- movie
	}()
}

func (c *Core) doMovieFound(msg *pubsub.Message) {
	movie := msg.Payload.(*model.Movie)

	check := &pubsub.Message{Payload: movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(check, "/command/movie/exists")

	reply := <-check.Reply
	exists := reply.(bool)

	if exists {
		mlog.Info("SKIPPED: exists [%s] (%s)", movie.Title, movie.Location)
	} else {
		lib.Notify(c.bus, "import:progress", fmt.Sprintf("NEW: [%s] (%s)", movie.Title, movie.Location))

		c.wg.Add(1)
		c.bus.Pub(msg, "/command/movie/scrape")
	}
}

func (c *Core) fixMovie(msg *pubsub.Message) {
	movie := msg.Payload.(*model.Movie)

	// 3 operations, rescrape, update and cache
	c.wg.Add(1)

	// rescrape
	scrape := &pubsub.Message{Payload: movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(scrape, "/command/movie/rescrape")

	go func() {
		c.wg.Wait()
		msg.Reply <- movie
	}()
	// go c.waitFixMovie(msg.Reply, movie)
}

func (c *Core) doMovieTmdbNotFound(msg *pubsub.Message) {
	dto := msg.Payload.(*dto.Scrape)

	store := &pubsub.Message{Payload: dto.Movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(store, "/command/movie/partialstore")
}

func (c *Core) doMovieScraped(msg *pubsub.Message) {
	dto := msg.Payload.(*dto.Scrape)

	mlog.Info("ScrapeDTO: %+v", dto)

	// I treat the following two commands as one, for the sake of the wg
	// now there are two oustanding locks, which will be decreased by each
	// responding service
	c.wg.Add(1)

	store := &pubsub.Message{Payload: dto.Movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(store, "/command/movie/store")

	cache := &pubsub.Message{Payload: dto, Reply: make(chan interface{}, 3)}
	c.bus.Pub(cache, "/command/movie/cache")

	// mlog.Info("ScrapeDTO: %+v", dto)
}

func (c *Core) doMovieReScraped(msg *pubsub.Message) {
	dto := msg.Payload.(*dto.Scrape)

	mlog.Info("ScrapeDTO: %+v", dto)

	// I treat the following two commands as one, for the sake of the wg
	// now there are two oustanding locks, which will be decreased by each
	// responding service
	c.wg.Add(1)

	store := &pubsub.Message{Payload: dto.Movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(store, "/command/movie/update")

	cache := &pubsub.Message{Payload: dto, Reply: make(chan interface{}, 3)}
	c.bus.Pub(cache, "/command/movie/cache")

	// mlog.Info("ScrapeDTO: %+v", dto)
}

func (c *Core) pruneMovies(msg *pubsub.Message) {
	t0 := time.Now()

	lib.Notify(c.bus, "prune:begin", "Started Prune Process")

	options := &lib.Options{Offset: 0, Limit: 99999999999999, SortBy: "title", SortOrder: "asc"}
	all := &pubsub.Message{Payload: options, Reply: make(chan interface{}, capacity)}
	c.bus.Pub(all, "/get/movies")

	reply := <-all.Reply
	dto := reply.(*model.MoviesDTO)

	if c.settings.UnraidMode {
		for _, item := range dto.Items {
			// mlog.Info("Item is %s (%s)", item.Title, item.Location)

			index := strings.Index(item.Location, ":")

			host := item.Location[:index]
			location := item.Location[index+1:]

			req := client.NewRequest("io.jbrodriguez.mediagui.agent."+host, "Agent.Exists", &agent.ExistsReq{
				Location: location,
			})

			rsp := &agent.ExistsRsp{}

			if err := client.Call(context.Background(), req, rsp); err != nil {
				mlog.Warning("Unable to connect to service (%s): %s", "io.jbrodriguez.mediagui.agent."+host, err)
				// lib.Notify(s.bus, "import:progress", "Unable to connect to host "+host)
				// lib.Notify(s.bus, "import:end", "Import process finished")
				// return
				continue
			}

			if !rsp.Exists {
				lib.Notify(c.bus, "prune:selected", fmt.Sprintf("UP FOR DELETION: [%d] %s (%s))", item.Id, item.Title, item.Location))

				movie := &pubsub.Message{Payload: item, Reply: make(chan interface{}, capacity)}
				c.bus.Pub(movie, "/command/movie/delete")
			}
		}
	} else {
		for _, item := range dto.Items {
			matches := c.re.FindStringSubmatch(item.Location)
			if len(matches) == 0 {
				continue
			}

			folder := filepath.Join("/Volumes", matches[1])
			if !c.maps[folder] {
				// mlog.Info("Folder not mapped (%s): %s", folder, item.Location)
				continue
			}

			// mlog.Info("Folder mapped (%s): %s", folder, item.Location)

			if _, err := os.Stat(item.Location); err != nil {
				if os.IsNotExist(err) {
					lib.Notify(c.bus, "prune:selected", fmt.Sprintf("UP FOR DELETION: [%d] %s (%s))", item.Id, item.Title, item.Location))

					movie := &pubsub.Message{Payload: item, Reply: make(chan interface{}, capacity)}
					c.bus.Pub(movie, "/command/movie/delete")
				}
			}
		}
	}

	// for _, folder := range c.settings.MediaFolders {
	// 	mlog.Info("folder: %s", folder)
	// 	if _, err := os.Stat(folder); err != nil {
	// 		if os.IsNotExist(err) {
	// 			lib.Notify(c.bus, "prune:error", fmt.Sprintf("Folder %s is not present. Aborting Prune process.", folder))
	// 			return
	// 		}
	// 	}
	// }

	// lib.Notify(c.bus, "prune:end", "Finished Prune Process")
	lib.Notify(c.bus, "prune:end", fmt.Sprintf("Prune process finished (%s elapsed)", time.Since(t0).String()))

}

func (c *Core) addMediaFolder(msg *pubsub.Message) {
	folder := msg.Payload.(string)
	mlog.Info("addMediaFolder: %s", folder)

	c.settings.MediaFolders = append(c.settings.MediaFolders, folder)
	c.settings.Save()

	cfg := &pubsub.Message{Payload: c.settings}
	c.bus.Pub(cfg, "/event/config/changed")

	msg.Reply <- &c.settings.Config
	// mlog.Info("Sent config")
}

// func (c *Core) waitFixMovie(ch chan interface{}, movie *model.Movie) {
// 	c.wg.Wait()
// 	ch <- movie
// }

// func (c *Core) doMovieUpdated(msg *pubsub.Message) {
// 	c.wg.Done()
// }

// func (c *Core) doMovieCachedForced(msg *pubsub.Message) {
// 	c.wg.Done()
// }

func (c *Core) doWorkUnitDone(msg *pubsub.Message) {
	c.wg.Done()
}
<commit_msg>Account for stub items when pruning<commit_after>package services

import (
	"jbrodriguez/mediagui/server/src/dto"
	"jbrodriguez/mediagui/server/src/lib"
	"jbrodriguez/mediagui/server/src/model"
	"jbrodriguez/mediagui/server/src/proto"

	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/jbrodriguez/actor"
	"github.com/jbrodriguez/mlog"
	"github.com/jbrodriguez/pubsub"
	"github.com/micro/go-micro/client"

	"golang.org/x/net/context"
)

// Core -
type Core struct {
	bus      *pubsub.PubSub
	settings *lib.Settings

	actor *actor.Actor
	re    *regexp.Regexp
	maps  map[string]bool

	wg sync.WaitGroup
}

// NewCore -
func NewCore(bus *pubsub.PubSub, settings *lib.Settings) *Core {
	core := &Core{
		bus:      bus,
		settings: settings,
		actor:    actor.NewActor(bus),
	}
	return core
}

// Start -
func (c *Core) Start() {
	mlog.Info("Starting service Core ...")

	c.actor.Register("/get/config", c.getConfig)
	c.actor.Register("/post/import", c.importMovies)
	c.actor.Register("/post/add", c.addMovie)
	c.actor.Register("/post/prune", c.pruneMovies)
	c.actor.Register("/put/config/folder", c.addMediaFolder)
	c.actor.Register("/put/movies/fix", c.fixMovie)

	c.actor.Register("/event/movie/found", c.doMovieFound)
	c.actor.Register("/event/movie/tmdbnotfound", c.doMovieTmdbNotFound)
	c.actor.Register("/event/movie/scraped", c.doMovieScraped)
	c.actor.Register("/event/movie/rescraped", c.doMovieReScraped)
	// c.actor.Register("/event/movie/updated", c.doMovieUpdated)
	// c.actor.Register("/event/movie/cached/forced", c.doMovieCachedForced)

	c.actor.Register("/event/workunit/done", c.doWorkUnitDone)

	c.re = regexp.MustCompile(`(?i)/Volumes/(.*?)/.*`)
	c.maps = make(map[string]bool)

	// maps["/Volumes/wopr-films"] = true // for example
	for _, folder := range c.settings.MediaFolders {
		c.maps[folder] = false
		if _, err := os.Stat(folder); err == nil {
			c.maps[folder] = true
		}
	}

	go c.actor.React()
}

// Stop -
func (c *Core) Stop() {
	mlog.Info("Stopped service Core ...")
}

func (c *Core) getConfig(msg *pubsub.Message) {
	msg.Reply <- &c.settings.Config
	mlog.Info("Sent config")
}

func (c *Core) importMovies(msg *pubsub.Message) {
	t0 := time.Now()
	// mlog.Info("Begin movie scanning ...")
	lib.Notify(c.bus, "import:begin", "Import process started")

	c.wg.Add(1)

	c.bus.Pub(nil, "/command/movie/scan")
	//	msg.Reply <- &c.settings.Config
	// mlog.Info("Import finished")

	go func() {
		c.wg.Wait()
		// 𝛥t := float64(time.Since(t0)) / 1e9
		lib.Notify(c.bus, "import:end", fmt.Sprintf("Import process finished (%s elapsed)", time.Since(t0).String()))
	}()

}

func (c *Core) addMovie(msg *pubsub.Message) {
	// I'm expecting movie to have the following fields filled
	// Title
	movie := msg.Payload.(*model.Movie)

	// Since it's coming in through this endpoint, we know it's a stub
	movie.File_Title = movie.Title
	movie.Resolution = "n/a"
	movie.FileType = "n/a"
	movie.Location = "n/a"
	movie.Stub = 1

	// 3 operations, rescrape, update and cache
	c.wg.Add(1)

	// rescrape
	scrape := &pubsub.Message{Payload: movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(scrape, "/command/movie/scrape")

	go func() {
		c.wg.Wait()
		msg.Reply <- movie
	}()
}

func (c *Core) doMovieFound(msg *pubsub.Message) {
	movie := msg.Payload.(*model.Movie)

	check := &pubsub.Message{Payload: movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(check, "/command/movie/exists")

	reply := <-check.Reply
	exists := reply.(bool)

	if exists {
		mlog.Info("SKIPPED: exists [%s] (%s)", movie.Title, movie.Location)
	} else {
		lib.Notify(c.bus, "import:progress", fmt.Sprintf("NEW: [%s] (%s)", movie.Title, movie.Location))

		c.wg.Add(1)
		c.bus.Pub(msg, "/command/movie/scrape")
	}
}

func (c *Core) fixMovie(msg *pubsub.Message) {
	movie := msg.Payload.(*model.Movie)

	// 3 operations, rescrape, update and cache
	c.wg.Add(1)

	// rescrape
	scrape := &pubsub.Message{Payload: movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(scrape, "/command/movie/rescrape")

	go func() {
		c.wg.Wait()
		msg.Reply <- movie
	}()
	// go c.waitFixMovie(msg.Reply, movie)
}

func (c *Core) doMovieTmdbNotFound(msg *pubsub.Message) {
	dto := msg.Payload.(*dto.Scrape)

	store := &pubsub.Message{Payload: dto.Movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(store, "/command/movie/partialstore")
}

func (c *Core) doMovieScraped(msg *pubsub.Message) {
	dto := msg.Payload.(*dto.Scrape)

	mlog.Info("ScrapeDTO: %+v", dto)

	// I treat the following two commands as one, for the sake of the wg
	// now there are two oustanding locks, which will be decreased by each
	// responding service
	c.wg.Add(1)

	store := &pubsub.Message{Payload: dto.Movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(store, "/command/movie/store")

	cache := &pubsub.Message{Payload: dto, Reply: make(chan interface{}, 3)}
	c.bus.Pub(cache, "/command/movie/cache")

	// mlog.Info("ScrapeDTO: %+v", dto)
}

func (c *Core) doMovieReScraped(msg *pubsub.Message) {
	dto := msg.Payload.(*dto.Scrape)

	mlog.Info("ScrapeDTO: %+v", dto)

	// I treat the following two commands as one, for the sake of the wg
	// now there are two oustanding locks, which will be decreased by each
	// responding service
	c.wg.Add(1)

	store := &pubsub.Message{Payload: dto.Movie, Reply: make(chan interface{}, 3)}
	c.bus.Pub(store, "/command/movie/update")

	cache := &pubsub.Message{Payload: dto, Reply: make(chan interface{}, 3)}
	c.bus.Pub(cache, "/command/movie/cache")

	// mlog.Info("ScrapeDTO: %+v", dto)
}

func (c *Core) pruneMovies(msg *pubsub.Message) {
	t0 := time.Now()

	lib.Notify(c.bus, "prune:begin", "Started Prune Process")

	options := &lib.Options{Offset: 0, Limit: 99999999999999, SortBy: "title", SortOrder: "asc"}
	all := &pubsub.Message{Payload: options, Reply: make(chan interface{}, capacity)}
	c.bus.Pub(all, "/get/movies")

	reply := <-all.Reply
	dto := reply.(*model.MoviesDTO)

	if c.settings.UnraidMode {
		for _, item := range dto.Items {
			// mlog.Info("Item is %s (%s)", item.Title, item.Location)

			index := strings.Index(item.Location, ":")
			if index == -1 {
				// a valid location is wopr:/mnt/user/films/bluray/22 Bullets (2010)/22.Bullets_BLURAY.iso
				// if a ':' isn't found, then this must be a stub
				continue
			}

			host := item.Location[:index]
			location := item.Location[index+1:]

			req := client.NewRequest("io.jbrodriguez.mediagui.agent."+host, "Agent.Exists", &agent.ExistsReq{
				Location: location,
			})

			rsp := &agent.ExistsRsp{}

			if err := client.Call(context.Background(), req, rsp); err != nil {
				mlog.Warning("Unable to connect to service (%s): %s", "io.jbrodriguez.mediagui.agent."+host, err)
				// lib.Notify(s.bus, "import:progress", "Unable to connect to host "+host)
				// lib.Notify(s.bus, "import:end", "Import process finished")
				// return
				continue
			}

			if !rsp.Exists {
				lib.Notify(c.bus, "prune:selected", fmt.Sprintf("UP FOR DELETION: [%d] %s (%s))", item.Id, item.Title, item.Location))

				movie := &pubsub.Message{Payload: item, Reply: make(chan interface{}, capacity)}
				c.bus.Pub(movie, "/command/movie/delete")
			}
		}
	} else {
		for _, item := range dto.Items {
			matches := c.re.FindStringSubmatch(item.Location)
			if len(matches) == 0 {
				continue
			}

			folder := filepath.Join("/Volumes", matches[1])
			if !c.maps[folder] {
				// mlog.Info("Folder not mapped (%s): %s", folder, item.Location)
				continue
			}

			// mlog.Info("Folder mapped (%s): %s", folder, item.Location)

			if _, err := os.Stat(item.Location); err != nil {
				if os.IsNotExist(err) {
					lib.Notify(c.bus, "prune:selected", fmt.Sprintf("UP FOR DELETION: [%d] %s (%s))", item.Id, item.Title, item.Location))

					movie := &pubsub.Message{Payload: item, Reply: make(chan interface{}, capacity)}
					c.bus.Pub(movie, "/command/movie/delete")
				}
			}
		}
	}

	// for _, folder := range c.settings.MediaFolders {
	// 	mlog.Info("folder: %s", folder)
	// 	if _, err := os.Stat(folder); err != nil {
	// 		if os.IsNotExist(err) {
	// 			lib.Notify(c.bus, "prune:error", fmt.Sprintf("Folder %s is not present. Aborting Prune process.", folder))
	// 			return
	// 		}
	// 	}
	// }

	// lib.Notify(c.bus, "prune:end", "Finished Prune Process")
	lib.Notify(c.bus, "prune:end", fmt.Sprintf("Prune process finished (%s elapsed)", time.Since(t0).String()))

}

func (c *Core) addMediaFolder(msg *pubsub.Message) {
	folder := msg.Payload.(string)
	mlog.Info("addMediaFolder: %s", folder)

	c.settings.MediaFolders = append(c.settings.MediaFolders, folder)
	c.settings.Save()

	cfg := &pubsub.Message{Payload: c.settings}
	c.bus.Pub(cfg, "/event/config/changed")

	msg.Reply <- &c.settings.Config
	// mlog.Info("Sent config")
}

// func (c *Core) waitFixMovie(ch chan interface{}, movie *model.Movie) {
// 	c.wg.Wait()
// 	ch <- movie
// }

// func (c *Core) doMovieUpdated(msg *pubsub.Message) {
// 	c.wg.Done()
// }

// func (c *Core) doMovieCachedForced(msg *pubsub.Message) {
// 	c.wg.Done()
// }

func (c *Core) doWorkUnitDone(msg *pubsub.Message) {
	c.wg.Done()
}
<|endoftext|>


<commit_before>package term

import (
	"syscall"
	"unsafe"
)

const (
	getTermios = syscall.TIOCGETA
	setTermios = syscall.TIOCSETA

	ECHO   = 0x00000008
	ONLCR  = 0x2
	ISTRIP = 0x20
	INLCR  = 0x40
	ISIG   = 0x80
	IGNCR  = 0x80
	ICANON = 0x100
	ICRNL  = 0x100
	IXOFF  = 0x400
	IXON   = 0x200
)

type Termios struct {
	Iflag  uint64
	Oflag  uint64
	Cflag  uint64
	Lflag  uint64
	Cc     [20]byte
	Ispeed uint64
	Ospeed uint64
}

// MakeRaw put the terminal connected to the given file descriptor into raw
// mode and returns the previous state of the terminal so that it can be
// restored.
func MakeRaw(fd uintptr) (*State, error) {
	var oldState State
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(getTermios), uintptr(unsafe.Pointer(&oldState.termios))); err != 0 {
		return nil, err
	}

	newState := oldState.termios
	newState.Iflag &^= (ISTRIP | INLCR | IGNCR | IXON | IXOFF)
	newState.Iflag |= ICRNL
	newState.Oflag |= ONLCR
	newState.Lflag &^= (ECHO | ICANON | ISIG)

	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(setTermios), uintptr(unsafe.Pointer(&newState))); err != 0 {
		return nil, err
	}

	return &oldState, nil
}
<commit_msg>Use BSD raw mode on darwin. Fixes nano, tmux and others<commit_after>package term

import (
	"syscall"
	"unsafe"
)

const (
	getTermios = syscall.TIOCGETA
	setTermios = syscall.TIOCSETA

	IGNBRK = syscall.IGNBRK
	PARMRK = syscall.PARMRK
	INLCR  = syscall.INLCR
	IGNCR  = syscall.IGNCR
	ECHONL = syscall.ECHONL
	CSIZE  = syscall.CSIZE
	ICRNL  = syscall.ICRNL
	ISTRIP = syscall.ISTRIP
	PARENB = syscall.PARENB
	ECHO   = syscall.ECHO
	ICANON = syscall.ICANON
	ISIG   = syscall.ISIG
	IXON   = syscall.IXON
	BRKINT = syscall.BRKINT
	INPCK  = syscall.INPCK
	OPOST  = syscall.OPOST
	CS8    = syscall.CS8
	IEXTEN = syscall.IEXTEN
)

type Termios struct {
	Iflag  uint64
	Oflag  uint64
	Cflag  uint64
	Lflag  uint64
	Cc     [20]byte
	Ispeed uint64
	Ospeed uint64
}

// MakeRaw put the terminal connected to the given file descriptor into raw
// mode and returns the previous state of the terminal so that it can be
// restored.
func MakeRaw(fd uintptr) (*State, error) {
	var oldState State
	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(getTermios), uintptr(unsafe.Pointer(&oldState.termios))); err != 0 {
		return nil, err
	}

	newState := oldState.termios
	newState.Iflag &^= (IGNBRK | BRKINT | PARMRK | ISTRIP | INLCR | IGNCR | ICRNL | IXON)
	newState.Oflag &^= OPOST
	newState.Lflag &^= (ECHO | ECHONL | ICANON | ISIG | IEXTEN)
	newState.Cflag &^= (CSIZE | PARENB)
	newState.Cflag |= CS8
	newState.Cc[syscall.VMIN] = 1
	newState.Cc[syscall.VTIME] = 0

	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd, uintptr(setTermios), uintptr(unsafe.Pointer(&newState))); err != 0 {
		return nil, err
	}

	return &oldState, nil
}
<|endoftext|>


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

package tracing

import (
	"bytes"
	"encoding/gob"
	"fmt"

	"golang.org/x/net/context"
	"golang.org/x/net/trace"

	"github.com/cockroachdb/cockroach/pkg/util/caller"
	"github.com/cockroachdb/cockroach/pkg/util/envutil"
	"github.com/lightstep/lightstep-tracer-go"
	basictracer "github.com/opentracing/basictracer-go"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	otlog "github.com/opentracing/opentracing-go/log"
)

// Snowball is set as Baggage on traces which are used for snowball tracing.
const Snowball = "sb"

// maxLogsPerSpan limits the number of logs in a Span; use a comfortable limit.
const maxLogsPerSpan = 1000

// A CallbackRecorder immediately invokes itself on received trace spans.
type CallbackRecorder func(sp basictracer.RawSpan)

// RecordSpan implements basictracer.SpanRecorder.
func (cr CallbackRecorder) RecordSpan(sp basictracer.RawSpan) {
	cr(sp)
}

// JoinOrNew creates a new Span joined to the provided DelegatingCarrier or
// creates Span from the given tracer.
func JoinOrNew(
	tr opentracing.Tracer, carrier *SpanContextCarrier, opName string,
) (opentracing.Span, error) {
	if carrier != nil {
		wireContext, err := tr.Extract(basictracer.Delegator, carrier)
		switch err {
		case nil:
			sp := tr.StartSpan(opName, opentracing.FollowsFrom(wireContext))

			// Copy baggage items to tags so they show up in the Lightstep UI.
			sp.Context().ForeachBaggageItem(func(k, v string) bool { sp.SetTag(k, v); return true })

			sp.LogEvent(opName)
			return sp, nil
		case opentracing.ErrSpanContextNotFound:
		default:
			return nil, err
		}
	}
	return tr.StartSpan(opName), nil
}

// JoinOrNewSnowball returns a Span which records directly via the specified
// callback. If the given DelegatingCarrier is nil, a new Span is created.
// otherwise, the created Span is a child.
//
// The recorder should be nil if we don't need to record spans.
//
// TODO(andrei): JoinOrNewSnowball creates a new tracer, which is not kosher.
// Also this can't use the lightstep tracer.
func JoinOrNewSnowball(
	opName string, carrier *SpanContextCarrier, recorder func(sp basictracer.RawSpan),
) (opentracing.Span, error) {
	tr := basictracer.NewWithOptions(basictracerOptions(recorder))
	sp, err := JoinOrNew(tr, carrier, opName)
	if err == nil {
		// We definitely want to sample a Snowball trace.
		// This must be set *before* SetBaggageItem, as that will otherwise be ignored.
		ext.SamplingPriority.Set(sp, 1)
		sp.SetBaggageItem(Snowball, "1")
	}
	return sp, err
}

// NewTracerAndSpanFor7881 creates a new tracer and a root span. The tracer is
// to be used for tracking down #7881; it runs a callback for each finished span
// (and the callback used accumulates the spans in a SQL txn).
func NewTracerAndSpanFor7881(
	callback func(sp basictracer.RawSpan),
) (opentracing.Span, opentracing.Tracer, error) {
	opts := basictracerOptions(callback)
	// Don't trim the logs in "unsampled" spans". Note that this tracer does not
	// use sampling; instead it uses an ad-hoc mechanism for marking spans of
	// interest.
	opts.TrimUnsampledSpans = false
	tr := basictracer.NewWithOptions(opts)
	sp, err := JoinOrNew(tr, nil, "sql txn")
	return sp, tr, err
}

// FinishSpan closes the given span (if not nil). It is a convenience wrapper
// for span.Finish() which tolerates nil spans.
func FinishSpan(span opentracing.Span) {
	if span != nil {
		span.Finish()
	}
}

// ForkCtxSpan checks if ctx has a Span open; if it does, it creates a new Span
// that follows from the original Span. This allows the resulting context to be
// used in an async task that might outlive the original operation.
//
// Returns the new context and the new span (if any). The span should be
// closed via FinishSpan.
func ForkCtxSpan(ctx context.Context, opName string) (context.Context, opentracing.Span) {
	if span := opentracing.SpanFromContext(ctx); span != nil {
		if span.BaggageItem(Snowball) == "1" {
			// If we are doing snowball tracing, the span might outlive the snowball
			// tracer (calling the record function when it is no longer legal to do
			// so). Return a context with no span in this case.
			return opentracing.ContextWithSpan(ctx, nil), nil
		}
		tr := span.Tracer()
		newSpan := tr.StartSpan(opName, opentracing.FollowsFrom(span.Context()))
		return opentracing.ContextWithSpan(ctx, newSpan), newSpan
	}
	return ctx, nil
}

// ChildSpan opens a span as a child of the current span in the context (if
// there is one).
//
// Returns the new context and the new span (if any). The span should be
// closed via FinishSpan.
func ChildSpan(ctx context.Context, opName string) (context.Context, opentracing.Span) {
	span := opentracing.SpanFromContext(ctx)
	if span == nil {
		return ctx, nil
	}
	newSpan := span.Tracer().StartSpan(opName, opentracing.ChildOf(span.Context()))
	return opentracing.ContextWithSpan(ctx, newSpan), newSpan
}

// netTraceIntegrator is passed into basictracer as NewSpanEventListener
// and causes all traces to be registered with the net/trace endpoint.
func netTraceIntegrator() func(basictracer.SpanEvent) {
	var tr trace.Trace
	return func(e basictracer.SpanEvent) {
		switch t := e.(type) {
		case basictracer.EventCreate:
			tr = trace.New("tracing", t.OperationName)
			tr.SetMaxEvents(maxLogsPerSpan)
		case basictracer.EventFinish:
			tr.Finish()
		case basictracer.EventTag:
			tr.LazyPrintf("%s:%v", t.Key, t.Value)
		case basictracer.EventLogFields:
			var buf bytes.Buffer
			for i, f := range t.Fields {
				if i > 0 {
					buf.WriteByte(' ')
				}
				fmt.Fprintf(&buf, "%s:%v", f.Key(), f.Value())
			}

			tr.LazyPrintf("%s", buf.String())
		case basictracer.EventLog:
			if t.Payload != nil {
				tr.LazyPrintf("%s (payload %v)", t.Event, t.Payload)
			} else {
				tr.LazyPrintf("%s", t.Event)
			}
		}
	}
}

// basicTracerOptions initializes options for basictracer.
// The recorder should be nil if we don't need to record spans.
func basictracerOptions(recorder func(basictracer.RawSpan)) basictracer.Options {
	opts := basictracer.DefaultOptions()
	opts.ShouldSample = func(traceID uint64) bool { return false }
	opts.TrimUnsampledSpans = true
	opts.NewSpanEventListener = netTraceIntegrator
	opts.DebugAssertUseAfterFinish = true // provoke crash on use-after-Finish
	if recorder == nil {
		opts.Recorder = CallbackRecorder(func(_ basictracer.RawSpan) {})
		// If we are not recording the spans, there is no need to keep them in
		// memory. Events still get passed to netTraceIntegrator.
		opts.DropAllLogs = true
	} else {
		opts.Recorder = CallbackRecorder(recorder)
		// Set a comfortable limit of log events per span.
		opts.MaxLogsPerSpan = maxLogsPerSpan
	}
	return opts
}

var lightstepToken = envutil.EnvOrDefaultString("COCKROACH_LIGHTSTEP_TOKEN", "")

// By default, if a lightstep token is specified we trace to both Lightstep and
// net/trace. If this flag is enabled, we will only trace to Lightstep.
var lightstepOnly = envutil.EnvOrDefaultBool("COCKROACH_LIGHTSTEP_ONLY", false)

// newTracer implements NewTracer and allows that function to be mocked out via Disable().
var newTracer = func() opentracing.Tracer {
	if lightstepToken != "" {
		lsTr := lightstep.NewTracer(lightstep.Options{
			AccessToken:    lightstepToken,
			MaxLogsPerSpan: maxLogsPerSpan,
			UseGRPC:        true,
		})
		if lightstepOnly {
			return lsTr
		}
		basicTr := basictracer.NewWithOptions(basictracerOptions(nil))
		// The TeeTracer uses the first tracer for serialization of span contexts;
		// lightspan needs to be first because it correlates spans between nodes.
		return NewTeeTracer(lsTr, basicTr)
	}
	return basictracer.NewWithOptions(basictracerOptions(nil))
}

// NewTracer creates a Tracer which records to the net/trace
// endpoint.
func NewTracer() opentracing.Tracer {
	return newTracer()
}

// EnsureContext checks whether the given context.Context contains a Span. If
// not, it creates one using the provided Tracer and wraps it in the returned
// Span. The returned closure must be called after the request has been fully
// processed.
func EnsureContext(ctx context.Context, tracer opentracing.Tracer) (context.Context, func()) {
	_, _, funcName := caller.Lookup(1)
	if opentracing.SpanFromContext(ctx) == nil {
		sp := tracer.StartSpan(funcName)
		return opentracing.ContextWithSpan(ctx, sp), sp.Finish
	}
	return ctx, func() {}
}

// Disable is for benchmarking use and causes all future tracers to deal in
// no-ops. Calling the returned closure undoes this effect. There is no
// synchronization, so no moving parts are allowed while Disable and the
// closure are called.
func Disable() func() {
	orig := newTracer
	newTracer = func() opentracing.Tracer { return opentracing.NoopTracer{} }
	return func() {
		newTracer = orig
	}
}

// EncodeRawSpan encodes a raw span into bytes, using the given dest slice
// as a buffer.
func EncodeRawSpan(rawSpan *basictracer.RawSpan, dest []byte) ([]byte, error) {
	// This is not a greatly efficient (but convenient) use of gob.
	buf := bytes.NewBuffer(dest[:0])
	e := gob.NewEncoder(buf)
	var err error

	encode := func(arg interface{}) {
		if err == nil {
			err = e.Encode(arg)
		}
	}

	// We cannot use gob for the Logs (because Field doesn't export the necessary
	// fields). We use it for the other fields.
	encode(rawSpan.Context)
	encode(rawSpan.ParentSpanID)
	encode(rawSpan.Operation)
	encode(rawSpan.Start)
	encode(rawSpan.Duration)
	encode(rawSpan.Tags)

	// Encode the number of LogRecords, then the records.
	encode(int32(len(rawSpan.Logs)))
	for _, lr := range rawSpan.Logs {
		encode(lr.Timestamp)
		// Encode the number of Fields.
		encode(int32(len(lr.Fields)))
		for _, f := range lr.Fields {
			encode(f.Key())
			// Encode the field value as a string.
			encode(fmt.Sprint(f.Value()))
		}
	}

	return buf.Bytes(), err
}

// DecodeRawSpan unmarshals into the given RawSpan.
func DecodeRawSpan(enc []byte, dest *basictracer.RawSpan) error {
	d := gob.NewDecoder(bytes.NewBuffer(enc))
	var err error
	decode := func(arg interface{}) {
		if err == nil {
			err = d.Decode(arg)
		}
	}

	decode(&dest.Context)
	decode(&dest.ParentSpanID)
	decode(&dest.Operation)
	decode(&dest.Start)
	decode(&dest.Duration)
	decode(&dest.Tags)

	var numLogs int32
	decode(&numLogs)
	dest.Logs = make([]opentracing.LogRecord, numLogs)
	for i := range dest.Logs {
		lr := &dest.Logs[i]
		decode(&lr.Timestamp)

		var numFields int32
		decode(&numFields)
		lr.Fields = make([]otlog.Field, numFields)
		for j := range lr.Fields {
			var key, val string
			decode(&key)
			decode(&val)
			lr.Fields[j] = otlog.String(key, val)
		}
	}

	return err
}

// contextTracerKeyType is an empty type for the handle associated with the
// tracer value (see context.Value).
type contextTracerKeyType struct{}

// WithTracer returns a context derived from the given context, for which
// TracerFromCtx returns the given tracer.
func WithTracer(ctx context.Context, tracer opentracing.Tracer) context.Context {
	return context.WithValue(ctx, contextTracerKeyType{}, tracer)
}

// TracerFromCtx returns the tracer set on the context (or a parent context) via
// WithTracer.
func TracerFromCtx(ctx context.Context) opentracing.Tracer {
	if tracerVal := ctx.Value(contextTracerKeyType{}); tracerVal != nil {
		return tracerVal.(opentracing.Tracer)
	}
	return nil
}
<commit_msg>Revert "util/tracing: use GPRC for lightstep reporting"<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Tobias Schottdorf (tobias.schottdorf@gmail.com)

package tracing

import (
	"bytes"
	"encoding/gob"
	"fmt"

	"golang.org/x/net/context"
	"golang.org/x/net/trace"

	"github.com/cockroachdb/cockroach/pkg/util/caller"
	"github.com/cockroachdb/cockroach/pkg/util/envutil"
	"github.com/lightstep/lightstep-tracer-go"
	basictracer "github.com/opentracing/basictracer-go"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	otlog "github.com/opentracing/opentracing-go/log"
)

// Snowball is set as Baggage on traces which are used for snowball tracing.
const Snowball = "sb"

// maxLogsPerSpan limits the number of logs in a Span; use a comfortable limit.
const maxLogsPerSpan = 1000

// A CallbackRecorder immediately invokes itself on received trace spans.
type CallbackRecorder func(sp basictracer.RawSpan)

// RecordSpan implements basictracer.SpanRecorder.
func (cr CallbackRecorder) RecordSpan(sp basictracer.RawSpan) {
	cr(sp)
}

// JoinOrNew creates a new Span joined to the provided DelegatingCarrier or
// creates Span from the given tracer.
func JoinOrNew(
	tr opentracing.Tracer, carrier *SpanContextCarrier, opName string,
) (opentracing.Span, error) {
	if carrier != nil {
		wireContext, err := tr.Extract(basictracer.Delegator, carrier)
		switch err {
		case nil:
			sp := tr.StartSpan(opName, opentracing.FollowsFrom(wireContext))

			// Copy baggage items to tags so they show up in the Lightstep UI.
			sp.Context().ForeachBaggageItem(func(k, v string) bool { sp.SetTag(k, v); return true })

			sp.LogEvent(opName)
			return sp, nil
		case opentracing.ErrSpanContextNotFound:
		default:
			return nil, err
		}
	}
	return tr.StartSpan(opName), nil
}

// JoinOrNewSnowball returns a Span which records directly via the specified
// callback. If the given DelegatingCarrier is nil, a new Span is created.
// otherwise, the created Span is a child.
//
// The recorder should be nil if we don't need to record spans.
//
// TODO(andrei): JoinOrNewSnowball creates a new tracer, which is not kosher.
// Also this can't use the lightstep tracer.
func JoinOrNewSnowball(
	opName string, carrier *SpanContextCarrier, recorder func(sp basictracer.RawSpan),
) (opentracing.Span, error) {
	tr := basictracer.NewWithOptions(basictracerOptions(recorder))
	sp, err := JoinOrNew(tr, carrier, opName)
	if err == nil {
		// We definitely want to sample a Snowball trace.
		// This must be set *before* SetBaggageItem, as that will otherwise be ignored.
		ext.SamplingPriority.Set(sp, 1)
		sp.SetBaggageItem(Snowball, "1")
	}
	return sp, err
}

// NewTracerAndSpanFor7881 creates a new tracer and a root span. The tracer is
// to be used for tracking down #7881; it runs a callback for each finished span
// (and the callback used accumulates the spans in a SQL txn).
func NewTracerAndSpanFor7881(
	callback func(sp basictracer.RawSpan),
) (opentracing.Span, opentracing.Tracer, error) {
	opts := basictracerOptions(callback)
	// Don't trim the logs in "unsampled" spans". Note that this tracer does not
	// use sampling; instead it uses an ad-hoc mechanism for marking spans of
	// interest.
	opts.TrimUnsampledSpans = false
	tr := basictracer.NewWithOptions(opts)
	sp, err := JoinOrNew(tr, nil, "sql txn")
	return sp, tr, err
}

// FinishSpan closes the given span (if not nil). It is a convenience wrapper
// for span.Finish() which tolerates nil spans.
func FinishSpan(span opentracing.Span) {
	if span != nil {
		span.Finish()
	}
}

// ForkCtxSpan checks if ctx has a Span open; if it does, it creates a new Span
// that follows from the original Span. This allows the resulting context to be
// used in an async task that might outlive the original operation.
//
// Returns the new context and the new span (if any). The span should be
// closed via FinishSpan.
func ForkCtxSpan(ctx context.Context, opName string) (context.Context, opentracing.Span) {
	if span := opentracing.SpanFromContext(ctx); span != nil {
		if span.BaggageItem(Snowball) == "1" {
			// If we are doing snowball tracing, the span might outlive the snowball
			// tracer (calling the record function when it is no longer legal to do
			// so). Return a context with no span in this case.
			return opentracing.ContextWithSpan(ctx, nil), nil
		}
		tr := span.Tracer()
		newSpan := tr.StartSpan(opName, opentracing.FollowsFrom(span.Context()))
		return opentracing.ContextWithSpan(ctx, newSpan), newSpan
	}
	return ctx, nil
}

// ChildSpan opens a span as a child of the current span in the context (if
// there is one).
//
// Returns the new context and the new span (if any). The span should be
// closed via FinishSpan.
func ChildSpan(ctx context.Context, opName string) (context.Context, opentracing.Span) {
	span := opentracing.SpanFromContext(ctx)
	if span == nil {
		return ctx, nil
	}
	newSpan := span.Tracer().StartSpan(opName, opentracing.ChildOf(span.Context()))
	return opentracing.ContextWithSpan(ctx, newSpan), newSpan
}

// netTraceIntegrator is passed into basictracer as NewSpanEventListener
// and causes all traces to be registered with the net/trace endpoint.
func netTraceIntegrator() func(basictracer.SpanEvent) {
	var tr trace.Trace
	return func(e basictracer.SpanEvent) {
		switch t := e.(type) {
		case basictracer.EventCreate:
			tr = trace.New("tracing", t.OperationName)
			tr.SetMaxEvents(maxLogsPerSpan)
		case basictracer.EventFinish:
			tr.Finish()
		case basictracer.EventTag:
			tr.LazyPrintf("%s:%v", t.Key, t.Value)
		case basictracer.EventLogFields:
			var buf bytes.Buffer
			for i, f := range t.Fields {
				if i > 0 {
					buf.WriteByte(' ')
				}
				fmt.Fprintf(&buf, "%s:%v", f.Key(), f.Value())
			}

			tr.LazyPrintf("%s", buf.String())
		case basictracer.EventLog:
			if t.Payload != nil {
				tr.LazyPrintf("%s (payload %v)", t.Event, t.Payload)
			} else {
				tr.LazyPrintf("%s", t.Event)
			}
		}
	}
}

// basicTracerOptions initializes options for basictracer.
// The recorder should be nil if we don't need to record spans.
func basictracerOptions(recorder func(basictracer.RawSpan)) basictracer.Options {
	opts := basictracer.DefaultOptions()
	opts.ShouldSample = func(traceID uint64) bool { return false }
	opts.TrimUnsampledSpans = true
	opts.NewSpanEventListener = netTraceIntegrator
	opts.DebugAssertUseAfterFinish = true // provoke crash on use-after-Finish
	if recorder == nil {
		opts.Recorder = CallbackRecorder(func(_ basictracer.RawSpan) {})
		// If we are not recording the spans, there is no need to keep them in
		// memory. Events still get passed to netTraceIntegrator.
		opts.DropAllLogs = true
	} else {
		opts.Recorder = CallbackRecorder(recorder)
		// Set a comfortable limit of log events per span.
		opts.MaxLogsPerSpan = maxLogsPerSpan
	}
	return opts
}

var lightstepToken = envutil.EnvOrDefaultString("COCKROACH_LIGHTSTEP_TOKEN", "")

// By default, if a lightstep token is specified we trace to both Lightstep and
// net/trace. If this flag is enabled, we will only trace to Lightstep.
var lightstepOnly = envutil.EnvOrDefaultBool("COCKROACH_LIGHTSTEP_ONLY", false)

// newTracer implements NewTracer and allows that function to be mocked out via Disable().
var newTracer = func() opentracing.Tracer {
	if lightstepToken != "" {
		lsTr := lightstep.NewTracer(lightstep.Options{
			AccessToken:    lightstepToken,
			MaxLogsPerSpan: maxLogsPerSpan,
		})
		if lightstepOnly {
			return lsTr
		}
		basicTr := basictracer.NewWithOptions(basictracerOptions(nil))
		// The TeeTracer uses the first tracer for serialization of span contexts;
		// lightspan needs to be first because it correlates spans between nodes.
		return NewTeeTracer(lsTr, basicTr)
	}
	return basictracer.NewWithOptions(basictracerOptions(nil))
}

// NewTracer creates a Tracer which records to the net/trace
// endpoint.
func NewTracer() opentracing.Tracer {
	return newTracer()
}

// EnsureContext checks whether the given context.Context contains a Span. If
// not, it creates one using the provided Tracer and wraps it in the returned
// Span. The returned closure must be called after the request has been fully
// processed.
func EnsureContext(ctx context.Context, tracer opentracing.Tracer) (context.Context, func()) {
	_, _, funcName := caller.Lookup(1)
	if opentracing.SpanFromContext(ctx) == nil {
		sp := tracer.StartSpan(funcName)
		return opentracing.ContextWithSpan(ctx, sp), sp.Finish
	}
	return ctx, func() {}
}

// Disable is for benchmarking use and causes all future tracers to deal in
// no-ops. Calling the returned closure undoes this effect. There is no
// synchronization, so no moving parts are allowed while Disable and the
// closure are called.
func Disable() func() {
	orig := newTracer
	newTracer = func() opentracing.Tracer { return opentracing.NoopTracer{} }
	return func() {
		newTracer = orig
	}
}

// EncodeRawSpan encodes a raw span into bytes, using the given dest slice
// as a buffer.
func EncodeRawSpan(rawSpan *basictracer.RawSpan, dest []byte) ([]byte, error) {
	// This is not a greatly efficient (but convenient) use of gob.
	buf := bytes.NewBuffer(dest[:0])
	e := gob.NewEncoder(buf)
	var err error

	encode := func(arg interface{}) {
		if err == nil {
			err = e.Encode(arg)
		}
	}

	// We cannot use gob for the Logs (because Field doesn't export the necessary
	// fields). We use it for the other fields.
	encode(rawSpan.Context)
	encode(rawSpan.ParentSpanID)
	encode(rawSpan.Operation)
	encode(rawSpan.Start)
	encode(rawSpan.Duration)
	encode(rawSpan.Tags)

	// Encode the number of LogRecords, then the records.
	encode(int32(len(rawSpan.Logs)))
	for _, lr := range rawSpan.Logs {
		encode(lr.Timestamp)
		// Encode the number of Fields.
		encode(int32(len(lr.Fields)))
		for _, f := range lr.Fields {
			encode(f.Key())
			// Encode the field value as a string.
			encode(fmt.Sprint(f.Value()))
		}
	}

	return buf.Bytes(), err
}

// DecodeRawSpan unmarshals into the given RawSpan.
func DecodeRawSpan(enc []byte, dest *basictracer.RawSpan) error {
	d := gob.NewDecoder(bytes.NewBuffer(enc))
	var err error
	decode := func(arg interface{}) {
		if err == nil {
			err = d.Decode(arg)
		}
	}

	decode(&dest.Context)
	decode(&dest.ParentSpanID)
	decode(&dest.Operation)
	decode(&dest.Start)
	decode(&dest.Duration)
	decode(&dest.Tags)

	var numLogs int32
	decode(&numLogs)
	dest.Logs = make([]opentracing.LogRecord, numLogs)
	for i := range dest.Logs {
		lr := &dest.Logs[i]
		decode(&lr.Timestamp)

		var numFields int32
		decode(&numFields)
		lr.Fields = make([]otlog.Field, numFields)
		for j := range lr.Fields {
			var key, val string
			decode(&key)
			decode(&val)
			lr.Fields[j] = otlog.String(key, val)
		}
	}

	return err
}

// contextTracerKeyType is an empty type for the handle associated with the
// tracer value (see context.Value).
type contextTracerKeyType struct{}

// WithTracer returns a context derived from the given context, for which
// TracerFromCtx returns the given tracer.
func WithTracer(ctx context.Context, tracer opentracing.Tracer) context.Context {
	return context.WithValue(ctx, contextTracerKeyType{}, tracer)
}

// TracerFromCtx returns the tracer set on the context (or a parent context) via
// WithTracer.
func TracerFromCtx(ctx context.Context) opentracing.Tracer {
	if tracerVal := ctx.Value(contextTracerKeyType{}); tracerVal != nil {
		return tracerVal.(opentracing.Tracer)
	}
	return nil
}
<|endoftext|>


<commit_before>package utils

import (
	"fmt"
	"io"
	"os"
	"os/signal"

	"golang.org/x/crypto/ssh/terminal"
)

// AttachConsole attaches stdin and stdout to the console
// in -> stdinWriter | stdinReader -> console
// out <- stdoutReader | stdoutWriter <- console
func AttachConsole(stdinReader, stdoutReader *io.PipeReader, stdinWriter, stdoutWriter *io.PipeWriter, message string, resChan <-chan error) error {
	stopChan := make(chan struct{}, 1)
	writeStop := make(chan error)
	readStop := make(chan error)
	state, err := terminal.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		return fmt.Errorf("Make raw terminal failed: %s", err)
	}
	fmt.Fprint(os.Stderr, message)

	in := os.Stdin
	out := os.Stdout

	go func() {
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt)
		<-interrupt
		close(stopChan)
	}()

	go func() {
		_, err := io.Copy(out, stdoutReader)
		readStop <- err
	}()

	go func() {
		defer close(writeStop)
		buf := make([]byte, 1024, 1024)
		for {
			// reading from stdin
			n, err := in.Read(buf)
			if err != nil && err != io.EOF {
				writeStop <- err
				return
			}
			if n == 0 && err == io.EOF {
				return
			}

			// the escape sequence
			if buf[0] == 29 {
				return
			}
			// Writing out to the console connection
			_, err = stdinWriter.Write(buf[0:n])
			if err == io.EOF {
				return
			}
		}
	}()

	select {
	case <-stopChan:
	case err = <-readStop:
	case err = <-writeStop:
	case err = <-resChan:
	}

	terminal.Restore(int(os.Stdin.Fd()), state)
	return err
}
<commit_msg>Don't change terminal settings if not run in a terminal<commit_after>package utils

import (
	"fmt"
	"io"
	"os"
	"os/signal"

	"golang.org/x/crypto/ssh/terminal"
)

// AttachConsole attaches stdin and stdout to the console
// in -> stdinWriter | stdinReader -> console
// out <- stdoutReader | stdoutWriter <- console
func AttachConsole(stdinReader, stdoutReader *io.PipeReader, stdinWriter, stdoutWriter *io.PipeWriter, message string, resChan <-chan error) (err error) {
	stopChan := make(chan struct{}, 1)
	writeStop := make(chan error)
	readStop := make(chan error)
	if terminal.IsTerminal(int(os.Stdin.Fd())) {
		state, err := terminal.MakeRaw(int(os.Stdin.Fd()))
		if err != nil {
			return fmt.Errorf("Make raw terminal failed: %s", err)
		}
		defer terminal.Restore(int(os.Stdin.Fd()), state)
	}
	fmt.Fprint(os.Stderr, message)

	in := os.Stdin
	out := os.Stdout

	go func() {
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt)
		<-interrupt
		close(stopChan)
	}()

	go func() {
		_, err := io.Copy(out, stdoutReader)
		readStop <- err
	}()

	go func() {
		defer close(writeStop)
		buf := make([]byte, 1024, 1024)
		for {
			// reading from stdin
			n, err := in.Read(buf)
			if err != nil && err != io.EOF {
				writeStop <- err
				return
			}
			if n == 0 && err == io.EOF {
				return
			}

			// the escape sequence
			if buf[0] == 29 {
				return
			}
			// Writing out to the console connection
			_, err = stdinWriter.Write(buf[0:n])
			if err == io.EOF {
				return
			}
		}
	}()

	select {
	case <-stopChan:
	case err = <-readStop:
	case err = <-writeStop:
	case err = <-resChan:
	}

	return err
}
<|endoftext|>


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

// Package wellknown contains common names for filters, listeners, etc.
package wellknown

// HTTP filter names
const (
	// Buffer HTTP filter
	Buffer = "envoy.buffer"
	// CORS HTTP filter
	CORS = "envoy.filters.http.cors"
	// Dynamo HTTP filter
	Dynamo = "envoy.http_dynamo_filter"
	// Fault HTTP filter
	Fault = "envoy.fault"
	// GRPCHTTP1Bridge HTTP filter
	GRPCHTTP1Bridge = "envoy.grpc_http1_bridge"
	// GRPCJSONTranscoder HTTP filter
	GRPCJSONTranscoder = "envoy.grpc_json_transcoder"
	// GRPCWeb HTTP filter
	GRPCWeb = "envoy.grpc_web"
	// Gzip HTTP filter
	Gzip = "envoy.gzip"
	// IPTagging HTTP filter
	IPTagging = "envoy.ip_tagging"
	// HTTPRateLimit filter
	HTTPRateLimit = "envoy.rate_limit"
	// Router HTTP filter
	Router = "envoy.router"
	// Health checking HTTP filter
	HealthCheck = "envoy.health_check"
	// Lua HTTP filter
	Lua = "envoy.lua"
	// Squash HTTP filter
	Squash = "envoy.squash"
	// HTTPExternalAuthorization HTTP filter
	HTTPExternalAuthorization = "envoy.filters.http.ext_authz"
	// HTTPRoleBasedAccessControl HTTP filter
	HTTPRoleBasedAccessControl = "envoy.filters.http.rbac"
	// HTTPGRPCStats HTTP filter
	HTTPGRPCStats = "envoy.filters.http.grpc_stats"
)

// Network filter names
const (
	// ClientSSLAuth network filter
	ClientSSLAuth = "envoy.client_ssl_auth"
	// Echo network filter
	Echo = "envoy.echo"
	// HTTPConnectionManager network filter
	HTTPConnectionManager = "envoy.http_connection_manager"
	// TCPProxy network filter
	TCPProxy = "envoy.tcp_proxy"
	// RateLimit network filter
	RateLimit = "envoy.ratelimit"
	// MongoProxy network filter
	MongoProxy = "envoy.mongo_proxy"
	// ThriftProxy network filter
	ThriftProxy = "envoy.filters.network.thrift_proxy"
	// RedisProxy network filter
	RedisProxy = "envoy.redis_proxy"
	// MySQLProxy network filter
	MySQLProxy = "envoy.filters.network.mysql_proxy"
	// ExternalAuthorization network filter
	ExternalAuthorization = "envoy.filters.network.ext_authz"
	// RoleBasedAccessControl network filter
	RoleBasedAccessControl = "envoy.filters.network.rbac"
)

// Listener filter names
const (
	// OriginalDestination listener filter
	OriginalDestination = "envoy.listener.original_dst"
	// ProxyProtocol listener filter
	ProxyProtocol = "envoy.listener.proxy_protocol"
	// TlsInspector listener filter
	TlsInspector = "envoy.listener.tls_inspector"
	// HttpInspector listener filter
	HttpInspector = "envoy.listener.http_inspector"
)

// Tracing provider names
const (
	// Lightstep tracer name
	Lightstep = "envoy.lightstep"
	// Zipkin tracer name
	Zipkin = "envoy.zipkin"
	// DynamicOT tracer name
	DynamicOT = "envoy.dynamic.ot"
)

// Stats sink names
const (
	// Statsd sink
	Statsd = "envoy.statsd"
	// DogStatsD compatible stastsd sink
	DogStatsd = "envoy.dog_statsd"
	// MetricsService sink
	MetricsService = "envoy.metrics_service"
)

// Access log sink names
const (
	// FileAccessLog sink name
	FileAccessLog = "envoy.file_access_log"
	// HTTPGRPCAccessLog sink for the HTTP gRPC access log service
	HTTPGRPCAccessLog = "envoy.http_grpc_access_log"
)

// Transport socket names
const (
	// TransportSocket Alts
	TransportSocketAlts = "envoy.transport_sockets.alts"
	// TransportSocket Tap
	TransportSocketTap = "envoy.transport_sockets.tap"
	// TransportSocket RawBuffer
	TransportSocketRawBuffer = "envoy.transport_sockets.raw_buffer"
	// TransportSocket Tls
	TransportSocketTls = "envoy.transport_sockets.tls"
	// TransportSocket Quic
	TransportSocketQuic = "envoy.transport_sockets.quic"
)
<commit_msg>pkg/wellknown: update wellknown const types (#294)<commit_after>// Copyright 2018 Envoyproxy Authors
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

// Package wellknown contains common names for filters, listeners, etc.
package wellknown

// HTTP filter names
const (
	// Buffer HTTP filter
	Buffer = "envoy.filters.http.buffer"
	// CORS HTTP filter
	CORS = "envoy.filters.http.cors"
	// Dynamo HTTP filter
	Dynamo = "envoy.filters.http.dynamo"
	// Fault HTTP filter
	Fault = "envoy.filters.http.fault"
	// GRPCHTTP1Bridge HTTP filter
	GRPCHTTP1Bridge = "envoy.filters.http.grpc_http1_bridge"
	// GRPCJSONTranscoder HTTP filter
	GRPCJSONTranscoder = "envoy.filters.http.grpc_json_transcoder"
	// GRPCWeb HTTP filter
	GRPCWeb = "envoy.filters.http.grpc_web"
	// Gzip HTTP filter
	Gzip = "envoy.filters.http.gzip"
	// IPTagging HTTP filter
	IPTagging = "envoy.filters.http.ip_tagging"
	// HTTPRateLimit filter
	HTTPRateLimit = "envoy.filters.http.ratelimit"
	// Router HTTP filter
	Router = "envoy.filters.http.router"
	// Health checking HTTP filter
	HealthCheck = "envoy.filters.http.health_check"
	// Lua HTTP filter
	Lua = "envoy.filters.http.lua"
	// Squash HTTP filter
	Squash = "envoy.filters.http.squash"
	// HTTPExternalAuthorization HTTP filter
	HTTPExternalAuthorization = "envoy.filters.http.ext_authz"
	// HTTPRoleBasedAccessControl HTTP filter
	HTTPRoleBasedAccessControl = "envoy.filters.http.rbac"
	// HTTPGRPCStats HTTP filter
	HTTPGRPCStats = "envoy.filters.http.grpc_stats"
)

// Network filter names
const (
	// ClientSSLAuth network filter
	ClientSSLAuth = "envoy.filters.network.client_ssl_auth"
	// Echo network filter
	Echo = "envoy.filters.network.echo"
	// HTTPConnectionManager network filter
	HTTPConnectionManager = "envoy.filters.network.http_connection_manager"
	// TCPProxy network filter
	TCPProxy = "envoy.filters.network.tcp_proxy"
	// RateLimit network filter
	RateLimit = "envoy.filters.network.ratelimit"
	// MongoProxy network filter
	MongoProxy = "envoy.filters.network.mongo_proxy"
	// ThriftProxy network filter
	ThriftProxy = "envoy.filters.network.thrift_proxy"
	// RedisProxy network filter
	RedisProxy = "envoy.filters.network.redis_proxy"
	// MySQLProxy network filter
	MySQLProxy = "envoy.filters.network.mysql_proxy"
	// ExternalAuthorization network filter
	ExternalAuthorization = "envoy.filters.network.ext_authz"
	// RoleBasedAccessControl network filter
	RoleBasedAccessControl = "envoy.filters.network.rbac"
)

// Listener filter names
const (
	// OriginalDestination listener filter
	OriginalDestination = "envoy.filters.listener.original_dst"
	// ProxyProtocol listener filter
	ProxyProtocol = "envoy.filters.listener.proxy_protocol"
	// TlsInspector listener filter
	TlsInspector = "envoy.filters.listener.tls_inspector"
	// HttpInspector listener filter
	HttpInspector = "envoy.filters.listener.http_inspector"
)

// Tracing provider names
const (
	// Lightstep tracer name
	Lightstep = "envoy.tracers.lightstep"
	// Zipkin tracer name
	Zipkin = "envoy.tracers.zipkin"
	// DynamicOT tracer name
	DynamicOT = "envoy.tracers.dynamic_ot"
)

// Stats sink names
const (
	// Statsd sink
	Statsd = "envoy.stat_sinks.statsd"
	// DogStatsD compatible stastsd sink
	DogStatsd = "envoy.stat_sinks.dog_statsd"
	// MetricsService sink
	MetricsService = "envoy.stat_sinks.metrics_service"
)

// Access log sink names
const (
	// FileAccessLog sink name
	FileAccessLog = "envoy.access_loggers.file"
	// HTTPGRPCAccessLog sink for the HTTP gRPC access log service
	HTTPGRPCAccessLog = "envoy.access_loggers.http_grpc"
)

// Transport socket names
const (
	// TransportSocket Alts
	TransportSocketAlts = "envoy.transport_sockets.alts"
	// TransportSocket Tap
	TransportSocketTap = "envoy.transport_sockets.tap"
	// TransportSocket RawBuffer
	TransportSocketRawBuffer = "envoy.transport_sockets.raw_buffer"
	// TransportSocket Tls
	TransportSocketTls = "envoy.transport_sockets.tls"
	// TransportSocket Quic
	TransportSocketQuic = "envoy.transport_sockets.quic"
)
<|endoftext|>


<commit_before>package scraper

import (
	"fmt"
	"log"
	"math"
	"sync"
	"sync/atomic"

	"github.com/influxdata/influxdb/models"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/storage"
)

var (
	_ Registry               = &Service{}
	_ storage.SampleAppender = &Service{}
)

// Service represents the scraper manager
type Service struct {
	PointsWriter interface {
		WritePoints(database, retentionPolicy string, consistencyLevel models.ConsistencyLevel, points []models.Point) error
	}
	mu sync.Mutex
	wg sync.WaitGroup

	open     bool
	running  bool
	closing  chan struct{}
	updating chan *config.Config

	configs atomic.Value // []Config

	logger *log.Logger

	discoverers []Discoverer

	// TargetManager represents a scraping/discovery manager
	mgr interface {
		ApplyConfig(cfg *config.Config) error
		Stop()
		Start()
		Wait()
	}
}

// NewService creates a new scraper service
func NewService(c []Config, l *log.Logger) *Service {
	s := &Service{
		logger: l,
	}
	s.storeConfigs(c)
	s.mgr = retrieval.NewTargetManager(s, NewLogger(l))
	return s
}

// Open starts the scraper service
func (s *Service) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.open {
		return nil
	}

	s.open = true
	s.updating = make(chan *config.Config)
	s.closing = make(chan struct{})

	go s.scrape()

	s.open = true
	return nil
}

// Close stops the scraper service
func (s *Service) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.open {
		return nil
	}

	s.open = false
	close(s.closing)

	s.wg.Wait()

	return nil
}

func (s *Service) loadConfigs() []Config {
	return s.configs.Load().([]Config)
}

func (s *Service) storeConfigs(c []Config) {
	s.configs.Store(c)
}

func (s *Service) scrape() {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		s.mu.Lock()
		pairs := s.pairs()
		conf := s.prom(pairs)
		s.mu.Unlock()

		s.mgr.ApplyConfig(conf)

		s.mu.Lock()
		// Need to check open if service was stopped just before this lock
		if s.open {
			s.mgr.Start()
			s.running = true
		}
		s.mu.Unlock()
		// Wait will block until mgr.Stop has been called.
		s.mgr.Wait()
	}()

	for {
		select {
		case <-s.closing:
			s.mu.Lock()
			// Need to check if the targetmanager is even running before Stopping
			if s.running {
				s.mgr.Stop()
				s.running = false
			}
			s.mu.Unlock()
			return
		case conf := <-s.updating:
			s.mgr.ApplyConfig(conf)
		}
	}
}

// Append tranforms prometheus samples and inserts data into the tasks pipeline
func (s *Service) Append(sample *model.Sample) error {
	value := float64(sample.Value)
	// Remove all NaN values
	// TODO: Add counter stat for this variable
	if math.IsNaN(value) {
		return nil
	}

	var err error
	db := ""
	rp := ""
	job := ""

	tags := map[string]string{}
	for name, value := range sample.Metric {
		if name == "job" {
			db, rp, job, err = decodeJobName(string(value))
			continue
		}
		tags[string(name)] = string(value)
	}

	// If instance is blacklisted then do not send to PointsWriter
	if instance, ok := tags["instance"]; ok {
		for _, c := range s.loadConfigs() {
			if c.Name == job {
				for _, listed := range c.Blacklist {
					if instance == listed {
						return nil
					}
				}
			}
		}
	}

	fields := models.Fields{
		"value": value,
	}

	pt, err := models.NewPoint(job, models.NewTags(tags), fields, sample.Timestamp.Time())
	if err != nil {
		return err
	}

	return s.PointsWriter.WritePoints(db, rp, models.ConsistencyLevelAny, []models.Point{pt})
}

// NeedsThrottling conforms to SampleAppender and never returns true currently.
func (s *Service) NeedsThrottling() bool {
	return false
}

type testOptions struct {
	Name string `json:"name"`
}

// TestOptions returns options that are allowed for the Test
func (s *Service) TestOptions() interface{} {
	return &testOptions{}
}

// Test tests the options for the scrapers
func (s *Service) Test(options interface{}) error {
	return nil
}

// Update will replace all scraper configurations and apply the configuration
// to the target manager
func (s *Service) Update(newConfigs []interface{}) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	configs := make([]Config, len(newConfigs))
	for i, c := range newConfigs {
		if config, ok := c.(Config); ok {
			configs[i] = config
		} else {
			return fmt.Errorf("unexpected config object type, got %T exp %T", c, config)
		}
	}

	s.storeConfigs(configs)
	if s.open {
		pairs := s.pairs()
		conf := s.prom(pairs)
		select {
		case <-s.closing:
			return fmt.Errorf("error writing configuration to closed scraper")
		case s.updating <- conf:
		}
	}
	return nil
}

// prom assumes service is locked
func (s *Service) prom(pairs []Pair) *config.Config {
	conf := &config.Config{
		ScrapeConfigs: make([]*config.ScrapeConfig, len(pairs)),
	}

	for i, pair := range pairs {
		sc := pair.Scraper.Prom()
		pair.Discoverer.Prom(sc)
		conf.ScrapeConfigs[i] = sc
	}
	return conf
}

// Commit applies the configuration to the scraper
func (s *Service) Commit() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return nil
	}

	pairs := s.pairs()
	conf := s.prom(pairs)
	select {
	case <-s.closing:
		return fmt.Errorf("error writing configuration to closed registry")
	case s.updating <- conf:
	}

	return nil
}

// AddDiscoverer adds discoverer to the registry
func (s *Service) AddDiscoverer(discoverer Discoverer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	s.discoverers = append(s.discoverers, discoverer)
}

// RemoveDiscoverer removes discoverer from the registry
func (s *Service) RemoveDiscoverer(rm Discoverer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	filtered := s.discoverers[0:0]
	for _, d := range s.discoverers {
		if d.ServiceID() != rm.ServiceID() || d.Service() != rm.Service() {
			filtered = append(filtered, d)
		}
	}
	s.discoverers = filtered
}

// AddScrapers adds scrapers to the registry
func (s *Service) AddScrapers(scrapers []Config) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	configs := s.loadConfigs()
	configs = append(configs, scrapers...)
	s.storeConfigs(configs)
}

// RemoveScrapers removes scrapers from the registry
func (s *Service) RemoveScrapers(scrapers []Config) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	configs := s.loadConfigs()
	for _, rm := range scrapers {
		for i, c := range configs {
			if c.Name == rm.Name {
				configs = append(configs, configs[i+1:]...)
			}
		}
	}
	s.storeConfigs(configs)
}

// pairs returns all named pairs of scrapers and discoverers from registry must be locked
func (s *Service) pairs() []Pair {
	pairs := []Pair{}
	for _, scr := range s.loadConfigs() {
		if !scr.Enabled {
			continue
		}
		for _, d := range s.discoverers {
			if scr.DiscoverID == d.ServiceID() && scr.DiscoverService == d.Service() {
				pairs = append(pairs, Pair{
					Scraper:    scr,
					Discoverer: d,
				})
			}
		}
	}
	return pairs
}

// Pairs returns all named pairs of scrapers and discoverers from registry must be locked
func (s *Service) Pairs() []Pair {
	s.mu.Lock()
	defer s.mu.Unlock()

	return s.pairs()
}
<commit_msg>Fix removal of scraper config<commit_after>package scraper

import (
	"fmt"
	"log"
	"math"
	"sync"
	"sync/atomic"

	"github.com/influxdata/influxdb/models"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/storage"
)

var (
	_ Registry               = &Service{}
	_ storage.SampleAppender = &Service{}
)

// Service represents the scraper manager
type Service struct {
	PointsWriter interface {
		WritePoints(database, retentionPolicy string, consistencyLevel models.ConsistencyLevel, points []models.Point) error
	}
	mu sync.Mutex
	wg sync.WaitGroup

	open     bool
	running  bool
	closing  chan struct{}
	updating chan *config.Config

	configs atomic.Value // []Config

	logger *log.Logger

	discoverers []Discoverer

	// TargetManager represents a scraping/discovery manager
	mgr interface {
		ApplyConfig(cfg *config.Config) error
		Stop()
		Start()
		Wait()
	}
}

// NewService creates a new scraper service
func NewService(c []Config, l *log.Logger) *Service {
	s := &Service{
		logger: l,
	}
	s.storeConfigs(c)
	s.mgr = retrieval.NewTargetManager(s, NewLogger(l))
	return s
}

// Open starts the scraper service
func (s *Service) Open() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.open {
		return nil
	}

	s.open = true
	s.updating = make(chan *config.Config)
	s.closing = make(chan struct{})

	go s.scrape()

	s.open = true
	return nil
}

// Close stops the scraper service
func (s *Service) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.open {
		return nil
	}

	s.open = false
	close(s.closing)

	s.wg.Wait()

	return nil
}

func (s *Service) loadConfigs() []Config {
	return s.configs.Load().([]Config)
}

func (s *Service) storeConfigs(c []Config) {
	s.configs.Store(c)
}

func (s *Service) scrape() {
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		s.mu.Lock()
		pairs := s.pairs()
		conf := s.prom(pairs)
		s.mu.Unlock()

		s.mgr.ApplyConfig(conf)

		s.mu.Lock()
		// Need to check open if service was stopped just before this lock
		if s.open {
			s.mgr.Start()
			s.running = true
		}
		s.mu.Unlock()
		// Wait will block until mgr.Stop has been called.
		s.mgr.Wait()
	}()

	for {
		select {
		case <-s.closing:
			s.mu.Lock()
			// Need to check if the targetmanager is even running before Stopping
			if s.running {
				s.mgr.Stop()
				s.running = false
			}
			s.mu.Unlock()
			return
		case conf := <-s.updating:
			s.mgr.ApplyConfig(conf)
		}
	}
}

// Append tranforms prometheus samples and inserts data into the tasks pipeline
func (s *Service) Append(sample *model.Sample) error {
	value := float64(sample.Value)
	// Remove all NaN values
	// TODO: Add counter stat for this variable
	if math.IsNaN(value) {
		return nil
	}

	var err error
	db := ""
	rp := ""
	job := ""

	tags := map[string]string{}
	for name, value := range sample.Metric {
		if name == "job" {
			db, rp, job, err = decodeJobName(string(value))
			continue
		}
		tags[string(name)] = string(value)
	}

	// If instance is blacklisted then do not send to PointsWriter
	if instance, ok := tags["instance"]; ok {
		for _, c := range s.loadConfigs() {
			if c.Name == job {
				for _, listed := range c.Blacklist {
					if instance == listed {
						return nil
					}
				}
			}
		}
	}

	fields := models.Fields{
		"value": value,
	}

	pt, err := models.NewPoint(job, models.NewTags(tags), fields, sample.Timestamp.Time())
	if err != nil {
		return err
	}

	return s.PointsWriter.WritePoints(db, rp, models.ConsistencyLevelAny, []models.Point{pt})
}

// NeedsThrottling conforms to SampleAppender and never returns true currently.
func (s *Service) NeedsThrottling() bool {
	return false
}

type testOptions struct {
	Name string `json:"name"`
}

// TestOptions returns options that are allowed for the Test
func (s *Service) TestOptions() interface{} {
	return &testOptions{}
}

// Test tests the options for the scrapers
func (s *Service) Test(options interface{}) error {
	return nil
}

// Update will replace all scraper configurations and apply the configuration
// to the target manager
func (s *Service) Update(newConfigs []interface{}) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	configs := make([]Config, len(newConfigs))
	for i, c := range newConfigs {
		if config, ok := c.(Config); ok {
			configs[i] = config
		} else {
			return fmt.Errorf("unexpected config object type, got %T exp %T", c, config)
		}
	}

	s.storeConfigs(configs)
	if s.open {
		pairs := s.pairs()
		conf := s.prom(pairs)
		select {
		case <-s.closing:
			return fmt.Errorf("error writing configuration to closed scraper")
		case s.updating <- conf:
		}
	}
	return nil
}

// prom assumes service is locked
func (s *Service) prom(pairs []Pair) *config.Config {
	conf := &config.Config{
		ScrapeConfigs: make([]*config.ScrapeConfig, len(pairs)),
	}

	for i, pair := range pairs {
		sc := pair.Scraper.Prom()
		pair.Discoverer.Prom(sc)
		conf.ScrapeConfigs[i] = sc
	}
	return conf
}

// Commit applies the configuration to the scraper
func (s *Service) Commit() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return nil
	}

	pairs := s.pairs()
	conf := s.prom(pairs)
	select {
	case <-s.closing:
		return fmt.Errorf("error writing configuration to closed registry")
	case s.updating <- conf:
	}

	return nil
}

// AddDiscoverer adds discoverer to the registry
func (s *Service) AddDiscoverer(discoverer Discoverer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	s.discoverers = append(s.discoverers, discoverer)
}

// RemoveDiscoverer removes discoverer from the registry
func (s *Service) RemoveDiscoverer(rm Discoverer) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	filtered := s.discoverers[0:0]
	for _, d := range s.discoverers {
		if d.ServiceID() != rm.ServiceID() || d.Service() != rm.Service() {
			filtered = append(filtered, d)
		}
	}
	s.discoverers = filtered
}

// AddScrapers adds scrapers to the registry
func (s *Service) AddScrapers(scrapers []Config) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	configs := s.loadConfigs()
	configs = append(configs, scrapers...)
	s.storeConfigs(configs)
}

// RemoveScrapers removes scrapers from the registry
func (s *Service) RemoveScrapers(scrapers []Config) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.open {
		return
	}

	configs := s.loadConfigs()[0:0]
	for _, rm := range scrapers {
		for _, c := range configs {
			if c.Name != rm.Name {
				configs = append(configs, c)
			}
		}
	}
	s.storeConfigs(configs)
}

// pairs returns all named pairs of scrapers and discoverers from registry must be locked
func (s *Service) pairs() []Pair {
	pairs := []Pair{}
	for _, scr := range s.loadConfigs() {
		if !scr.Enabled {
			continue
		}
		for _, d := range s.discoverers {
			if scr.DiscoverID == d.ServiceID() && scr.DiscoverService == d.Service() {
				pairs = append(pairs, Pair{
					Scraper:    scr,
					Discoverer: d,
				})
			}
		}
	}
	return pairs
}

// Pairs returns all named pairs of scrapers and discoverers from registry must be locked
func (s *Service) Pairs() []Pair {
	s.mu.Lock()
	defer s.mu.Unlock()

	return s.pairs()
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

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

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

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

package event

var OnStartList []func()

func OnStart(index int, fn ...func()) {
	if len(fn) == 0 {
		return
	}
	if index < 0 {
		OnStartList = append(OnStartList, fn...)
		return
	}
	size := len(OnStartList)
	if size > index {
		OnStartList[index] = fn[0]
		if len(fn) > 1 {
			OnStart(index+1, fn[1:]...)
		}
		return
	}
	for start, end := size, index-1; start < end; start++ {
		OnStartList = append(OnStartList, nil)
	}
	OnStartList = append(OnStartList, fn...)
}

func Start() {
	for _, fn := range OnStartList {
		if fn == nil {
			continue
		}
		fn()
	}
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

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

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

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

package event

var OnStartList []func()

func OnStart(index int, fn ...func()) {
	if len(fn) == 0 {
		return
	}
	if index < 0 {
		OnStartList = append(OnStartList, fn...)
		return
	}
	size := len(OnStartList)
	if size > index {
		fn = append(fn, OnStartList[index])
		OnStartList[index] = fn[0]
		if len(fn) > 1 {
			OnStart(index+1, fn[1:]...)
		}
		return
	}
	for start, end := size, index-1; start < end; start++ {
		OnStartList = append(OnStartList, nil)
	}
	OnStartList = append(OnStartList, fn...)
}

func Start() {
	for _, fn := range OnStartList {
		if fn == nil {
			continue
		}
		fn()
	}
}
<|endoftext|>


<commit_before>package discovery

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"time"

	"github.com/flynn/flynn/Godeps/_workspace/src/golang.org/x/crypto/ssh"
	"github.com/flynn/flynn/pkg/version"
)

type Info struct {
	ClusterURL  string
	InstanceURL string
	Name        string
}

type Instance struct {
	ID            string         `json:"id,omitempty"`
	ClusterID     string         `json:"cluster_id,omitempty"`
	FlynnVersion  string         `json:"flynn_version,omitempty"`
	SSHPublicKeys []SSHPublicKey `json:"ssh_public_keys,omitempty"`
	URL           string         `json:"url,omitempty"`
	Name          string         `json:"name,omitempty"`
	CreatedAt     *time.Time     `json:"created_at,omitempty"`
}

type SSHPublicKey struct {
	Type string `json:"type"`
	Data []byte `json:"data"`
}

func RegisterInstance(info Info) (string, error) {
	data := struct {
		Data Instance `json:"data"`
	}{Instance{
		Name:          info.Name,
		URL:           info.InstanceURL,
		SSHPublicKeys: make([]SSHPublicKey, 0, 4),
		FlynnVersion:  version.String(),
	}}

	for _, t := range []string{"dsa", "rsa", "ecdsa", "ed25519"} {
		keyData, err := ioutil.ReadFile(fmt.Sprintf("/etc/ssh/ssh_host_%s_key.pub", t))
		if err != nil {
			// TODO(titanous): log this?
			continue
		}
		k, _, _, _, err := ssh.ParseAuthorizedKey(keyData)
		if err != nil {
			// TODO(titanous): log this?
			continue
		}
		data.Data.SSHPublicKeys = append(data.Data.SSHPublicKeys, SSHPublicKey{Type: t, Data: k.Marshal()})
	}

	jsonData, err := json.Marshal(&data)
	if err != nil {
		return "", err
	}
	// TODO(titanous): retry
	uri := info.ClusterURL + "/instances"
	res, err := http.Post(uri, "application/json", bytes.NewReader(jsonData))
	if err != nil {
		return "", err
	}
	if res.StatusCode != http.StatusCreated && res.StatusCode != http.StatusConflict {
		return "", &url.Error{
			Op:  "POST",
			URL: uri,
			Err: fmt.Errorf("unexpected status %d", res.StatusCode),
		}
	}
	if err := json.NewDecoder(res.Body).Decode(&data); err != nil {
		return "", err
	}
	return data.Data.ID, nil
}

func GetCluster(uri string) ([]*Instance, error) {
	uri += "/instances"
	res, err := http.Get(uri)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != 200 {
		return nil, urlError(uri, res.StatusCode)
	}
	defer res.Body.Close()

	var data struct {
		Data []*Instance `json:"data"`
	}
	err = json.NewDecoder(res.Body).Decode(&data)
	return data.Data, err
}

func NewToken() (string, error) {
	uri := "https://discovery.flynn.io/clusters"
	if base := os.Getenv("DISCOVERY_SERVER"); base != "" {
		uri = base + "/clusters"
	}

	req, err := http.NewRequest("POST", uri, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("User-Agent", fmt.Sprintf("flynn-host/%s %s-%s", version.String(), runtime.GOOS, runtime.GOARCH))
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	if res.StatusCode != http.StatusCreated {
		return "", urlError(uri, res.StatusCode)
	}

	base, err := url.Parse(uri)
	if err != nil {
		return "", err
	}
	cluster, err := url.Parse(res.Header.Get("Location"))
	if err != nil {
		return "", err
	}

	return base.ResolveReference(cluster).String(), nil
}

func urlError(uri string, status int) error {
	return &url.Error{
		Op:  "GET",
		URL: uri,
		Err: fmt.Errorf("unexpected status %d", status),
	}

}
<commit_msg>bootstrap/discovery: Refactor URL errors<commit_after>package discovery

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"time"

	"github.com/flynn/flynn/Godeps/_workspace/src/golang.org/x/crypto/ssh"
	"github.com/flynn/flynn/pkg/version"
)

type Info struct {
	ClusterURL  string
	InstanceURL string
	Name        string
}

type Instance struct {
	ID            string         `json:"id,omitempty"`
	ClusterID     string         `json:"cluster_id,omitempty"`
	FlynnVersion  string         `json:"flynn_version,omitempty"`
	SSHPublicKeys []SSHPublicKey `json:"ssh_public_keys,omitempty"`
	URL           string         `json:"url,omitempty"`
	Name          string         `json:"name,omitempty"`
	CreatedAt     *time.Time     `json:"created_at,omitempty"`
}

type SSHPublicKey struct {
	Type string `json:"type"`
	Data []byte `json:"data"`
}

func RegisterInstance(info Info) (string, error) {
	data := struct {
		Data Instance `json:"data"`
	}{Instance{
		Name:          info.Name,
		URL:           info.InstanceURL,
		SSHPublicKeys: make([]SSHPublicKey, 0, 4),
		FlynnVersion:  version.String(),
	}}

	for _, t := range []string{"dsa", "rsa", "ecdsa", "ed25519"} {
		keyData, err := ioutil.ReadFile(fmt.Sprintf("/etc/ssh/ssh_host_%s_key.pub", t))
		if err != nil {
			// TODO(titanous): log this?
			continue
		}
		k, _, _, _, err := ssh.ParseAuthorizedKey(keyData)
		if err != nil {
			// TODO(titanous): log this?
			continue
		}
		data.Data.SSHPublicKeys = append(data.Data.SSHPublicKeys, SSHPublicKey{Type: t, Data: k.Marshal()})
	}

	jsonData, err := json.Marshal(&data)
	if err != nil {
		return "", err
	}
	// TODO(titanous): retry
	uri := info.ClusterURL + "/instances"
	res, err := http.Post(uri, "application/json", bytes.NewReader(jsonData))
	if err != nil {
		return "", err
	}
	if res.StatusCode != http.StatusCreated && res.StatusCode != http.StatusConflict {
		return "", urlError("POST", uri, res.StatusCode)
	}
	if err := json.NewDecoder(res.Body).Decode(&data); err != nil {
		return "", err
	}
	return data.Data.ID, nil
}

func GetCluster(uri string) ([]*Instance, error) {
	uri += "/instances"
	res, err := http.Get(uri)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != 200 {
		return nil, urlError("GET", uri, res.StatusCode)
	}
	defer res.Body.Close()

	var data struct {
		Data []*Instance `json:"data"`
	}
	err = json.NewDecoder(res.Body).Decode(&data)
	return data.Data, err
}

func NewToken() (string, error) {
	uri := "https://discovery.flynn.io/clusters"
	if base := os.Getenv("DISCOVERY_SERVER"); base != "" {
		uri = base + "/clusters"
	}

	req, err := http.NewRequest("POST", uri, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("User-Agent", fmt.Sprintf("flynn-host/%s %s-%s", version.String(), runtime.GOOS, runtime.GOARCH))
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	if res.StatusCode != http.StatusCreated {
		return "", urlError("POST", uri, res.StatusCode)
	}

	base, err := url.Parse(uri)
	if err != nil {
		return "", err
	}
	cluster, err := url.Parse(res.Header.Get("Location"))
	if err != nil {
		return "", err
	}

	return base.ResolveReference(cluster).String(), nil
}

func urlError(method, uri string, status int) error {
	return &url.Error{
		Op:  method,
		URL: uri,
		Err: fmt.Errorf("unexpected status %d", status),
	}

}
<|endoftext|>