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

package mongo

import (
	"reflect"
	"time"

	"github.com/juju/errors"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"launchpad.net/tomb"
)

// OplogDoc represents a document in the oplog.rs collection.
// See: http://www.kchodorow.com/blog/2010/10/12/replication-internals/
//
// The Object and UpdateObject fields are returned raw to allow
// unmarshalling into arbitrary types. Use the UnmarshalObject and
// UnmarshalUpdate methods to unmarshall these fields.
type OplogDoc struct {
	Timestamp    bson.MongoTimestamp `bson:"ts"`
	OperationId  int64               `bson:"h"`
	MongoVersion int                 `bson:"v"`
	Operation    string              `bson:"op"` // "i" - insert, "u" - update, "d" - delete
	Namespace    string              `bson:"ns"`
	Object       *bson.Raw           `bson:"o"`
	UpdateObject *bson.Raw           `bson:"o2"`
}

// UnmarshalObject unmarshals the Object field into out. The out
// argument should be a pointer or a suitable map.
func (d *OplogDoc) UnmarshalObject(out interface{}) error {
	return d.unmarshal(d.Object, out)
}

// UnmarshalUpdate unmarshals the UpdateObject field into out. The out
// argument should be a pointer or a suitable map.
func (d *OplogDoc) UnmarshalUpdate(out interface{}) error {
	return d.unmarshal(d.UpdateObject, out)
}

func (d *OplogDoc) unmarshal(raw *bson.Raw, out interface{}) error {
	if raw == nil {
		// If the field is not set, set out to the zero value for its type.
		v := reflect.ValueOf(out)
		switch v.Kind() {
		case reflect.Ptr:
			v = v.Elem()
			v.Set(reflect.Zero(v.Type()))
		case reflect.Map:
			// Empty the map.
			for _, k := range v.MapKeys() {
				v.SetMapIndex(k, reflect.Value{})
			}
		default:
			return errors.New("output must be a pointer or map")
		}
		return nil
	}
	return raw.Unmarshal(out)
}

// NewMongoTimestamp returns a bson.MongoTimestamp repesentation for
// the time.Time given. Note that these timestamps are not the same
// the usual MongoDB time fields. These are an internal format used
// only in a few places such as the replication oplog.
//
// See: http://docs.mongodb.org/manual/reference/bson-types/#timestamps
func NewMongoTimestamp(t time.Time) bson.MongoTimestamp {
	unixTime := t.Unix()
	if unixTime < 0 {
		unixTime = 0
	}
	return bson.MongoTimestamp(unixTime << 32)
}

// GetOplog returns the the oplog collection in the local database.
func GetOplog(session *mgo.Session) *mgo.Collection {
	return session.DB("local").C("oplog.rs")
}

func isRealOplog(c *mgo.Collection) bool {
	return c.Database.Name == "local" && c.Name == "oplog.rs"
}

// The parts of the mgo.Iter that we use - this interface will allow
// us to switch out the querying for testing.
type OplogIterator interface {
	Next(interface{}) bool
	Err() error
	Timeout() bool
}

type OplogSession interface {
	NewIter(bson.MongoTimestamp, []int64) OplogIterator
	Close()
}

type oplogSession struct {
	session    *mgo.Session
	collection *mgo.Collection
	query      bson.D
}

func newOplogSession(collection *mgo.Collection, query bson.D) *oplogSession {
	// Use a fresh session for the tailer.
	session := collection.Database.Session.Copy()
	return &oplogSession{
		session:    session,
		collection: collection.With(session),
		query:      query,
	}
}

const oplogTailTimeout = time.Second

func (s *oplogSession) NewIter(fromTimestamp bson.MongoTimestamp, excludeIds []int64) OplogIterator {
	// When recreating the iterator (required when the cursor
	// is invalidated) avoid reporting oplog entries that have
	// already been reported.
	sel := append(s.query,
		bson.DocElem{"ts", bson.D{{"$gte", fromTimestamp}}},
		bson.DocElem{"h", bson.D{{"$nin", excludeIds}}},
	)
	// Time the tail call out every second so that requests to
	// stop can be honoured.
	//
	// TODO(mjs): Ideally -1 (no timeout) could be used here,
	// with session.Close() being used to unblock Next() if
	// the tailer should stop (these semantics are hinted at
	// by the mgo docs). Unfortunately this can trigger
	// panics. See: https://github.com/go-mgo/mgo/issues/121
	query := s.collection.Find(sel)
	if isRealOplog(s.collection) {
		// Apply an optmisation that is only supported with
		// the real oplog.
		query = query.LogReplay()
	}
	return query.Tail(oplogTailTimeout)
}

func (s *oplogSession) Close() {
	s.session.Close()
}

// NewOplogTailer returns a new OplogTailer.
//
// Arguments:
// - "oplog" is the collection to use for the oplog. Typically this
//   would be the result of GetOpLog.
// - "query" can be used to limit the returned oplog entries. A
//    typical filter would limit based on ns ("<database>.<collection>")
//    and o (object).
// - "initialTs" sets the operation timestamp to start returning
//    results from. This can be used to avoid an expensive initial search
//    through the oplog when the tailer first starts.
//
// Remember to call Stop on the returned OplogTailer when it is no
// longer needed.
func NewOplogTailer(
	oplog *mgo.Collection,
	query bson.D,
	initialTs time.Time,
) *OplogTailer {
	return newOplogTailerWithSession(newOplogSession(oplog, query), initialTs)
}

func newOplogTailerWithSession(
	session OplogSession,
	initialTs time.Time,
) *OplogTailer {
	t := &OplogTailer{
		session:   session,
		initialTs: NewMongoTimestamp(initialTs),
		outCh:     make(chan *OplogDoc),
	}
	go func() {
		defer func() {
			close(t.outCh)
			t.tomb.Done()
			session.Close()
		}()
		t.tomb.Kill(t.loop())
	}()
	return t
}

// OplogTailer tails MongoDB's replication oplog.
type OplogTailer struct {
	tomb      tomb.Tomb
	session   OplogSession
	initialTs bson.MongoTimestamp
	outCh     chan *OplogDoc
}

// Out returns a channel that reports the oplog entries matching the
// query passed to NewOplogTailer as they appear.
func (t *OplogTailer) Out() <-chan *OplogDoc {
	return t.outCh
}

// Dying returns a channel that will be closed with the OplogTailer is
// shutting down.
func (t *OplogTailer) Dying() <-chan struct{} {
	return t.tomb.Dying()
}

// Stop shuts down the OplogTailer. It will block until shutdown is
// complete.
func (t *OplogTailer) Stop() error {
	t.tomb.Kill(nil)
	return t.tomb.Wait()
}

// Err returns the error that caused the OplogTailer to stop. If it
// finished normally or hasn't stopped then nil will be returned.
func (t *OplogTailer) Err() error {
	return t.tomb.Err()
}

func (t *OplogTailer) loop() error {
	var iter OplogIterator

	// lastTimestamp tracks the most recent oplog timestamp reported.
	lastTimestamp := t.initialTs

	// idsForLastTimestamp records the unique operation ids that have
	// been reported for the most recently reported oplog
	// timestamp. This is used to avoid re-reporting oplog entries
	// when the iterator is restarted. These timestamps are unique for
	// a given mongod but when there's multiple replicaset members
	// it's possible for there to be multiple oplog entries for a
	// given timestamp.
	//
	// See: http://docs.mongodb.org/v2.4/reference/bson-types/#timestamps
	var idsForLastTimestamp []int64

	for {
		if t.dying() {
			return tomb.ErrDying
		}

		if iter == nil {
			iter = t.session.NewIter(lastTimestamp, idsForLastTimestamp)
		}

		var doc OplogDoc
		if iter.Next(&doc) {
			select {
			case <-t.tomb.Dying():
				return tomb.ErrDying
			case t.outCh <- &doc:
			}

			if doc.Timestamp > lastTimestamp {
				lastTimestamp = doc.Timestamp
				idsForLastTimestamp = nil
			}
			idsForLastTimestamp = append(idsForLastTimestamp, doc.OperationId)
		} else {
			if err := iter.Err(); err != nil && err != mgo.ErrCursor {
				return err
			}
			if iter.Timeout() {
				continue
			}
			// Either there's no error or the error is an expired
			// cursor. Force recreating the iterator next loop by
			// marking it as nil.
			iter = nil
		}
	}
}

func (t *OplogTailer) dying() bool {
	select {
	case <-t.tomb.Dying():
		return true
	default:
		return false
	}
}
<commit_msg>Comment OplogSession<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package mongo

import (
	"reflect"
	"time"

	"github.com/juju/errors"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"launchpad.net/tomb"
)

// OplogDoc represents a document in the oplog.rs collection.
// See: http://www.kchodorow.com/blog/2010/10/12/replication-internals/
//
// The Object and UpdateObject fields are returned raw to allow
// unmarshalling into arbitrary types. Use the UnmarshalObject and
// UnmarshalUpdate methods to unmarshall these fields.
type OplogDoc struct {
	Timestamp    bson.MongoTimestamp `bson:"ts"`
	OperationId  int64               `bson:"h"`
	MongoVersion int                 `bson:"v"`
	Operation    string              `bson:"op"` // "i" - insert, "u" - update, "d" - delete
	Namespace    string              `bson:"ns"`
	Object       *bson.Raw           `bson:"o"`
	UpdateObject *bson.Raw           `bson:"o2"`
}

// UnmarshalObject unmarshals the Object field into out. The out
// argument should be a pointer or a suitable map.
func (d *OplogDoc) UnmarshalObject(out interface{}) error {
	return d.unmarshal(d.Object, out)
}

// UnmarshalUpdate unmarshals the UpdateObject field into out. The out
// argument should be a pointer or a suitable map.
func (d *OplogDoc) UnmarshalUpdate(out interface{}) error {
	return d.unmarshal(d.UpdateObject, out)
}

func (d *OplogDoc) unmarshal(raw *bson.Raw, out interface{}) error {
	if raw == nil {
		// If the field is not set, set out to the zero value for its type.
		v := reflect.ValueOf(out)
		switch v.Kind() {
		case reflect.Ptr:
			v = v.Elem()
			v.Set(reflect.Zero(v.Type()))
		case reflect.Map:
			// Empty the map.
			for _, k := range v.MapKeys() {
				v.SetMapIndex(k, reflect.Value{})
			}
		default:
			return errors.New("output must be a pointer or map")
		}
		return nil
	}
	return raw.Unmarshal(out)
}

// NewMongoTimestamp returns a bson.MongoTimestamp repesentation for
// the time.Time given. Note that these timestamps are not the same
// the usual MongoDB time fields. These are an internal format used
// only in a few places such as the replication oplog.
//
// See: http://docs.mongodb.org/manual/reference/bson-types/#timestamps
func NewMongoTimestamp(t time.Time) bson.MongoTimestamp {
	unixTime := t.Unix()
	if unixTime < 0 {
		unixTime = 0
	}
	return bson.MongoTimestamp(unixTime << 32)
}

// GetOplog returns the the oplog collection in the local database.
func GetOplog(session *mgo.Session) *mgo.Collection {
	return session.DB("local").C("oplog.rs")
}

func isRealOplog(c *mgo.Collection) bool {
	return c.Database.Name == "local" && c.Name == "oplog.rs"
}

// The parts of the mgo.Iter that we use - this interface will allow
// us to switch out the querying for testing.
type OplogIterator interface {
	Next(interface{}) bool
	Err() error
	Timeout() bool
}

// The connection to the oplog database, used to create an iterator to
// get oplog documents (and recreate it if it gets killed or times
// out).
type OplogSession interface {
	NewIter(bson.MongoTimestamp, []int64) OplogIterator
	Close()
}

type oplogSession struct {
	session    *mgo.Session
	collection *mgo.Collection
	query      bson.D
}

func newOplogSession(collection *mgo.Collection, query bson.D) *oplogSession {
	// Use a fresh session for the tailer.
	session := collection.Database.Session.Copy()
	return &oplogSession{
		session:    session,
		collection: collection.With(session),
		query:      query,
	}
}

const oplogTailTimeout = time.Second

func (s *oplogSession) NewIter(fromTimestamp bson.MongoTimestamp, excludeIds []int64) OplogIterator {
	// When recreating the iterator (required when the cursor
	// is invalidated) avoid reporting oplog entries that have
	// already been reported.
	sel := append(s.query,
		bson.DocElem{"ts", bson.D{{"$gte", fromTimestamp}}},
		bson.DocElem{"h", bson.D{{"$nin", excludeIds}}},
	)
	// Time the tail call out every second so that requests to
	// stop can be honoured.
	//
	// TODO(mjs): Ideally -1 (no timeout) could be used here,
	// with session.Close() being used to unblock Next() if
	// the tailer should stop (these semantics are hinted at
	// by the mgo docs). Unfortunately this can trigger
	// panics. See: https://github.com/go-mgo/mgo/issues/121
	query := s.collection.Find(sel)
	if isRealOplog(s.collection) {
		// Apply an optmisation that is only supported with
		// the real oplog.
		query = query.LogReplay()
	}
	return query.Tail(oplogTailTimeout)
}

func (s *oplogSession) Close() {
	s.session.Close()
}

// NewOplogTailer returns a new OplogTailer.
//
// Arguments:
// - "oplog" is the collection to use for the oplog. Typically this
//   would be the result of GetOpLog.
// - "query" can be used to limit the returned oplog entries. A
//    typical filter would limit based on ns ("<database>.<collection>")
//    and o (object).
// - "initialTs" sets the operation timestamp to start returning
//    results from. This can be used to avoid an expensive initial search
//    through the oplog when the tailer first starts.
//
// Remember to call Stop on the returned OplogTailer when it is no
// longer needed.
func NewOplogTailer(
	oplog *mgo.Collection,
	query bson.D,
	initialTs time.Time,
) *OplogTailer {
	return newOplogTailerWithSession(newOplogSession(oplog, query), initialTs)
}

func newOplogTailerWithSession(
	session OplogSession,
	initialTs time.Time,
) *OplogTailer {
	t := &OplogTailer{
		session:   session,
		initialTs: NewMongoTimestamp(initialTs),
		outCh:     make(chan *OplogDoc),
	}
	go func() {
		defer func() {
			close(t.outCh)
			t.tomb.Done()
			session.Close()
		}()
		t.tomb.Kill(t.loop())
	}()
	return t
}

// OplogTailer tails MongoDB's replication oplog.
type OplogTailer struct {
	tomb      tomb.Tomb
	session   OplogSession
	initialTs bson.MongoTimestamp
	outCh     chan *OplogDoc
}

// Out returns a channel that reports the oplog entries matching the
// query passed to NewOplogTailer as they appear.
func (t *OplogTailer) Out() <-chan *OplogDoc {
	return t.outCh
}

// Dying returns a channel that will be closed with the OplogTailer is
// shutting down.
func (t *OplogTailer) Dying() <-chan struct{} {
	return t.tomb.Dying()
}

// Stop shuts down the OplogTailer. It will block until shutdown is
// complete.
func (t *OplogTailer) Stop() error {
	t.tomb.Kill(nil)
	return t.tomb.Wait()
}

// Err returns the error that caused the OplogTailer to stop. If it
// finished normally or hasn't stopped then nil will be returned.
func (t *OplogTailer) Err() error {
	return t.tomb.Err()
}

func (t *OplogTailer) loop() error {
	var iter OplogIterator

	// lastTimestamp tracks the most recent oplog timestamp reported.
	lastTimestamp := t.initialTs

	// idsForLastTimestamp records the unique operation ids that have
	// been reported for the most recently reported oplog
	// timestamp. This is used to avoid re-reporting oplog entries
	// when the iterator is restarted. These timestamps are unique for
	// a given mongod but when there's multiple replicaset members
	// it's possible for there to be multiple oplog entries for a
	// given timestamp.
	//
	// See: http://docs.mongodb.org/v2.4/reference/bson-types/#timestamps
	var idsForLastTimestamp []int64

	for {
		if t.dying() {
			return tomb.ErrDying
		}

		if iter == nil {
			iter = t.session.NewIter(lastTimestamp, idsForLastTimestamp)
		}

		var doc OplogDoc
		if iter.Next(&doc) {
			select {
			case <-t.tomb.Dying():
				return tomb.ErrDying
			case t.outCh <- &doc:
			}

			if doc.Timestamp > lastTimestamp {
				lastTimestamp = doc.Timestamp
				idsForLastTimestamp = nil
			}
			idsForLastTimestamp = append(idsForLastTimestamp, doc.OperationId)
		} else {
			if err := iter.Err(); err != nil && err != mgo.ErrCursor {
				return err
			}
			if iter.Timeout() {
				continue
			}
			// Either there's no error or the error is an expired
			// cursor. Force recreating the iterator next loop by
			// marking it as nil.
			iter = nil
		}
	}
}

func (t *OplogTailer) dying() bool {
	select {
	case <-t.tomb.Dying():
		return true
	default:
		return false
	}
}
<|endoftext|>


<commit_before>package monit

import (
	"fmt"
	"github.com/tolexo/aero/conf"
	"github.com/tolexo/aero/panik"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"strconv"
	"strings"
	"time"
)

func NotFound(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(404)
	fmt.Fprintf(w, "%s", "4041 page not found")

	dataDogAgent := GetDataDogAgent()
	var statusCode int64 = 4041
	dataDogAgent.Count(FormatHttpStatusCode(statusCode), 1, nil, 1)
}

func FormatHttpStatusCode(httpStatusCode int64) string {
	return "http_" + strconv.FormatInt(httpStatusCode, 10)
}

func GetTimeIntervalTag(respTime float64) (ret []string) {
	interval := conf.String("monitor.interval", "")
	if interval != "" {
		intervalArr := strings.Split(interval, ",")
		var lowerLimit, higherLimit float64
		var matched bool
		for i := range intervalArr {
			higherLimit, _ = strconv.ParseFloat(intervalArr[i], 64)
			if respTime <= higherLimit {
				matched = true
				break
			} else {
				lowerLimit = higherLimit
			}
		}
		var intervalGroup string
		if matched == true {
			intervalGroup = "from_" + strconv.FormatInt(int64(lowerLimit), 10) + "_to_" + strconv.FormatInt(int64(higherLimit), 10)
		} else {
			intervalGroup = "above_" + strconv.FormatInt(int64(higherLimit), 10)
		}
		ret = append(ret, intervalGroup)
	}
	return
}

func ModRecorder() func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			rec := httptest.NewRecorder()
			next.ServeHTTP(rec, r)

			for k, v := range rec.Header() {
				w.Header()[k] = v
			}

			len := rec.Body.Len()
			w.Header().Set("Content-Length", strconv.Itoa(len))
			w.WriteHeader(rec.Code)
			w.Write(rec.Body.Bytes())

			go func() {
				dur := time.Since(start).Seconds() * 1000
				intervalTag := GetTimeIntervalTag(float64(dur))
				statusCode := FormatHttpStatusCode(int64(rec.Code))
				//TODO: Remove debugging info post qa verification
				f, err := os.OpenFile("datadog_event.csv", os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
				panik.On(err)
				l := log.New(f, "", log.LstdFlags)
				l.Printf("%s %f %s", statusCode, dur, intervalTag)

				dataDogAgent := GetDataDogAgent()
				dataDogAgent.Count("throughput", 1, nil, 1)
				dataDogAgent.Count(statusCode, 1, nil, 1)
				dataDogAgent.Histogram("resptime", dur, nil, 1)
				dataDogAgent.Histogram("resptimeinterval", dur, intervalTag, 1)
			}()

		})
	}
}
<commit_msg>TEC-113 log file based debugging info removed<commit_after>package monit

import (
	"fmt"
	"github.com/tolexo/aero/conf"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"time"
)

func NotFound(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(404)
	fmt.Fprintf(w, "%s", "4041 page not found")

	dataDogAgent := GetDataDogAgent()
	var statusCode int64 = 4041
	dataDogAgent.Count(FormatHttpStatusCode(statusCode), 1, nil, 1)
}

func FormatHttpStatusCode(httpStatusCode int64) string {
	return "http_" + strconv.FormatInt(httpStatusCode, 10)
}

func GetTimeIntervalTag(respTime float64) (ret []string) {
	interval := conf.String("monitor.interval", "")
	if interval != "" {
		intervalArr := strings.Split(interval, ",")
		var lowerLimit, higherLimit float64
		var matched bool
		for i := range intervalArr {
			higherLimit, _ = strconv.ParseFloat(intervalArr[i], 64)
			if respTime <= higherLimit {
				matched = true
				break
			} else {
				lowerLimit = higherLimit
			}
		}
		var intervalGroup string
		if matched == true {
			intervalGroup = "from_" + strconv.FormatInt(int64(lowerLimit), 10) + "_to_" + strconv.FormatInt(int64(higherLimit), 10)
		} else {
			intervalGroup = "above_" + strconv.FormatInt(int64(higherLimit), 10)
		}
		ret = append(ret, intervalGroup)
	}
	return
}

func ModRecorder() func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			rec := httptest.NewRecorder()
			next.ServeHTTP(rec, r)

			for k, v := range rec.Header() {
				w.Header()[k] = v
			}

			len := rec.Body.Len()
			w.Header().Set("Content-Length", strconv.Itoa(len))
			w.WriteHeader(rec.Code)
			w.Write(rec.Body.Bytes())

			go func() {
				dur := time.Since(start).Seconds() * 1000
				intervalTag := GetTimeIntervalTag(float64(dur))
				statusCode := FormatHttpStatusCode(int64(rec.Code))

				dataDogAgent := GetDataDogAgent()
				dataDogAgent.Count("throughput", 1, nil, 1)
				dataDogAgent.Count(statusCode, 1, nil, 1)
				dataDogAgent.Histogram("resptime", dur, nil, 1)
				dataDogAgent.Histogram("resptimeinterval", dur, intervalTag, 1)
			}()

		})
	}
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/cri-o/cri-o/server/cri/types"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// ExecSync runs a command in a container synchronously.
func (s *Server) ExecSync(ctx context.Context, req *types.ExecSyncRequest) (*types.ExecSyncResponse, error) {
	c, err := s.GetContainerFromShortID(req.ContainerID)
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "could not find container %q: %v", req.ContainerID, err)
	}

	if err := c.IsAlive(); err != nil {
		return nil, status.Errorf(codes.NotFound, "container is not created or running: %v", err)
	}

	cmd := req.Cmd
	if cmd == nil {
		return nil, errors.New("exec command cannot be empty")
	}

	execResp, err := s.Runtime().ExecSyncContainer(ctx, c, cmd, req.Timeout)
	if err != nil {
		return nil, err
	}
	return &types.ExecSyncResponse{
		Stdout:   execResp.Stdout,
		Stderr:   execResp.Stderr,
		ExitCode: execResp.ExitCode,
	}, nil
}
<commit_msg>server: don't unnecessarily copy struct<commit_after>package server

import (
	"github.com/cri-o/cri-o/server/cri/types"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// ExecSync runs a command in a container synchronously.
func (s *Server) ExecSync(ctx context.Context, req *types.ExecSyncRequest) (*types.ExecSyncResponse, error) {
	c, err := s.GetContainerFromShortID(req.ContainerID)
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "could not find container %q: %v", req.ContainerID, err)
	}

	if err := c.IsAlive(); err != nil {
		return nil, status.Errorf(codes.NotFound, "container is not created or running: %v", err)
	}

	cmd := req.Cmd
	if cmd == nil {
		return nil, errors.New("exec command cannot be empty")
	}

	return s.Runtime().ExecSyncContainer(ctx, c, cmd, req.Timeout)
}
<|endoftext|>


<commit_before>package dockerface

import (
	"bytes"
	"os"
	"os/exec"
	"strings"

	log "github.com/Sirupsen/logrus"
)

func getDaemonHostIp() string {
	var hostIp string

	// try check DOCKER_IP envar first
	value := os.Getenv("DOCKER_IP")
	if value != "" {
		hostIp = strings.TrimSpace(value)
	} else {
		fail := func(err error) {
			log.Fatalf("Error retrieving host ip: %v", err)
		}

		// if that didn't work, get ip of host where this container is running
		phonecmd := exec.Command("bash", "-c", "ip route | awk '/default/ { print $3 }'")

		var outBuff bytes.Buffer
		phonecmd.Stdout = &outBuff
		if err := phonecmd.Start(); err != nil {
			fail(err)
		}
		if err := phonecmd.Wait(); err != nil {
			fail(err)
		}
		hostIp = strings.TrimSpace(string(outBuff.Bytes()))
	}

	return hostIp
}
<commit_msg>add GetEnv util<commit_after>package dockerface

import (
	"bytes"
	"os"
	"os/exec"
	"strings"

	log "github.com/Sirupsen/logrus"
)

func getDaemonHostIp() string {
	var hostIp string

	// try check DOCKER_IP envar first
	value := os.Getenv("DOCKER_IP")
	if value != "" {
		hostIp = strings.TrimSpace(value)
	} else {
		fail := func(err error) {
			log.Fatalf("Error retrieving host ip: %v", err)
		}

		// if that didn't work, get ip of host where this container is running
		phonecmd := exec.Command("bash", "-c", "ip route | awk '/default/ { print $3 }'")

		var outBuff bytes.Buffer
		phonecmd.Stdout = &outBuff
		if err := phonecmd.Start(); err != nil {
			fail(err)
		}
		if err := phonecmd.Wait(); err != nil {
			fail(err)
		}
		hostIp = strings.TrimSpace(string(outBuff.Bytes()))
	}

	return hostIp
}

func GetEnv(name string, alt string) string {
	value := os.Getenv(name)
	if value == "" && alt != "" {
		value = alt
	}
	return value
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/levenlabs/go-srvclient"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s <hostname>\n", os.Args[0])
		os.Exit(1)
	}

	r, err := srvclient.SRV(os.Args[1])
	if err != nil {
		fmt.Fprintf(os.Stderr, "error resolving %q: %s\n", os.Args[1], err)
		os.Exit(2)
	}

	fmt.Println(r)
}
<commit_msg>add --resolvers option to srvclient binary<commit_after>package main

import (
	"fmt"
	"net"
	"os"
	"strings"
	"time"

	"github.com/levenlabs/go-srvclient"
	"github.com/mediocregopher/lever"
)

func main() {
	l := lever.New("srvclient", &lever.Opts{
		HelpHeader:         "Usage: srvclient [options] <hostname>\n",
		DisallowConfigFile: true,
	})
	l.Add(lever.Param{
		Name:        "--resolvers",
		Description: "Comma separated list of resolver ips or addresses (ip:port) which should be used instead of /etc/resolv.conf",
	})
	l.Parse()
	argv := l.ParamRest()

	if len(argv) < 1 {
		fmt.Print(l.Help())
		exit(1)
	}

	sc := srvclient.SRVClient{}
	resolvers, _ := l.ParamStr("--resolvers")
	for _, r := range strings.Split(resolvers, ",") {
		if net.ParseIP(r) != nil {
			r += ":53"
		}
		if r != "" {
			sc.ResolverAddrs = append(sc.ResolverAddrs, r)
		}
	}

	r, err := sc.SRV(argv[0])
	if err != nil {
		fmt.Fprintf(os.Stderr, "error resolving %q: %s\n", os.Args[1], err)
		os.Exit(2)
	}

	fmt.Println(r)
}

func exit(i int) {
	time.Sleep(100 * time.Millisecond)
	os.Exit(i)
}
<|endoftext|>


<commit_before>package client

import (
	"context"

	"google.golang.org/grpc"

	pb "github.com/slok/ragnarok/grpc/nodestatus"
	"github.com/slok/ragnarok/log"
)

// Status interface will implement the required methods to be able to communicate
// with a node status server
type Status interface {
	// RegisterNode registers a node as available on the server
	RegisterNode(id string, tags map[string]string) error
}

// StatusGRPC satisfies Status interface with GRPC communication
type StatusGRPC struct {
	c      pb.NodeStatusClient
	logger log.Logger
}

// NewStatusGRPCFromConnection returns a new Status GRPC client based on the grpc connection
func NewStatusGRPCFromConnection(connection *grpc.ClientConn, logger log.Logger) (*StatusGRPC, error) {
	c := pb.NewNodeStatusClient(connection)
	return NewStatusGRPC(c, logger)
}

// NewStatusGRPC returns a new Status GRPC client
func NewStatusGRPC(client pb.NodeStatusClient, logger log.Logger) (*StatusGRPC, error) {
	logger = logger.WithField("service", "node-status").WithField("service-kind", "grpc")
	return &StatusGRPC{
		c:      client,
		logger: logger,
	}, nil
}

// RegisterNode satisfies Status interface
func (s *StatusGRPC) RegisterNode(id string, tags map[string]string) error {
	logger := s.logger.WithField("call", "register-node").WithField("id", id)
	logger.Debug("making GRPC service call")

	// Create the request objects
	ni := &pb.NodeInfo{
		Node: &pb.Node{
			Id: id,
		},
		Address: "0.0.0.0", // TODO: Set up correct address
		Tags:    tags,
	}

	// Make the request synchronously
	resp, err := s.c.Register(context.Background(), ni)

	// Call error
	if err != nil {
		return err
	}
	// If we have a response then (we should, just checking)
	if resp != nil {
		logger.Debug("call response: %s", resp.Message)
	}

	return nil
}
<commit_msg>Minor changes<commit_after>package client

import (
	"context"

	"google.golang.org/grpc"

	pb "github.com/slok/ragnarok/grpc/nodestatus"
	"github.com/slok/ragnarok/log"
)

// Status interface will implement the required methods to be able to communicate
// with a node status server
type Status interface {
	// RegisterNode registers a node as available on the server
	RegisterNode(id string, tags map[string]string) error
}

// StatusGRPC satisfies Status interface with GRPC communication
type StatusGRPC struct {
	c      pb.NodeStatusClient
	logger log.Logger
}

// NewStatusGRPCFromConnection returns a new Status GRPC client based on the grpc connection
func NewStatusGRPCFromConnection(connection *grpc.ClientConn, logger log.Logger) (*StatusGRPC, error) {
	c := pb.NewNodeStatusClient(connection)
	return NewStatusGRPC(c, logger)
}

// NewStatusGRPC returns a new Status GRPC client
func NewStatusGRPC(client pb.NodeStatusClient, logger log.Logger) (*StatusGRPC, error) {
	logger = logger.WithField("service", "node-status").WithField("service-kind", "grpc")
	return &StatusGRPC{
		c:      client,
		logger: logger,
	}, nil
}

// RegisterNode satisfies Status interface
func (s *StatusGRPC) RegisterNode(id string, tags map[string]string) error {
	logger := s.logger.WithField("call", "register-node").WithField("id", id)
	logger.Debug("making GRPC service call")

	// Create the request objects
	ni := &pb.NodeInfo{
		Node: &pb.Node{
			Id: id,
		},
		Address: "0.0.0.0", // TODO: Set up correct address
		Tags:    tags,
	}

	// Make the request synchronously
	resp, err := s.c.Register(context.Background(), ni)

	// Call error
	if err != nil {
		return err
	}
	// If we have a response then (we should, just checking)
	if resp != nil {
		logger.Debugf("call response: %s", resp.Message)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import . "syscall"
import "fmt"
import "unsafe"

type NetlinkAuditRequest struct {
    Header NlMsghdr
    Data   RtGenmsg
}

func (rr *NetlinkAuditRequest) toWireFormat() []byte {
    b := make([]byte, rr.Header.Len)
    *(*uint32)(unsafe.Pointer(&b[0:4][0])) = rr.Header.Len
    *(*uint16)(unsafe.Pointer(&b[4:6][0])) = rr.Header.Type
    *(*uint16)(unsafe.Pointer(&b[6:8][0])) = rr.Header.Flags
    *(*uint32)(unsafe.Pointer(&b[8:12][0])) = rr.Header.Seq
    *(*uint32)(unsafe.Pointer(&b[12:16][0])) = rr.Header.Pid
    b[16] = byte(rr.Data.Family)
    return b
}

func newNetlinkAuditRequest(proto, seq, family int) []byte {
    rr := &NetlinkAuditRequest{}
    rr.Header.Len = uint32(NLMSG_HDRLEN + SizeofRtGenmsg)
    rr.Header.Type = uint16(AF_UNSPEC)
    rr.Header.Flags = NLM_F_REQUEST
    rr.Header.Seq = uint32(seq)
    rr.Data.Family = uint8(family)
    return rr.toWireFormat()
}

func nlmAlignOf(msglen int) int {
    return (msglen + NLMSG_ALIGNTO - 1) & ^(NLMSG_ALIGNTO - 1)
}

func main() {

    // Create a netlink socket
    s, err := Socket(AF_NETLINK, SOCK_RAW, NETLINK_AUDIT)
    if err != nil {
        fmt.Println("scoket error:")
        fmt.Println(err)
        return
    }
    defer Close(s)

    lsa := &SockaddrNetlink{Family: AF_NETLINK}
    if err := Bind(s, lsa); err != nil {
        fmt.Println("bind error:")
        fmt.Println(err)
        return
    }
    fmt.Println("**** Starting Netlink")

    // this should be 10 or not??
    wb := newNetlinkAuditRequest(1000, 1, NETLINK_AUDIT)
    if err := Sendto(s, wb, 0, lsa); err != nil {
        fmt.Println("sending error: ", err)
        return
    }
    var tab []byte
done:
    for {

        rb := make([]byte, Getpagesize())
        nr, _, err := Recvfrom(s, rb, 0)
        if err != nil {
            fmt.Println("rec err: ", err)
            return
        }
        if nr < NLMSG_HDRLEN {
            fmt.Println(EINVAL)
            fmt.Println("nr < nlmsg hdrlen")
            return
        }
        rb = rb[:nr]
        tab = append(tab, rb...)
        msgs, err := ParseNetlinkMessage(rb)

        if err != nil {
            fmt.Println("parse error:", err)
            return
        }

        for _, m := range msgs {

            lsa, err := Getsockname(s)
            if err != nil {
                fmt.Println("get socket name error:", err)
                return
            }
            switch v := lsa.(type) {
            case *SockaddrNetlink:
                if m.Header.Seq != 1 || m.Header.Pid != v.Pid {
                    fmt.Println("case")
                    fmt.Println(EINVAL)
                    return
                }
            default:
                fmt.Println(EINVAL)
                fmt.Println("default case")
                return
            }
            if m.Header.Type == NLMSG_DONE {
                break done
            }

            if m.Header.Type == NLMSG_ERROR {
                fmt.Println(EINVAL)
                fmt.Println("somewhere below")
                return
            }

        }
    }
    fmt.Println(tab);
}
<commit_msg>few fixes for proto<commit_after>package main

import s "syscall"
import "fmt"
import "unsafe"

type NetlinkAuditRequest struct {
    Header s.NlMsghdr
    Data   byte  //s.RtGenmsg
}

func (rr *NetlinkAuditRequest) toWireFormat() []byte {
    b := make([]byte, rr.Header.Len)
    *(*uint32)(unsafe.Pointer(&b[0:4][0])) = rr.Header.Len
    *(*uint16)(unsafe.Pointer(&b[4:6][0])) = rr.Header.Type
    *(*uint16)(unsafe.Pointer(&b[6:8][0])) = rr.Header.Flags
    *(*uint32)(unsafe.Pointer(&b[8:12][0])) = rr.Header.Seq
    *(*uint32)(unsafe.Pointer(&b[12:16][0])) = rr.Header.Pid
    //b[16] = byte(rr.Data.Family)
    return b
}

func newNetlinkAuditRequest(proto, seq, family int) []byte {
    rr := &NetlinkAuditRequest{}
    rr.Header.Len = uint32(s.NLMSG_HDRLEN) //+ s.SizeofRtGenmsg)
    rr.Header.Type = uint16(proto)
    rr.Header.Flags = s.NLM_F_REQUEST//s.MSG_PEEK|s.MSG_DONTWAIT//s.NLM_F_REQUEST //| s.NLM_F_ACK
    rr.Header.Seq = uint32(seq)
    //rr.Data.Family = uint8(family)
    return rr.toWireFormat()
}

func nlmAlignOf(msglen int) int {
    return (msglen + s.NLMSG_ALIGNTO - 1) & ^(s.NLMSG_ALIGNTO - 1)
}

func main() {

    // Create a netlink socket
    sock, err := s.Socket(s.AF_NETLINK, s.SOCK_RAW, s.NETLINK_AUDIT)
    if err != nil {
        fmt.Println("scoket error:")
        fmt.Println(err)
        return
    }
    defer s.Close(sock)

    lsa := &s.SockaddrNetlink{Family: s.AF_NETLINK}
    if err := s.Bind(sock, lsa); err != nil {
        fmt.Println("bind error:")
        fmt.Println(err)
        return
    }
    fmt.Println("**** Starting Netlink")

    // sendto(fd, &req, req.nlh.nlmsg_len, 0,(struct sockaddr*)&addr, sizeof(addr))
    // sendto(3, "\20\0\0\0\350\3\5\0\1\0\0\0\0\0\0\0", 16, 0, {sa_family=AF_NETLINK, pid=0, groups=00000000}, 12) = 16
    // sendto(3, "\20\0\0\0\350\3\1\0\1\0\0\0\0\0\0\0", 16, 0, {sa_family=AF_NETLINK, pid=0, groups=00000000}, 12) = 16
    wb := newNetlinkAuditRequest(1000, 1, s.NETLINK_AUDIT)
    if err := s.Sendto(sock, wb, 0, lsa); err != nil {
        fmt.Println("sending error: ", err)
        return
    }
    var tab []byte
done:
    for {

        // recvfrom(fd, &rep->msg, sizeof(rep->msg), block|peek,   (struct sockaddr*)&nladdr, &nladdrlen);
        // recvfrom(3, "$\0\0\0\2\0\0\0\1\0\0\0\234K\0\0\0\0\0\0\20\0\0\0\350\3\5\0\1\0\0\0"..., 8988, MSG_PEEK|MSG_DONTWAIT, {sa_family=AF_NETLINK, pid=0, groups=00000000}, [12]) = 36
        // recvfrom(3, "0\0\0\0\350\3\0\0\1\0\0\0r\177\0\0\0\0\0\0\0\0\0\0\1\0\0\0\0\0\0\0"..., 8988, MSG_PEEK|MSG_DONTWAIT, {sa_family=AF_NETLINK, pid=0, groups=00000000}, [12]) = 48
        rb := make([]byte, s.Getpagesize())

        //nr, _, err := Recvfrom(s, rb, MSG_PEEK|MSG_DONTWAIT)
        nr, _, err := s.Recvfrom(sock, rb, 0)
        if err != nil {
            fmt.Println("rec err: ", err)
            return
        }
        if nr < s.NLMSG_HDRLEN {
            fmt.Println(s.EINVAL)
            fmt.Println("nr < nlmsg hdrlen")
            return
        }
        rb = rb[:nr]
        tab = append(tab, rb...)
        msgs, err := s.ParseNetlinkMessage(rb)

        if err != nil {
            fmt.Println("parse error:", err)
            return
        }

        for _, m := range msgs {

            fmt.Println(m.Header.Pid)

            lsa, err := s.Getsockname(sock)
            if err != nil {
                fmt.Println("get socket name error:", err)
                return
            }

            switch v := lsa.(type) {
            case *s.SockaddrNetlink:
                if m.Header.Seq != 1 || m.Header.Pid != v.Pid {
                    fmt.Println("case")
                    fmt.Println(s.EINVAL)
                    return
                }
            default:
                fmt.Println(s.EINVAL)
                fmt.Println("default case")
                return
            }
            if m.Header.Type == s.NLMSG_DONE {
                break done
            }

            if m.Header.Type == s.NLMSG_ERROR {
                fmt.Println(s.EINVAL)
                fmt.Println("somewhere below")
                return
            }

        }
        fmt.Println(string(tab[:]));
    }
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/prometheus/procfs"
)

var (
	binary = filepath.Join(os.Getenv("GOPATH"), "bin/node_exporter")
)
const (
	address = "localhost:19100"
)

func TestFileDescriptorLeak(t *testing.T) {
	if _, err := os.Stat(binary); err != nil {
		t.Skipf("node_exporter binary not available, try to run `make build` first: %s", err)
	}
	if _, err := procfs.NewStat(); err != nil {
		t.Skipf("proc filesystem is not available, but currently required to read number of open file descriptors: %s", err)
	}

	exporter := exec.Command(binary, "--web.listen-address", address)
	test := func(pid int) error {
		if err := queryExporter(address); err != nil {
			return err
		}
		proc, err := procfs.NewProc(pid)
		if err != nil {
			return err
		}
		fdsBefore, err := proc.FileDescriptors()
		if err != nil {
			return err
		}
		for i := 0; i < 5; i++ {
			if err := queryExporter(address); err != nil {
				return err
			}
		}
		fdsAfter, err := proc.FileDescriptors()
		if err != nil {
			return err
		}
		if want, have := len(fdsBefore), len(fdsAfter); want != have {
			return fmt.Errorf("want %d open file descriptors after metrics scrape, have %d", want, have)
		}
		return nil
	}

	if err := runCommandAndTests(exporter, address, test); err != nil {
		t.Error(err)
	}
}

func TestHandlingOfDuplicatedMetrics(t *testing.T) {
	if _, err := os.Stat(binary); err != nil {
		t.Skipf("node_exporter binary not available, try to run `make build` first: %s", err)
	}

	dir, err := ioutil.TempDir("", "node-exporter")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	content := []byte("dummy_metric 1\n")
	if err := ioutil.WriteFile(filepath.Join(dir, "a.prom"), content, 0600); err != nil {
		t.Fatal(err)
	}
	if err := ioutil.WriteFile(filepath.Join(dir, "b.prom"), content, 0600); err != nil {
		t.Fatal(err)
	}

	exporter := exec.Command(binary, "--web.listen-address", address, "--collector.textfile.directory", dir)
	test := func(_ int) error {
		return queryExporter(address)
	}

	if err := runCommandAndTests(exporter, address, test); err != nil {
		t.Error(err)
	}
}

func queryExporter(address string) error {
	resp, err := http.Get(fmt.Sprintf("http://%s/metrics", address))
	if err != nil {
		return err
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	if err := resp.Body.Close(); err != nil {
		return err
	}
	if want, have := http.StatusOK, resp.StatusCode; want != have {
		return fmt.Errorf("want /metrics status code %d, have %d. Body:\n%s", want, have, b)
	}
	return nil
}

func runCommandAndTests(cmd *exec.Cmd, address string, fn func(pid int) error) error {
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("failed to start command: %s", err)
	}
	time.Sleep(50 * time.Millisecond)
	for i := 0; i < 10; i++ {
		if err := queryExporter(address); err == nil {
			break
		}
		time.Sleep(500 * time.Millisecond)
		if cmd.Process== nil || i == 9 {
			return fmt.Errorf("can't start command")
		}
	}

	errc := make(chan error)
	go func(pid int) {
		errc <- fn(pid)
	}(cmd.Process.Pid)

	select {
	case err := <-errc:
		if cmd.Process != nil {
			cmd.Process.Kill()
		}
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>gofmt node_exporter_test<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/prometheus/procfs"
)

var (
	binary = filepath.Join(os.Getenv("GOPATH"), "bin/node_exporter")
)

const (
	address = "localhost:19100"
)

func TestFileDescriptorLeak(t *testing.T) {
	if _, err := os.Stat(binary); err != nil {
		t.Skipf("node_exporter binary not available, try to run `make build` first: %s", err)
	}
	if _, err := procfs.NewStat(); err != nil {
		t.Skipf("proc filesystem is not available, but currently required to read number of open file descriptors: %s", err)
	}

	exporter := exec.Command(binary, "--web.listen-address", address)
	test := func(pid int) error {
		if err := queryExporter(address); err != nil {
			return err
		}
		proc, err := procfs.NewProc(pid)
		if err != nil {
			return err
		}
		fdsBefore, err := proc.FileDescriptors()
		if err != nil {
			return err
		}
		for i := 0; i < 5; i++ {
			if err := queryExporter(address); err != nil {
				return err
			}
		}
		fdsAfter, err := proc.FileDescriptors()
		if err != nil {
			return err
		}
		if want, have := len(fdsBefore), len(fdsAfter); want != have {
			return fmt.Errorf("want %d open file descriptors after metrics scrape, have %d", want, have)
		}
		return nil
	}

	if err := runCommandAndTests(exporter, address, test); err != nil {
		t.Error(err)
	}
}

func TestHandlingOfDuplicatedMetrics(t *testing.T) {
	if _, err := os.Stat(binary); err != nil {
		t.Skipf("node_exporter binary not available, try to run `make build` first: %s", err)
	}

	dir, err := ioutil.TempDir("", "node-exporter")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)

	content := []byte("dummy_metric 1\n")
	if err := ioutil.WriteFile(filepath.Join(dir, "a.prom"), content, 0600); err != nil {
		t.Fatal(err)
	}
	if err := ioutil.WriteFile(filepath.Join(dir, "b.prom"), content, 0600); err != nil {
		t.Fatal(err)
	}

	exporter := exec.Command(binary, "--web.listen-address", address, "--collector.textfile.directory", dir)
	test := func(_ int) error {
		return queryExporter(address)
	}

	if err := runCommandAndTests(exporter, address, test); err != nil {
		t.Error(err)
	}
}

func queryExporter(address string) error {
	resp, err := http.Get(fmt.Sprintf("http://%s/metrics", address))
	if err != nil {
		return err
	}
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	if err := resp.Body.Close(); err != nil {
		return err
	}
	if want, have := http.StatusOK, resp.StatusCode; want != have {
		return fmt.Errorf("want /metrics status code %d, have %d. Body:\n%s", want, have, b)
	}
	return nil
}

func runCommandAndTests(cmd *exec.Cmd, address string, fn func(pid int) error) error {
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("failed to start command: %s", err)
	}
	time.Sleep(50 * time.Millisecond)
	for i := 0; i < 10; i++ {
		if err := queryExporter(address); err == nil {
			break
		}
		time.Sleep(500 * time.Millisecond)
		if cmd.Process == nil || i == 9 {
			return fmt.Errorf("can't start command")
		}
	}

	errc := make(chan error)
	go func(pid int) {
		errc <- fn(pid)
	}(cmd.Process.Pid)

	select {
	case err := <-errc:
		if cmd.Process != nil {
			cmd.Process.Kill()
		}
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package nsqcrawl

import (
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/crackcomm/crawl"
	"github.com/crackcomm/nsqueue/consumer"
	"github.com/crackcomm/nsqueue/producer"
)

// NewQueue - Creates nsq consumer and producer.
func NewQueue(topic, channel string, maxInFlight int) *Queue {
	q := &Queue{
		Consumer: consumer.New(),
		Producer: producer.New(),
		memq:     crawl.NewQueue(maxInFlight + 1),
		topic:    topic,
	}
	q.Consumer.Register(topic, channel, maxInFlight, q.nsqHandler)
	return q
}

// NewProducer - Creates queue producer.
func NewProducer(topic string) *Queue {
	return &Queue{
		Producer: producer.New(),
		topic:    topic,
	}
}

// Queue - NSQ Queue.
type Queue struct {
	*consumer.Consumer
	*producer.Producer

	topic string
	memq  crawl.Queue
}

// Schedule - Schedules job in nsq.
// It will not call job.Done ever.
func (queue *Queue) Schedule(ctx context.Context, r *crawl.Request) (err error) {
	r := &request{Request: r}
	if deadline, ok := ctx.Deadline(); ok {
		r.Deadline = deadline
	}
	return queue.Producer.PublishJSON(queue.topic, r)
}

// Get - Gets job from channel.
func (queue *Queue) Get() (job crawl.Job, err error) {
	return queue.memq.Get()
}

// Close - Closes consumer and producer.
func (queue *Queue) Close() (err error) {
	if queue.Producer != nil {
		queue.Producer.Stop()
	}
	if queue.Consumer != nil {
		queue.Consumer.Stop()
	}
	if queue.memq != nil {
		err = queue.memq.Close()
	}
	return
}

func (queue *Queue) nsqHandler(msg *consumer.Message) {
	req := new(request)
	err := msg.ReadJSON(req)
	if err != nil {
		glog.V(3).Infof("nsq json (%s) error: %v", msg.Body, err)
		msg.GiveUp()
		return
	}

	// Check if deadline exceeded
	if !req.Deadline.IsZero() && time.Now().After(req.Deadline) {
		glog.V(3).Infof("request deadline exceeded (%s)", msg.Body)
		msg.GiveUp()
		return
	}

	// Request context
	ctx := context.Background()

	// Set request deadline
	if !req.Deadline.IsZero() {
		ctx, _ = context.WithDeadline(ctx, req.Deadline)
	}

	// Set nsq message in context
	ctx = consumer.WithMessage(ctx, msg)

	// Schedule job in memory
	queue.memq.Schedule(&nsqJob{msg: msg, req: req.Request, ctx: ctx})
}

type request struct {
	Request  *crawl.Request `json:"request,omitempty"`
	Deadline time.Time      `json:"deadline,omitempty"`
}

type nsqJob struct {
	msg *consumer.Message
	req *crawl.Request
	ctx context.Context
}

func (job *nsqJob) Context() context.Context { return job.ctx }
func (job *nsqJob) Request() *crawl.Request  { return job.req }
func (job *nsqJob) Done()                    { job.msg.Success() }
<commit_msg>nsq queue: use internal channel<commit_after>package nsqcrawl

import (
	"io"
	"time"

	"github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/crackcomm/crawl"
	"github.com/crackcomm/nsqueue/consumer"
	"github.com/crackcomm/nsqueue/producer"
)

// NewQueue - Creates nsq consumer and producer.
func NewQueue(topic, channel string, maxInFlight int) *Queue {
	q := &Queue{
		Consumer: consumer.New(),
		Producer: producer.New(),
		channel:  make(chan *nsqJob, maxInFlight+1),
		topic:    topic,
	}
	q.Consumer.Register(topic, channel, maxInFlight, q.nsqHandler)
	return q
}

// NewProducer - Creates queue producer.
func NewProducer(topic string) *Queue {
	return &Queue{
		Producer: producer.New(),
		topic:    topic,
	}
}

// Queue - NSQ Queue.
type Queue struct {
	*consumer.Consumer
	*producer.Producer

	topic   string
	channel chan *nsqJob
}

// Schedule - Schedules job in nsq.
// It will not call job.Done ever.
func (queue *Queue) Schedule(ctx context.Context, req *crawl.Request) (err error) {
	r := &request{Request: req}
	if deadline, ok := ctx.Deadline(); ok {
		r.Deadline = deadline
	}
	return queue.Producer.PublishJSON(queue.topic, r)
}

// Get - Gets job from channel.
func (queue *Queue) Get() (crawl.Job, error) {
	job, ok := <-queue.channel
	if !ok {
		return nil, io.EOF
	}
	return job, nil
}

// Close - Closes consumer and producer.
func (queue *Queue) Close() (err error) {
	if queue.Producer != nil {
		queue.Producer.Stop()
	}
	if queue.Consumer != nil {
		queue.Consumer.Stop()
	}
	if queue.channel != nil {
		close(queue.channel)
	}
	return
}

func (queue *Queue) nsqHandler(msg *consumer.Message) {
	req := new(request)
	err := msg.ReadJSON(req)
	if err != nil {
		glog.V(3).Infof("nsq json (%s) error: %v", msg.Body, err)
		msg.GiveUp()
		return
	}

	// Check if deadline exceeded
	if !req.Deadline.IsZero() && time.Now().After(req.Deadline) {
		glog.V(3).Infof("request deadline exceeded (%s)", msg.Body)
		msg.GiveUp()
		return
	}

	// Request context
	ctx := context.Background()

	// Set request deadline
	if !req.Deadline.IsZero() {
		ctx, _ = context.WithDeadline(ctx, req.Deadline)
	}

	// Set nsq message in context
	ctx = consumer.WithMessage(ctx, msg)

	// Schedule job in memory
	queue.channel <- &nsqJob{msg: msg, req: req.Request, ctx: ctx}
}

type request struct {
	Request  *crawl.Request `json:"request,omitempty"`
	Deadline time.Time      `json:"deadline,omitempty"`
}

type nsqJob struct {
	msg *consumer.Message
	req *crawl.Request
	ctx context.Context
}

func (job *nsqJob) Context() context.Context { return job.ctx }
func (job *nsqJob) Request() *crawl.Request  { return job.req }
func (job *nsqJob) Done()                    { job.msg.Success() }
<|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"regexp"
	"strings"
	"time"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-docopt"
	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
	"github.com/flynn/flynn/host/types"
	"github.com/flynn/flynn/pkg/cluster"
	"github.com/flynn/flynn/pkg/exec"
	"github.com/flynn/flynn/pkg/random"
	"github.com/flynn/flynn/pkg/version"
)

var clusterc = cluster.NewClient()

func init() {
	log.SetFlags(0)
}

var typesPattern = regexp.MustCompile("types.* -> (.+)\n")

const blobstoreURL = "http://blobstore.discoverd"
const scaleTimeout = 20 * time.Second

func parsePairs(args *docopt.Args, str string) (map[string]string, error) {
	pairs := args.All[str].([]string)
	item := make(map[string]string, len(pairs))
	for _, s := range pairs {
		v := strings.SplitN(s, "=", 2)
		if len(v) != 2 {
			return nil, fmt.Errorf("invalid var format: %q", s)
		}
		item[v[0]] = v[1]
	}
	return item, nil
}

func main() {
	client, err := controller.NewClient("", os.Getenv("CONTROLLER_KEY"))
	if err != nil {
		log.Fatalln("Unable to connect to controller:", err)
	}

	usage := `
Usage: flynn-receiver <app> <rev> [-e <var>=<val>]... [-m <key>=<val>]...

Options:
	-e,--env <var>=<val>
	-m,--meta <key>=<val>
`[1:]
	args, _ := docopt.Parse(usage, nil, true, version.String(), false)

	appName := args.String["<app>"]
	env, err := parsePairs(args, "--env")
	if err != nil {
		log.Fatal(err)
	}
	meta, err := parsePairs(args, "--meta")
	if err != nil {
		log.Fatal(err)
	}

	app, err := client.GetApp(appName)
	if err == controller.ErrNotFound {
		log.Fatalf("Unknown app %q", appName)
	} else if err != nil {
		log.Fatalln("Error retrieving app:", err)
	}
	prevRelease, err := client.GetAppRelease(app.Name)
	if err == controller.ErrNotFound {
		prevRelease = &ct.Release{}
	} else if err != nil {
		log.Fatalln("Error getting current app release:", err)
	}

	fmt.Printf("-----> Building %s...\n", app.Name)

	var output bytes.Buffer
	slugURL := fmt.Sprintf("%s/%s.tgz", blobstoreURL, random.UUID())
	cmd := exec.Command(exec.DockerImage(os.Getenv("SLUGBUILDER_IMAGE_URI")), slugURL)
	cmd.Stdout = io.MultiWriter(os.Stdout, &output)
	cmd.Stderr = os.Stderr
	cmd.Meta = map[string]string{
		"flynn-controller.app":      app.ID,
		"flynn-controller.app_name": app.Name,
		"flynn-controller.release":  prevRelease.ID,
		"flynn-controller.type":     "slugbuilder",
	}
	if len(prevRelease.Env) > 0 {
		stdin, err := cmd.StdinPipe()
		if err != nil {
			log.Fatalln(err)
		}
		go appendEnvDir(os.Stdin, stdin, prevRelease.Env)
	} else {
		cmd.Stdin = os.Stdin
	}
	cmd.Env = make(map[string]string)
	cmd.Env["BUILD_CACHE_URL"] = fmt.Sprintf("%s/%s-cache.tgz", blobstoreURL, app.ID)
	if buildpackURL, ok := env["BUILDPACK_URL"]; ok {
		cmd.Env["BUILDPACK_URL"] = buildpackURL
	} else if buildpackURL, ok := prevRelease.Env["BUILDPACK_URL"]; ok {
		cmd.Env["BUILDPACK_URL"] = buildpackURL
	}
	for _, k := range []string{"SSH_CLIENT_KEY", "SSH_CLIENT_HOSTS"} {
		if v := os.Getenv(k); v != "" {
			cmd.Env[k] = v
		}
	}

	if err := cmd.Run(); err != nil {
		log.Fatalln("Build failed:", err)
	}

	var types []string
	if match := typesPattern.FindSubmatch(output.Bytes()); match != nil {
		types = strings.Split(string(match[1]), ", ")
	}

	fmt.Printf("-----> Creating release...\n")

	artifact := &ct.Artifact{Type: "docker", URI: os.Getenv("SLUGRUNNER_IMAGE_URI")}
	if err := client.CreateArtifact(artifact); err != nil {
		log.Fatalln("Error creating artifact:", err)
	}

	release := &ct.Release{
		ArtifactID: artifact.ID,
		Env:        prevRelease.Env,
		Meta:       prevRelease.Meta,
	}
	if release.Meta == nil {
		release.Meta = make(map[string]string, len(meta))
	}
	if release.Env == nil {
		release.Env = make(map[string]string, len(env))
	}
	for k, v := range env {
		release.Env[k] = v
	}
	for k, v := range meta {
		release.Meta[k] = v
	}
	procs := make(map[string]ct.ProcessType)
	for _, t := range types {
		proc := prevRelease.Processes[t]
		proc.Cmd = []string{"start", t}
		if t == "web" || strings.HasSuffix(t, "-web") {
			proc.Service = app.Name + "-" + t
			proc.Ports = []ct.Port{{
				Port:  8080,
				Proto: "tcp",
				Service: &host.Service{
					Name:   proc.Service,
					Create: true,
					Check:  &host.HealthCheck{Type: "tcp"},
				},
			}}
		}
		procs[t] = proc
	}
	release.Processes = procs
	if release.Env == nil {
		release.Env = make(map[string]string)
	}
	release.Env["SLUG_URL"] = slugURL

	if err := client.CreateRelease(release); err != nil {
		log.Fatalln("Error creating release:", err)
	}
	if err := client.DeployAppRelease(app.Name, release.ID); err != nil {
		log.Fatalln("Error deploying app release:", err)
	}

	fmt.Println("=====> Application deployed")

	if needsDefaultScale(app.ID, prevRelease.ID, procs, client) {
		formation := &ct.Formation{
			AppID:     app.ID,
			ReleaseID: release.ID,
			Processes: map[string]int{"web": 1},
		}

		watcher, err := client.WatchJobEvents(app.ID, release.ID)
		if err != nil {
			log.Fatalln("Error streaming job events", err)
			return
		}
		defer watcher.Close()

		if err := client.PutFormation(formation); err != nil {
			log.Fatalln("Error putting formation:", err)
		}
		fmt.Println("=====> Waiting for web job to start...")

		err = watcher.WaitFor(ct.JobEvents{"web": ct.JobUpEvents(1)}, scaleTimeout, func(e *ct.Job) error {
			switch e.State {
			case ct.JobStateUp:
				fmt.Println("=====> Default web formation scaled to 1")
			case ct.JobStateDown:
				return fmt.Errorf("Failed to scale web process type")
			}
			return nil
		})
		if err != nil {
			log.Fatalln(err.Error())
		}
	}
}

// needsDefaultScale indicates whether a release needs a default scale based on
// whether it has a web process type and either has no previous release or no
// previous scale.
func needsDefaultScale(appID, prevReleaseID string, procs map[string]ct.ProcessType, client *controller.Client) bool {
	if _, ok := procs["web"]; !ok {
		return false
	}
	if prevReleaseID == "" {
		return true
	}
	_, err := client.GetFormation(appID, prevReleaseID)
	return err == controller.ErrNotFound
}

func appendEnvDir(stdin io.Reader, pipe io.WriteCloser, env map[string]string) {
	defer pipe.Close()
	tr := tar.NewReader(stdin)
	tw := tar.NewWriter(pipe)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			log.Fatalln(err)
		}
		hdr.Name = path.Join("app", hdr.Name)
		if err := tw.WriteHeader(hdr); err != nil {
			log.Fatalln(err)
		}
		if _, err := io.Copy(tw, tr); err != nil {
			log.Fatalln(err)
		}
	}
	// append env dir
	for key, value := range env {
		hdr := &tar.Header{
			Name:    path.Join("env", key),
			Mode:    0644,
			ModTime: time.Now(),
			Size:    int64(len(value)),
		}

		if err := tw.WriteHeader(hdr); err != nil {
			log.Fatalln(err)
		}
		if _, err := tw.Write([]byte(value)); err != nil {
			log.Fatalln(err)
		}
	}
	hdr := &tar.Header{
		Name:    ".ENV_DIR_bdca46b87df0537eaefe79bb632d37709ff1df18",
		Mode:    0400,
		ModTime: time.Now(),
		Size:    0,
	}
	if err := tw.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
}
<commit_msg>gitreceive/receiver: Move slugbuilder jobs to background partition<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"regexp"
	"strings"
	"time"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-docopt"
	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
	"github.com/flynn/flynn/host/types"
	"github.com/flynn/flynn/pkg/cluster"
	"github.com/flynn/flynn/pkg/exec"
	"github.com/flynn/flynn/pkg/random"
	"github.com/flynn/flynn/pkg/version"
)

var clusterc = cluster.NewClient()

func init() {
	log.SetFlags(0)
}

var typesPattern = regexp.MustCompile("types.* -> (.+)\n")

const blobstoreURL = "http://blobstore.discoverd"
const scaleTimeout = 20 * time.Second

func parsePairs(args *docopt.Args, str string) (map[string]string, error) {
	pairs := args.All[str].([]string)
	item := make(map[string]string, len(pairs))
	for _, s := range pairs {
		v := strings.SplitN(s, "=", 2)
		if len(v) != 2 {
			return nil, fmt.Errorf("invalid var format: %q", s)
		}
		item[v[0]] = v[1]
	}
	return item, nil
}

func main() {
	client, err := controller.NewClient("", os.Getenv("CONTROLLER_KEY"))
	if err != nil {
		log.Fatalln("Unable to connect to controller:", err)
	}

	usage := `
Usage: flynn-receiver <app> <rev> [-e <var>=<val>]... [-m <key>=<val>]...

Options:
	-e,--env <var>=<val>
	-m,--meta <key>=<val>
`[1:]
	args, _ := docopt.Parse(usage, nil, true, version.String(), false)

	appName := args.String["<app>"]
	env, err := parsePairs(args, "--env")
	if err != nil {
		log.Fatal(err)
	}
	meta, err := parsePairs(args, "--meta")
	if err != nil {
		log.Fatal(err)
	}

	app, err := client.GetApp(appName)
	if err == controller.ErrNotFound {
		log.Fatalf("Unknown app %q", appName)
	} else if err != nil {
		log.Fatalln("Error retrieving app:", err)
	}
	prevRelease, err := client.GetAppRelease(app.Name)
	if err == controller.ErrNotFound {
		prevRelease = &ct.Release{}
	} else if err != nil {
		log.Fatalln("Error getting current app release:", err)
	}

	fmt.Printf("-----> Building %s...\n", app.Name)

	jobEnv := make(map[string]string)
	jobEnv["BUILD_CACHE_URL"] = fmt.Sprintf("%s/%s-cache.tgz", blobstoreURL, app.ID)
	if buildpackURL, ok := env["BUILDPACK_URL"]; ok {
		jobEnv["BUILDPACK_URL"] = buildpackURL
	} else if buildpackURL, ok := prevRelease.Env["BUILDPACK_URL"]; ok {
		jobEnv["BUILDPACK_URL"] = buildpackURL
	}
	for _, k := range []string{"SSH_CLIENT_KEY", "SSH_CLIENT_HOSTS"} {
		if v := os.Getenv(k); v != "" {
			jobEnv[k] = v
		}
	}
	slugURL := fmt.Sprintf("%s/%s.tgz", blobstoreURL, random.UUID())

	cmd := exec.Job(exec.DockerImage(os.Getenv("SLUGBUILDER_IMAGE_URI")), &host.Job{
		Config: host.ContainerConfig{
			Cmd:        []string{slugURL},
			Env:        jobEnv,
			Stdin:      true,
			DisableLog: true,
		},
		Partition: "background",
		Metadata: map[string]string{
			"flynn-controller.app":      app.ID,
			"flynn-controller.app_name": app.Name,
			"flynn-controller.release":  prevRelease.ID,
			"flynn-controller.type":     "slugbuilder",
		},
	})
	var output bytes.Buffer
	cmd.Stdout = io.MultiWriter(os.Stdout, &output)
	cmd.Stderr = os.Stderr

	if len(prevRelease.Env) > 0 {
		stdin, err := cmd.StdinPipe()
		if err != nil {
			log.Fatalln(err)
		}
		go appendEnvDir(os.Stdin, stdin, prevRelease.Env)
	} else {
		cmd.Stdin = os.Stdin
	}

	if err := cmd.Run(); err != nil {
		log.Fatalln("Build failed:", err)
	}

	var types []string
	if match := typesPattern.FindSubmatch(output.Bytes()); match != nil {
		types = strings.Split(string(match[1]), ", ")
	}

	fmt.Printf("-----> Creating release...\n")

	artifact := &ct.Artifact{Type: "docker", URI: os.Getenv("SLUGRUNNER_IMAGE_URI")}
	if err := client.CreateArtifact(artifact); err != nil {
		log.Fatalln("Error creating artifact:", err)
	}

	release := &ct.Release{
		ArtifactID: artifact.ID,
		Env:        prevRelease.Env,
		Meta:       prevRelease.Meta,
	}
	if release.Meta == nil {
		release.Meta = make(map[string]string, len(meta))
	}
	if release.Env == nil {
		release.Env = make(map[string]string, len(env))
	}
	for k, v := range env {
		release.Env[k] = v
	}
	for k, v := range meta {
		release.Meta[k] = v
	}
	procs := make(map[string]ct.ProcessType)
	for _, t := range types {
		proc := prevRelease.Processes[t]
		proc.Cmd = []string{"start", t}
		if t == "web" || strings.HasSuffix(t, "-web") {
			proc.Service = app.Name + "-" + t
			proc.Ports = []ct.Port{{
				Port:  8080,
				Proto: "tcp",
				Service: &host.Service{
					Name:   proc.Service,
					Create: true,
					Check:  &host.HealthCheck{Type: "tcp"},
				},
			}}
		}
		procs[t] = proc
	}
	release.Processes = procs
	if release.Env == nil {
		release.Env = make(map[string]string)
	}
	release.Env["SLUG_URL"] = slugURL

	if err := client.CreateRelease(release); err != nil {
		log.Fatalln("Error creating release:", err)
	}
	if err := client.DeployAppRelease(app.Name, release.ID); err != nil {
		log.Fatalln("Error deploying app release:", err)
	}

	fmt.Println("=====> Application deployed")

	if needsDefaultScale(app.ID, prevRelease.ID, procs, client) {
		formation := &ct.Formation{
			AppID:     app.ID,
			ReleaseID: release.ID,
			Processes: map[string]int{"web": 1},
		}

		watcher, err := client.WatchJobEvents(app.ID, release.ID)
		if err != nil {
			log.Fatalln("Error streaming job events", err)
			return
		}
		defer watcher.Close()

		if err := client.PutFormation(formation); err != nil {
			log.Fatalln("Error putting formation:", err)
		}
		fmt.Println("=====> Waiting for web job to start...")

		err = watcher.WaitFor(ct.JobEvents{"web": ct.JobUpEvents(1)}, scaleTimeout, func(e *ct.Job) error {
			switch e.State {
			case ct.JobStateUp:
				fmt.Println("=====> Default web formation scaled to 1")
			case ct.JobStateDown:
				return fmt.Errorf("Failed to scale web process type")
			}
			return nil
		})
		if err != nil {
			log.Fatalln(err.Error())
		}
	}
}

// needsDefaultScale indicates whether a release needs a default scale based on
// whether it has a web process type and either has no previous release or no
// previous scale.
func needsDefaultScale(appID, prevReleaseID string, procs map[string]ct.ProcessType, client *controller.Client) bool {
	if _, ok := procs["web"]; !ok {
		return false
	}
	if prevReleaseID == "" {
		return true
	}
	_, err := client.GetFormation(appID, prevReleaseID)
	return err == controller.ErrNotFound
}

func appendEnvDir(stdin io.Reader, pipe io.WriteCloser, env map[string]string) {
	defer pipe.Close()
	tr := tar.NewReader(stdin)
	tw := tar.NewWriter(pipe)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			log.Fatalln(err)
		}
		hdr.Name = path.Join("app", hdr.Name)
		if err := tw.WriteHeader(hdr); err != nil {
			log.Fatalln(err)
		}
		if _, err := io.Copy(tw, tr); err != nil {
			log.Fatalln(err)
		}
	}
	// append env dir
	for key, value := range env {
		hdr := &tar.Header{
			Name:    path.Join("env", key),
			Mode:    0644,
			ModTime: time.Now(),
			Size:    int64(len(value)),
		}

		if err := tw.WriteHeader(hdr); err != nil {
			log.Fatalln(err)
		}
		if _, err := tw.Write([]byte(value)); err != nil {
			log.Fatalln(err)
		}
	}
	hdr := &tar.Header{
		Name:    ".ENV_DIR_bdca46b87df0537eaefe79bb632d37709ff1df18",
		Mode:    0400,
		ModTime: time.Now(),
		Size:    0,
	}
	if err := tw.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
}
<|endoftext|>


<commit_before>package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
)

type Broker struct {
	Name               string
	ServiceGenericName string
	IP                 string
	Port               int
	CertFile           string
	KeyFile            string
	AuthExchange       string
	AuthAllExchange    string
	WebProtocol        string
}

type Config struct {
	BuildNumber int
	Environment string
	Regions     struct {
		Vagrant string
		SJ      string
		AWS     string
		Premium string
	}
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	VmPool          string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
		RuntimeOptions     struct {
			Kites struct {
				DisableWebSocketByDefault bool `json:"disableWebSocketByDefault"`
				Stack                     struct {
					Force    bool `json:"force"`
					NewKites bool `json:"newKites"`
				} `json:"stack"`
				Kontrol struct {
					Username string `json:"username"`
				} `json:"kontrol"`
				Os struct {
					Version string `json:"version"`
				} `json:"os"`
				Terminal struct {
					Version string `json:"version"`
				} `json:"terminal"`
				Klient struct {
					Version string `json:"version"`
				} `json:"klient"`
				Kloud struct {
					Version string `json:"version"`
				} `json:"kloud"`
			} `json:"kites"`
			Algolia struct {
				AppId       string `json:"appId"`
				ApiKey      string `json:"apiKey"`
				IndexSuffix string `json:"indexSuffix"`
			} `json:"algolia"`
			LogToExternal   bool   `json:"logToExternal"`
			SuppressLogs    bool   `json:"suppressLogs"`
			LogToInternal   bool   `json:"logToInternal"`
			AuthExchange    string `json:"authExchange"`
			Environment     string `json:"environment"`
			Version         string `json:"version"`
			ResourceName    string `json:"resourceName"`
			UserSitesDomain string `json:"userSitesDomain"`
			LogResourceName string `json:"logResourceName"`
			SocialApiUri    string `json:"socialApiUri"`
			ApiUri          string `json:"apiUri"`
			MainUri         string `json:"mainUri"`
			SourceMapsUri   string `json:"sourceMapsUri"`
			Broker          struct {
				Uri string `json:"uri"`
			} `json:"broker"`
			AppsUri            string `json:"appsUri"`
			UploadsUri         string `json:"uploadsUri"`
			UploadsUriForGroup string `json:"uploadsUriForGroup"`
			FileFetchTimeout   int    `json:"fileFetchTimeout"`
			UserIdleMs         int    `json:"userIdleMs"`
			Embedly            struct {
				ApiKey string `json:"apiKey"`
			} `json:"embedly"`
			Github struct {
				ClientId string `json:"clientId"`
			} `json:"github"`
			Newkontrol struct {
				Url string `json:"url"`
			} `json:"newkontrol"`
			SessionCookie struct {
				MaxAge int  `json:"maxAge"`
				Secure bool `json:"secure"`
			} `json:"sessionCookie"`
			Troubleshoot struct {
				IdleTime    int    `json:"idleTime"`
				ExternalUrl string `json:"externalUrl"`
			} `json:"troubleshoot"`
			Recaptcha string `json:"recaptcha"`
			Stripe    struct {
				Token string `json:"token"`
			} `json:"stripe"`
			ExternalProfiles struct {
				Google struct {
					Nicename string `json:"nicename"`
				} `json:"google"`
				Linkedin struct {
					Nicename string `json:"nicename"`
				} `json:"linkedin"`
				Twitter struct {
					Nicename string `json:"nicename"`
				} `json:"twitter"`
				Odesk struct {
					Nicename    string `json:"nicename"`
					UrlLocation string `json:"urlLocation"`
				} `json:"odesk"`
				Facebook struct {
					Nicename    string `json:"nicename"`
					UrlLocation string `json:"urlLocation"`
				} `json:"facebook"`
				Github struct {
					Nicename    string `json:"nicename"`
					UrlLocation string `json:"urlLocation"`
				} `json:"github"`
			} `json:"externalProfiles"`
			EntryPoint struct {
				Slug string `json:"slug"`
				Type string `json:"type"`
			} `json:"entryPoint"`
			Roles       []string      `json:"roles"`
			Permissions []interface{} `json:"permissions"`
		}
	}
	Mongo          string
	MongoKontrol   string
	MongoMinWrites int
	Mq             struct {
		Host     string
		Port     int
		Login    string
		Password string
		Vhost    string
		LogLevel string
	}
	Neo4j struct {
		Read    string
		Write   string
		Port    int
		Enabled bool
	}
	GoLogLevel        string
	Broker            Broker
	PremiumBroker     Broker
	BrokerKite        Broker
	PremiumBrokerKite Broker
	Loggr             struct {
		Push   bool
		Url    string
		ApiKey string
	}
	Librato struct {
		Push     bool
		Email    string
		Token    string
		Interval int
	}
	Opsview struct {
		Push bool
		Host string
	}
	ElasticSearch struct {
		Host  string
		Port  int
		Queue string
	}
	NewKites struct {
		UseTLS   bool
		CertFile string
		KeyFile  string
	}
	NewKontrol struct {
		Port           int
		UseTLS         bool
		CertFile       string
		KeyFile        string
		PublicKeyFile  string
		PrivateKeyFile string
	}
	ProxyKite struct {
		Domain   string
		CertFile string
		KeyFile  string
	}
	Etcd []struct {
		Host string
		Port int
	}
	Kontrold struct {
		Vhost    string
		Overview struct {
			ApiPort    int
			ApiHost    string
			Port       int
			KodingHost string
			SocialHost string
		}
		Api struct {
			Port int
			URL  string
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
		}
	}
	FollowFeed struct {
		Host          string
		Port          int
		ComponentUser string
		Password      string
		Vhost         string
	}
	Statsd struct {
		Use  bool
		Ip   string
		Port int
	}
	TopicModifier struct {
		CronSchedule string
	}
	Slack struct {
		Token   string
		Channel string
	}
	Graphite struct {
		Use  bool
		Host string
		Port int
	}
	LogLevel             map[string]string
	Redis                string
	SubscriptionEndpoint string
	Gowebserver          struct {
		Port int
	}
}

// TODO: THIS IS ADDED SO ALL GO PACKAGES CLEANLY EXIT EVEN WHEN
// RUN WITH RERUN

func init() {

	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				os.Exit(0)
			}
		}
	}()
}

func MustConfig(profile string) *Config {
	conf, err := readConfig("", profile)
	if err != nil {
		panic(err)
	}

	return conf
}

// MustEnv is like Env, but panics if the Config cannot be read successfully.
func MustEnv() *Config {
	conf, err := Env()
	if err != nil {
		panic(err)
	}

	return conf
}

// Env reads from the KONFIG_JSON environment variable and intitializes the
// Config struct
func Env() (*Config, error) {
	return readConfig("", "")
}

// TODO: Fix this shit below where dir and profile is not even used ...
func MustConfigDir(dir, profile string) *Config {
	conf, err := readConfig(dir, profile)
	if err != nil {
		panic(err)
	}

	return conf
}

func readConfig(configDir, profile string) (*Config, error) {
	jsonData := os.Getenv("KONFIG_JSON")
	if jsonData == "" {
		return nil, errors.New("KONFIG_JSON is not set")
	}

	conf := new(Config)
	err := json.Unmarshal([]byte(jsonData), &conf)
	if err != nil {
		return nil, fmt.Errorf("Configuration error, make sure KONFIG_JSON is set: %s\nConfiguration source output:\n%s\n",
			err.Error(), string(jsonData))
	}

	return conf, nil
}
<commit_msg>Config: add rerouting port into config<commit_after>package config

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"syscall"
)

type Broker struct {
	Name               string
	ServiceGenericName string
	IP                 string
	Port               int
	CertFile           string
	KeyFile            string
	AuthExchange       string
	AuthAllExchange    string
	WebProtocol        string
}

type Config struct {
	BuildNumber int
	Environment string
	Regions     struct {
		Vagrant string
		SJ      string
		AWS     string
		Premium string
	}
	ProjectRoot     string
	UserSitesDomain string
	ContainerSubnet string
	VmPool          string
	Version         string
	Client          struct {
		StaticFilesBaseUrl string
		RuntimeOptions     struct {
			Kites struct {
				DisableWebSocketByDefault bool `json:"disableWebSocketByDefault"`
				Stack                     struct {
					Force    bool `json:"force"`
					NewKites bool `json:"newKites"`
				} `json:"stack"`
				Kontrol struct {
					Username string `json:"username"`
				} `json:"kontrol"`
				Os struct {
					Version string `json:"version"`
				} `json:"os"`
				Terminal struct {
					Version string `json:"version"`
				} `json:"terminal"`
				Klient struct {
					Version string `json:"version"`
				} `json:"klient"`
				Kloud struct {
					Version string `json:"version"`
				} `json:"kloud"`
			} `json:"kites"`
			Algolia struct {
				AppId       string `json:"appId"`
				ApiKey      string `json:"apiKey"`
				IndexSuffix string `json:"indexSuffix"`
			} `json:"algolia"`
			LogToExternal   bool   `json:"logToExternal"`
			SuppressLogs    bool   `json:"suppressLogs"`
			LogToInternal   bool   `json:"logToInternal"`
			AuthExchange    string `json:"authExchange"`
			Environment     string `json:"environment"`
			Version         string `json:"version"`
			ResourceName    string `json:"resourceName"`
			UserSitesDomain string `json:"userSitesDomain"`
			LogResourceName string `json:"logResourceName"`
			SocialApiUri    string `json:"socialApiUri"`
			ApiUri          string `json:"apiUri"`
			MainUri         string `json:"mainUri"`
			SourceMapsUri   string `json:"sourceMapsUri"`
			Broker          struct {
				Uri string `json:"uri"`
			} `json:"broker"`
			AppsUri            string `json:"appsUri"`
			UploadsUri         string `json:"uploadsUri"`
			UploadsUriForGroup string `json:"uploadsUriForGroup"`
			FileFetchTimeout   int    `json:"fileFetchTimeout"`
			UserIdleMs         int    `json:"userIdleMs"`
			Embedly            struct {
				ApiKey string `json:"apiKey"`
			} `json:"embedly"`
			Github struct {
				ClientId string `json:"clientId"`
			} `json:"github"`
			Newkontrol struct {
				Url string `json:"url"`
			} `json:"newkontrol"`
			SessionCookie struct {
				MaxAge int  `json:"maxAge"`
				Secure bool `json:"secure"`
			} `json:"sessionCookie"`
			Troubleshoot struct {
				IdleTime    int    `json:"idleTime"`
				ExternalUrl string `json:"externalUrl"`
			} `json:"troubleshoot"`
			Recaptcha string `json:"recaptcha"`
			Stripe    struct {
				Token string `json:"token"`
			} `json:"stripe"`
			ExternalProfiles struct {
				Google struct {
					Nicename string `json:"nicename"`
				} `json:"google"`
				Linkedin struct {
					Nicename string `json:"nicename"`
				} `json:"linkedin"`
				Twitter struct {
					Nicename string `json:"nicename"`
				} `json:"twitter"`
				Odesk struct {
					Nicename    string `json:"nicename"`
					UrlLocation string `json:"urlLocation"`
				} `json:"odesk"`
				Facebook struct {
					Nicename    string `json:"nicename"`
					UrlLocation string `json:"urlLocation"`
				} `json:"facebook"`
				Github struct {
					Nicename    string `json:"nicename"`
					UrlLocation string `json:"urlLocation"`
				} `json:"github"`
			} `json:"externalProfiles"`
			EntryPoint struct {
				Slug string `json:"slug"`
				Type string `json:"type"`
			} `json:"entryPoint"`
			Roles       []string      `json:"roles"`
			Permissions []interface{} `json:"permissions"`
		}
	}
	Mongo          string
	MongoKontrol   string
	MongoMinWrites int
	Mq             struct {
		Host     string
		Port     int
		Login    string
		Password string
		Vhost    string
		LogLevel string
	}
	Neo4j struct {
		Read    string
		Write   string
		Port    int
		Enabled bool
	}
	GoLogLevel        string
	Broker            Broker
	PremiumBroker     Broker
	BrokerKite        Broker
	PremiumBrokerKite Broker
	Loggr             struct {
		Push   bool
		Url    string
		ApiKey string
	}
	Librato struct {
		Push     bool
		Email    string
		Token    string
		Interval int
	}
	Opsview struct {
		Push bool
		Host string
	}
	ElasticSearch struct {
		Host  string
		Port  int
		Queue string
	}
	NewKites struct {
		UseTLS   bool
		CertFile string
		KeyFile  string
	}
	NewKontrol struct {
		Port           int
		UseTLS         bool
		CertFile       string
		KeyFile        string
		PublicKeyFile  string
		PrivateKeyFile string
	}
	ProxyKite struct {
		Domain   string
		CertFile string
		KeyFile  string
	}
	Etcd []struct {
		Host string
		Port int
	}
	Kontrold struct {
		Vhost    string
		Overview struct {
			ApiPort    int
			ApiHost    string
			Port       int
			KodingHost string
			SocialHost string
		}
		Api struct {
			Port int
			URL  string
		}
		Proxy struct {
			Port    int
			PortSSL int
			FTPIP   string
		}
	}
	FollowFeed struct {
		Host          string
		Port          int
		ComponentUser string
		Password      string
		Vhost         string
	}
	Statsd struct {
		Use  bool
		Ip   string
		Port int
	}
	TopicModifier struct {
		CronSchedule string
	}
	Slack struct {
		Token   string
		Channel string
	}
	Graphite struct {
		Use  bool
		Host string
		Port int
	}
	LogLevel             map[string]string
	Redis                string
	SubscriptionEndpoint string
	Gowebserver          struct {
		Port int
	}
	Rerouting struct {
		Port int
	}
}

// TODO: THIS IS ADDED SO ALL GO PACKAGES CLEANLY EXIT EVEN WHEN
// RUN WITH RERUN

func init() {

	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals)
		for {
			signal := <-signals
			switch signal {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGSTOP:
				os.Exit(0)
			}
		}
	}()
}

func MustConfig(profile string) *Config {
	conf, err := readConfig("", profile)
	if err != nil {
		panic(err)
	}

	return conf
}

// MustEnv is like Env, but panics if the Config cannot be read successfully.
func MustEnv() *Config {
	conf, err := Env()
	if err != nil {
		panic(err)
	}

	return conf
}

// Env reads from the KONFIG_JSON environment variable and intitializes the
// Config struct
func Env() (*Config, error) {
	return readConfig("", "")
}

// TODO: Fix this shit below where dir and profile is not even used ...
func MustConfigDir(dir, profile string) *Config {
	conf, err := readConfig(dir, profile)
	if err != nil {
		panic(err)
	}

	return conf
}

func readConfig(configDir, profile string) (*Config, error) {
	jsonData := os.Getenv("KONFIG_JSON")
	if jsonData == "" {
		return nil, errors.New("KONFIG_JSON is not set")
	}

	conf := new(Config)
	err := json.Unmarshal([]byte(jsonData), &conf)
	if err != nil {
		return nil, fmt.Errorf("Configuration error, make sure KONFIG_JSON is set: %s\nConfiguration source output:\n%s\n",
			err.Error(), string(jsonData))
	}

	return conf, nil
}
<|endoftext|>


<commit_before>package vtctld

import (
	"fmt"
	"strconv"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/golang/protobuf/proto"
	"github.com/youtube/vitess/go/vt/discovery"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
	"github.com/youtube/vitess/go/vt/tabletserver/grpcqueryservice"
	"github.com/youtube/vitess/go/vt/tabletserver/queryservice/fakes"
	"github.com/youtube/vitess/go/vt/vttest/fakesqldb"
	"github.com/youtube/vitess/go/vt/wrangler"
	"github.com/youtube/vitess/go/vt/wrangler/testlib"
	"github.com/youtube/vitess/go/vt/zktopo/zktestserver"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// TestRealtimeStats tests the functionality of the realtimeStats object without using the HealthCheck object.
func TestRealtimeStats(t *testing.T) {
	tabletType := topodatapb.TabletType_REPLICA.String()
	ctx := context.Background()
	cells := []string{"cell1", "cell2"}
	ts := zktestserver.New(t, cells)

	// Populate topo.
	ts.CreateKeyspace(ctx, "ks1", &topodatapb.Keyspace{ShardingColumnName: "shardcol"})
	ts.Impl.CreateShard(ctx, "ks1", "-80", &topodatapb.Shard{
		Cells:    cells,
		KeyRange: &topodatapb.KeyRange{Start: nil, End: []byte{0x80}},
	})
	ts.Impl.CreateShard(ctx, "ks1", "80-", &topodatapb.Shard{
		Cells:    cells,
		KeyRange: &topodatapb.KeyRange{Start: []byte{0x80}, End: nil},
	})

	tablet1 := &topodatapb.Tablet{
		Alias:    &topodatapb.TabletAlias{Cell: "cell1", Uid: 100},
		Keyspace: "ks1",
		Shard:    "-80",
		Type:     topodatapb.TabletType_REPLICA,
		KeyRange: &topodatapb.KeyRange{Start: nil, End: []byte{0x80}},
		PortMap:  map[string]int32{"vt": 100},
	}
	ts.CreateTablet(ctx, tablet1)

	tablet2 := &topodatapb.Tablet{
		Alias:    &topodatapb.TabletAlias{Cell: "cell2", Uid: 200},
		Keyspace: "ks1",
		Shard:    "-80",
		Type:     topodatapb.TabletType_REPLICA,
		KeyRange: &topodatapb.KeyRange{Start: nil, End: []byte{0x80}},
		PortMap:  map[string]int32{"vt": 200},
	}
	ts.CreateTablet(ctx, tablet2)

	realtimeStats := newRealtimeStatsForTesting()

	target := &querypb.Target{
		Keyspace:   "ks1",
		Shard:      "-80",
		TabletType: topodatapb.TabletType_REPLICA,
	}

	stats1 := &querypb.RealtimeStats{
		HealthError:         "",
		SecondsBehindMaster: 2,
		BinlogPlayersCount:  0,
		CpuUsage:            12.1,
		Qps:                 5.6,
	}

	// Test 1: tablet1's stats should be updated with the one received by the HealthCheck object.
	want1 := &discovery.TabletStats{
		Tablet:  tablet1,
		Target:  target,
		Up:      true,
		Serving: true,
		TabletExternallyReparentedTimestamp: 5,
		Stats:     stats1,
		LastError: nil,
	}
	realtimeStats.tabletStats.StatsUpdate(want1)
	result := realtimeStats.tabletStatuses("cell1", "ks1", "-80", tabletType)
	checkResult(t, tablet1.Alias.Uid, result, want1)

	// Test 2: tablet1's stats should be updated with the new one received by the HealthCheck object.
	stats2 := &querypb.RealtimeStats{
		HealthError:         "Unhealthy tablet",
		SecondsBehindMaster: 15,
		BinlogPlayersCount:  0,
		CpuUsage:            56.5,
		Qps:                 7.9,
	}
	want2 := &discovery.TabletStats{
		Tablet:  tablet1,
		Target:  target,
		Up:      true,
		Serving: true,
		TabletExternallyReparentedTimestamp: 5,
		Stats:     stats2,
		LastError: nil,
	}
	realtimeStats.tabletStats.StatsUpdate(want2)
	result = realtimeStats.tabletStatuses("cell1", "ks1", "-80", tabletType)
	checkResult(t, tablet1.Alias.Uid, result, want2)

	// Test 3: tablet2's stats should be updated with the one received by the HealthCheck object,
	// leaving tablet1's stats unchanged.
	stats3 := &querypb.RealtimeStats{
		HealthError:         "Unhealthy tablet",
		SecondsBehindMaster: 15,
		BinlogPlayersCount:  0,
		CpuUsage:            56.5,
		Qps:                 7.9,
	}
	want3 := &discovery.TabletStats{
		Tablet:  tablet2,
		Target:  target,
		Up:      true,
		Serving: true,
		TabletExternallyReparentedTimestamp: 5,
		Stats:     stats3,
		LastError: nil,
	}
	realtimeStats.tabletStats.StatsUpdate(want3)
	result = realtimeStats.tabletStatuses("cell1", "ks1", "-80", tabletType)
	checkResult(t, tablet1.Alias.Uid, result, want2)
}

// TestRealtimeStatsWithQueryService uses fakeTablets and the fakeQueryService to
// copy the environment needed for the HealthCheck object.
func TestRealtimeStatsWithQueryService(t *testing.T) {
	// Set up testing keyspace with 2 tablets within 2 cells.
	keyspace := "ks"
	shard := "-80"
	tabletType := topodatapb.TabletType_REPLICA.String()
	ctx := context.Background()
	db := fakesqldb.Register()
	ts := zktestserver.New(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient())

	if err := ts.CreateKeyspace(context.Background(), keyspace, &topodatapb.Keyspace{
		ShardingColumnName: "keyspace_id",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
	}); err != nil {
		t.Fatalf("CreateKeyspace failed: %v", err)
	}

	t1 := testlib.NewFakeTablet(t, wr, "cell1", 0, topodatapb.TabletType_REPLICA, db,
		testlib.TabletKeyspaceShard(t, keyspace, shard))
	t2 := testlib.NewFakeTablet(t, wr, "cell2", 1, topodatapb.TabletType_REPLICA, db,
		testlib.TabletKeyspaceShard(t, keyspace, shard))
	for _, ft := range []*(testlib.FakeTablet){t1, t2} {
		ft.StartActionLoop(t, wr)
		defer ft.StopActionLoop(t)
	}

	target := querypb.Target{
		Keyspace:   keyspace,
		Shard:      shard,
		TabletType: topodatapb.TabletType_REPLICA,
	}
	fqs1 := fakes.NewStreamHealthQueryService(target)
	fqs2 := fakes.NewStreamHealthQueryService(target)
	grpcqueryservice.Register(t1.RPCServer, fqs1)
	grpcqueryservice.Register(t2.RPCServer, fqs2)

	fqs1.AddDefaultHealthResponse()

	realtimeStats, err := newRealtimeStats(ts)
	if err != nil {
		t.Fatalf("newRealtimeStats error: %v", err)
	}

	if err := discovery.WaitForTablets(ctx, realtimeStats.healthCheck, "cell1", keyspace, shard, []topodatapb.TabletType{topodatapb.TabletType_REPLICA}); err != nil {
		t.Fatalf("waitForTablets failed: %v", err)
	}

	// Test 1: tablet1's stats should be updated with the one received by the HealthCheck object.
	result := realtimeStats.tabletStatuses("cell1", keyspace, shard, tabletType)

	fmt.Println("PRINTING: ")
	for a, b := range result {
		fmt.Printf("a: %v b: %v ", a, b)
	}

	got := result["0"].Stats
	want := &querypb.RealtimeStats{
		SecondsBehindMaster: 1,
	}
	if !proto.Equal(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Test 2: tablet1's stats should be updated with the new one received by the HealthCheck object.
	fqs1.AddHealthResponseWithQPS(2.0)
	want2 := &querypb.RealtimeStats{
		SecondsBehindMaster: 1,
		Qps:                 2.0,
	}
	if err := checkStats(realtimeStats, "0", "cell1", keyspace, shard, tabletType, want2); err != nil {
		t.Errorf("%v", err)
	}

	// Test 3: tablet2's stats should be updated with the one received by the HealthCheck object,
	// leaving tablet1's stats unchanged.
	fqs2.AddHealthResponseWithQPS(3.0)
	want3 := &querypb.RealtimeStats{
		SecondsBehindMaster: 1,
		Qps:                 3.0,
	}
	if err := checkStats(realtimeStats, "1", "cell2", keyspace, shard, tabletType, want3); err != nil {
		t.Errorf("%v", err)
	}

	if err := checkStats(realtimeStats, "0", "cell1", keyspace, shard, tabletType, want2); err != nil {
		t.Errorf("%v", err)
	}
}

// checkResult checks to see that the TabletStats received are as expected.
func checkResult(t *testing.T, wantedUID uint32, resultMap map[string]*discovery.TabletStats, original *discovery.TabletStats) {
	result, ok := resultMap[strconv.FormatUint(uint64(wantedUID), 10)]
	if !ok {
		t.Errorf("No such tablet in tabletStatsCache")
	}
	if got, want := result.String(), original.String(); got != want {
		t.Errorf("got: %#v, want: %#v", got, want)
	}
}

// waitForTest ensures that the HealthCheck object received an update and passed
// that information to the correct tablet.
func checkStats(realtimeStats *realtimeStats, tabletUid, cell, keyspace, shard, tabletType string, want *querypb.RealtimeStats) error {
	deadline := time.Now().Add(time.Second * 5)
	for time.Now().Before(deadline) {
		result, ok := (realtimeStats.tabletStatuses(cell, keyspace, shard, tabletType))[tabletUid]
		if !ok {
			continue
		}
		got := result.Stats
		if proto.Equal(got, want) {
			return nil
		}
		time.Sleep(1 * time.Millisecond)
	}
	return fmt.Errorf("timeout error when getting tabletStatuses")
}

// newRealtimeStatsForTesting creates a new realtimeStats object without creating a HealthCheck object.
func newRealtimeStatsForTesting() *realtimeStats {
	tabletStatsCache := &tabletStatsCache{
		statuses: make(map[string]map[string]*discovery.TabletStats),
	}
	return &realtimeStats{
		tabletStats: tabletStatsCache,
	}
}
<commit_msg>Removing debugging statements<commit_after>package vtctld

import (
	"fmt"
	"strconv"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/golang/protobuf/proto"
	"github.com/youtube/vitess/go/vt/discovery"
	"github.com/youtube/vitess/go/vt/logutil"
	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
	"github.com/youtube/vitess/go/vt/tabletserver/grpcqueryservice"
	"github.com/youtube/vitess/go/vt/tabletserver/queryservice/fakes"
	"github.com/youtube/vitess/go/vt/vttest/fakesqldb"
	"github.com/youtube/vitess/go/vt/wrangler"
	"github.com/youtube/vitess/go/vt/wrangler/testlib"
	"github.com/youtube/vitess/go/vt/zktopo/zktestserver"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
)

// TestRealtimeStats tests the functionality of the realtimeStats object without using the HealthCheck object.
func TestRealtimeStats(t *testing.T) {
	tabletType := topodatapb.TabletType_REPLICA.String()
	ctx := context.Background()
	cells := []string{"cell1", "cell2"}
	ts := zktestserver.New(t, cells)

	// Populate topo.
	ts.CreateKeyspace(ctx, "ks1", &topodatapb.Keyspace{ShardingColumnName: "shardcol"})
	ts.Impl.CreateShard(ctx, "ks1", "-80", &topodatapb.Shard{
		Cells:    cells,
		KeyRange: &topodatapb.KeyRange{Start: nil, End: []byte{0x80}},
	})
	ts.Impl.CreateShard(ctx, "ks1", "80-", &topodatapb.Shard{
		Cells:    cells,
		KeyRange: &topodatapb.KeyRange{Start: []byte{0x80}, End: nil},
	})

	tablet1 := &topodatapb.Tablet{
		Alias:    &topodatapb.TabletAlias{Cell: "cell1", Uid: 100},
		Keyspace: "ks1",
		Shard:    "-80",
		Type:     topodatapb.TabletType_REPLICA,
		KeyRange: &topodatapb.KeyRange{Start: nil, End: []byte{0x80}},
		PortMap:  map[string]int32{"vt": 100},
	}
	ts.CreateTablet(ctx, tablet1)

	tablet2 := &topodatapb.Tablet{
		Alias:    &topodatapb.TabletAlias{Cell: "cell2", Uid: 200},
		Keyspace: "ks1",
		Shard:    "-80",
		Type:     topodatapb.TabletType_REPLICA,
		KeyRange: &topodatapb.KeyRange{Start: nil, End: []byte{0x80}},
		PortMap:  map[string]int32{"vt": 200},
	}
	ts.CreateTablet(ctx, tablet2)

	realtimeStats := newRealtimeStatsForTesting()

	target := &querypb.Target{
		Keyspace:   "ks1",
		Shard:      "-80",
		TabletType: topodatapb.TabletType_REPLICA,
	}

	stats1 := &querypb.RealtimeStats{
		HealthError:         "",
		SecondsBehindMaster: 2,
		BinlogPlayersCount:  0,
		CpuUsage:            12.1,
		Qps:                 5.6,
	}

	// Test 1: tablet1's stats should be updated with the one received by the HealthCheck object.
	want1 := &discovery.TabletStats{
		Tablet:  tablet1,
		Target:  target,
		Up:      true,
		Serving: true,
		TabletExternallyReparentedTimestamp: 5,
		Stats:     stats1,
		LastError: nil,
	}
	realtimeStats.tabletStats.StatsUpdate(want1)
	result := realtimeStats.tabletStatuses("cell1", "ks1", "-80", tabletType)
	checkResult(t, tablet1.Alias.Uid, result, want1)

	// Test 2: tablet1's stats should be updated with the new one received by the HealthCheck object.
	stats2 := &querypb.RealtimeStats{
		HealthError:         "Unhealthy tablet",
		SecondsBehindMaster: 15,
		BinlogPlayersCount:  0,
		CpuUsage:            56.5,
		Qps:                 7.9,
	}
	want2 := &discovery.TabletStats{
		Tablet:  tablet1,
		Target:  target,
		Up:      true,
		Serving: true,
		TabletExternallyReparentedTimestamp: 5,
		Stats:     stats2,
		LastError: nil,
	}
	realtimeStats.tabletStats.StatsUpdate(want2)
	result = realtimeStats.tabletStatuses("cell1", "ks1", "-80", tabletType)
	checkResult(t, tablet1.Alias.Uid, result, want2)

	// Test 3: tablet2's stats should be updated with the one received by the HealthCheck object,
	// leaving tablet1's stats unchanged.
	stats3 := &querypb.RealtimeStats{
		HealthError:         "Unhealthy tablet",
		SecondsBehindMaster: 15,
		BinlogPlayersCount:  0,
		CpuUsage:            56.5,
		Qps:                 7.9,
	}
	want3 := &discovery.TabletStats{
		Tablet:  tablet2,
		Target:  target,
		Up:      true,
		Serving: true,
		TabletExternallyReparentedTimestamp: 5,
		Stats:     stats3,
		LastError: nil,
	}
	realtimeStats.tabletStats.StatsUpdate(want3)
	result = realtimeStats.tabletStatuses("cell1", "ks1", "-80", tabletType)
	checkResult(t, tablet1.Alias.Uid, result, want2)
}

// TestRealtimeStatsWithQueryService uses fakeTablets and the fakeQueryService to
// copy the environment needed for the HealthCheck object.
func TestRealtimeStatsWithQueryService(t *testing.T) {
	// Set up testing keyspace with 2 tablets within 2 cells.
	keyspace := "ks"
	shard := "-80"
	tabletType := topodatapb.TabletType_REPLICA.String()
	ctx := context.Background()
	db := fakesqldb.Register()
	ts := zktestserver.New(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient())

	if err := ts.CreateKeyspace(context.Background(), keyspace, &topodatapb.Keyspace{
		ShardingColumnName: "keyspace_id",
		ShardingColumnType: topodatapb.KeyspaceIdType_UINT64,
	}); err != nil {
		t.Fatalf("CreateKeyspace failed: %v", err)
	}

	t1 := testlib.NewFakeTablet(t, wr, "cell1", 0, topodatapb.TabletType_REPLICA, db,
		testlib.TabletKeyspaceShard(t, keyspace, shard))
	t2 := testlib.NewFakeTablet(t, wr, "cell2", 1, topodatapb.TabletType_REPLICA, db,
		testlib.TabletKeyspaceShard(t, keyspace, shard))
	for _, ft := range []*(testlib.FakeTablet){t1, t2} {
		ft.StartActionLoop(t, wr)
		defer ft.StopActionLoop(t)
	}

	target := querypb.Target{
		Keyspace:   keyspace,
		Shard:      shard,
		TabletType: topodatapb.TabletType_REPLICA,
	}
	fqs1 := fakes.NewStreamHealthQueryService(target)
	fqs2 := fakes.NewStreamHealthQueryService(target)
	grpcqueryservice.Register(t1.RPCServer, fqs1)
	grpcqueryservice.Register(t2.RPCServer, fqs2)

	fqs1.AddDefaultHealthResponse()

	realtimeStats, err := newRealtimeStats(ts)
	if err != nil {
		t.Fatalf("newRealtimeStats error: %v", err)
	}

	if err := discovery.WaitForTablets(ctx, realtimeStats.healthCheck, "cell1", keyspace, shard, []topodatapb.TabletType{topodatapb.TabletType_REPLICA}); err != nil {
		t.Fatalf("waitForTablets failed: %v", err)
	}

	// Test 1: tablet1's stats should be updated with the one received by the HealthCheck object.
	result := realtimeStats.tabletStatuses("cell1", keyspace, shard, tabletType)
	got := result["0"].Stats
	want := &querypb.RealtimeStats{
		SecondsBehindMaster: 1,
	}
	if !proto.Equal(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}

	// Test 2: tablet1's stats should be updated with the new one received by the HealthCheck object.
	fqs1.AddHealthResponseWithQPS(2.0)
	want2 := &querypb.RealtimeStats{
		SecondsBehindMaster: 1,
		Qps:                 2.0,
	}
	if err := checkStats(realtimeStats, "0", "cell1", keyspace, shard, tabletType, want2); err != nil {
		t.Errorf("%v", err)
	}

	// Test 3: tablet2's stats should be updated with the one received by the HealthCheck object,
	// leaving tablet1's stats unchanged.
	fqs2.AddHealthResponseWithQPS(3.0)
	want3 := &querypb.RealtimeStats{
		SecondsBehindMaster: 1,
		Qps:                 3.0,
	}
	if err := checkStats(realtimeStats, "1", "cell2", keyspace, shard, tabletType, want3); err != nil {
		t.Errorf("%v", err)
	}

	if err := checkStats(realtimeStats, "0", "cell1", keyspace, shard, tabletType, want2); err != nil {
		t.Errorf("%v", err)
	}
}

// checkResult checks to see that the TabletStats received are as expected.
func checkResult(t *testing.T, wantedUID uint32, resultMap map[string]*discovery.TabletStats, original *discovery.TabletStats) {
	result, ok := resultMap[strconv.FormatUint(uint64(wantedUID), 10)]
	if !ok {
		t.Errorf("No such tablet in tabletStatsCache")
	}
	if got, want := result.String(), original.String(); got != want {
		t.Errorf("got: %#v, want: %#v", got, want)
	}
}

// checkStats ensures that the HealthCheck object received an update and passed
// that information to the correct tablet.
func checkStats(realtimeStats *realtimeStats, tabletUid, cell, keyspace, shard, tabletType string, want *querypb.RealtimeStats) error {
	deadline := time.Now().Add(time.Second * 5)
	for time.Now().Before(deadline) {
		result, ok := (realtimeStats.tabletStatuses(cell, keyspace, shard, tabletType))[tabletUid]
		if !ok {
			continue
		}
		got := result.Stats
		if proto.Equal(got, want) {
			return nil
		}
		time.Sleep(1 * time.Millisecond)
	}
	return fmt.Errorf("timeout error when getting tabletStatuses")
}

// newRealtimeStatsForTesting creates a new realtimeStats object without creating a HealthCheck object.
func newRealtimeStatsForTesting() *realtimeStats {
	tabletStatsCache := &tabletStatsCache{
		statuses: make(map[string]map[string]*discovery.TabletStats),
	}
	return &realtimeStats{
		tabletStats: tabletStatsCache,
	}
}
<|endoftext|>