<commit_before>// Copyright 2014 The mqrouter Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package dbrouter

import (
	"context"
	"fmt"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/log"
	"github.com/shawnfeng/sutil/slog/slog"
	stat "github.com/shawnfeng/sutil/stat"
	"github.com/shawnfeng/sutil/stime"
	"gopkg.in/mgo.v2"
)

const (
	spanLogKeyCluster = "cluster"
	spanLogKeyTable   = "table"
)

type Router struct {
	configer  Configer
	instances *InstanceManager
	report    *stat.StatReport
}

type dbConfigChange struct {
	dbInstanceChange map[string][]string
	dbGroups         []string
}

func NewRouter(data []byte) (*Router, error) {
	// TODO config type由哪里决定
	var dbChangeChan = make(chan dbConfigChange)
	configer, err := NewConfiger(CONFIG_TYPE_ETCD, data, dbChangeChan)
	if err != nil {
		return nil, err
	}

	factory := func(configer Configer) func(ctx context.Context, key, group string) (in Instancer, err error) {
		return func(ctx context.Context, key, group string) (in Instancer, err error) {
			return Factory(ctx, key, group, configer)
		}
	}(configer)

	return &Router{
		configer:  configer,
		instances: NewInstanceManager(factory, dbChangeChan, configer.GetGroups(context.TODO())),
		report:    stat.NewStat(),
	}, nil
}

func (m *Router) StatInfo() []*stat.QueryStat {
	return m.report.StatInfo()
}

func (m *Router) sqlPrepare(ctx context.Context, cluster, table string) (db *DB, err error) {
	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.sqlPrepare")
	defer span.Finish()

	instance := m.configer.GetInstance(ctx, cluster, table)
	in := m.instances.Get(ctx, generateKey(instance))
	if in == nil {
		err = fmt.Errorf("db instance not find: instance:%s", instance)
		return
	}

	dbsql, ok := in.(*Sql)
	if !ok {
		err = fmt.Errorf("db instance type error: instance:%s, dbtype:%s", instance, in.GetType())
		return
	}

	db = dbsql.getDB()
	return
}

func (m *Router) SqlExec(ctx context.Context, cluster string, query func(*DB, []interface{}) error, tables ...string) error {
	fun := "Router.SqlExec -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.SqlExec")
	defer span.Finish()

	st := stime.NewTimeStat()

	if len(tables) <= 0 {
		return fmt.Errorf("tables is empty")
	}
	table := tables[0]

	span.LogFields(
		log.String(spanLogKeyCluster, cluster),
		log.String(spanLogKeyTable, table))

	db, err := m.sqlPrepare(ctx, cluster, table)
	if err != nil {
		return err
	}

	defer func() {
		dur := st.Duration()
		m.report.IncQuery(cluster, table, st.Duration())
		slog.Tracef(ctx, "%s cls:%s table:%s dur:%d", fun, cluster, table, dur)
	}()

	var tmptables []interface{}
	for _, item := range tables {
		tmptables = append(tmptables, item)
	}

	return query(db, tmptables)
}

func (m *Router) ormPrepare(ctx context.Context, cluster, table string) (db *GormDB, err error) {
	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.ormPrepare")
	defer span.Finish()

	instance := m.configer.GetInstance(ctx, cluster, table)
	in := m.instances.Get(ctx, generateKey(instance))
	if in == nil {
		err = fmt.Errorf("db instance not find: instance:%s", instance)
		return
	}

	dbsql, ok := in.(*Sql)
	if !ok {
		err = fmt.Errorf("db instance type error: instance:%s, dbtype:%s", instance, in.GetType())
		return
	}

	db = dbsql.getGormDB()
	return
}

func (m *Router) OrmExec(ctx context.Context, cluster string, query func(*GormDB, []interface{}) error, tables ...string) error {
	fun := "Router.OrmExec -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.OrmExec")
	defer span.Finish()

	st := stime.NewTimeStat()

	if len(tables) <= 0 {
		return fmt.Errorf("tables is empty")
	}

	table := tables[0]
	span.LogFields(
		log.String(spanLogKeyCluster, cluster),
		log.String(spanLogKeyTable, table))

	db, err := m.ormPrepare(ctx, cluster, table)
	if err != nil {
		return err
	}

	defer func() {
		dur := st.Duration()
		m.report.IncQuery(cluster, table, st.Duration())
		slog.Tracef(ctx, "%s cls:%s table:%s dur:%d", fun, cluster, table, dur)
	}()

	var tmptables []interface{}
	for _, item := range tables {
		tmptables = append(tmptables, item)
	}

	return query(db, tmptables)
}

func (m *Router) MongoExecEventual(ctx context.Context, cluster, table string, query func(*mgo.Collection) error) error {
	return m.mongoExec(ctx, eventual, cluster, table, query)
}

func (m *Router) MongoExecMonotonic(ctx context.Context, cluster, table string, query func(*mgo.Collection) error) error {
	return m.mongoExec(ctx, monotonic, cluster, table, query)
}

func (m *Router) MongoExecStrong(ctx context.Context, cluster, table string, query func(*mgo.Collection) error) error {
	return m.mongoExec(ctx, strong, cluster, table, query)
}

func (m *Router) mongoPrepare(ctx context.Context, consistency mode, cluster, table string) (sess *mgo.Session, err error) {
	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.mongoPrepare")
	defer span.Finish()

	instance := m.configer.GetInstance(ctx, cluster, table)
	in := m.instances.Get(ctx, generateKey(instance))
	if in == nil {
		err = fmt.Errorf("db instance not find: cluster:%s table:%s", cluster, table)
		return
	}

	db, ok := in.(*dbMongo)
	if !ok {
		err = fmt.Errorf("db instance type error: cluster:%s table:%s type:%s", cluster, table, in.GetType())
		return
	}

	ss, err := db.getSession(consistency)
	if err != nil {
		return
	}

	if ss == nil {
		err = fmt.Errorf("db instance session empty: cluster:%s table:%s type:%s", cluster, table, in.GetType())
		return
	}

	sess = ss.Copy()
	return
}

func (m *Router) mongoExec(ctx context.Context, consistency mode, cluster, table string, query func(*mgo.Collection) error) error {
	fun := "Router.mongoExec -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.mongoExec")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyCluster, cluster),
		log.String(spanLogKeyTable, table))

	st := stime.NewTimeStat()

	sess, err := m.mongoPrepare(ctx, consistency, cluster, table)
	if err != nil {
		return err
	}
	defer sess.Close()
	coll := sess.DB("").C(table)

	defer func() {
		dur := st.Duration()
		m.report.IncQuery(cluster, table, st.Duration())
		slog.Tracef(ctx, "%s const:%d cls:%s table:%s dur:%d", fun, consistency, cluster, table, dur)
	}()

	return query(coll)
}
<commit_msg>mod tidb not found instance error<commit_after>// Copyright 2014 The mqrouter Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package dbrouter

import (
	"context"
	"fmt"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/log"
	"github.com/shawnfeng/sutil/slog/slog"
	"github.com/shawnfeng/sutil/stat"
	"github.com/shawnfeng/sutil/stime"
	"gopkg.in/mgo.v2"
)

const (
	spanLogKeyCluster = "cluster"
	spanLogKeyTable   = "table"
)

type Router struct {
	configer  Configer
	instances *InstanceManager
	report    *stat.StatReport
}

type dbConfigChange struct {
	dbInstanceChange map[string][]string
	dbGroups         []string
}

func NewRouter(data []byte) (*Router, error) {
	// TODO config type由哪里决定
	var dbChangeChan = make(chan dbConfigChange)
	configer, err := NewConfiger(CONFIG_TYPE_ETCD, data, dbChangeChan)
	if err != nil {
		return nil, err
	}

	factory := func(configer Configer) func(ctx context.Context, key, group string) (in Instancer, err error) {
		return func(ctx context.Context, key, group string) (in Instancer, err error) {
			return Factory(ctx, key, group, configer)
		}
	}(configer)

	return &Router{
		configer:  configer,
		instances: NewInstanceManager(factory, dbChangeChan, configer.GetGroups(context.TODO())),
		report:    stat.NewStat(),
	}, nil
}

func (m *Router) StatInfo() []*stat.QueryStat {
	return m.report.StatInfo()
}

func (m *Router) sqlPrepare(ctx context.Context, cluster, table string) (db *DB, err error) {
	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.sqlPrepare")
	defer span.Finish()

	instance := m.configer.GetInstance(ctx, cluster, table)
	in := m.instances.Get(ctx, generateKey(instance))
	if in == nil {
		err = fmt.Errorf("db instance not find: cluster:%s table:%s instance:%s", cluster, table, instance)
		return
	}

	dbsql, ok := in.(*Sql)
	if !ok {
		err = fmt.Errorf("db instance type error: cluster:%s table:%s instance:%s, dbtype:%s", cluster, table, instance, in.GetType())
		return
	}

	db = dbsql.getDB()
	return
}

func (m *Router) SqlExec(ctx context.Context, cluster string, query func(*DB, []interface{}) error, tables ...string) error {
	fun := "Router.SqlExec -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.SqlExec")
	defer span.Finish()

	st := stime.NewTimeStat()

	if len(tables) <= 0 {
		return fmt.Errorf("tables is empty")
	}
	table := tables[0]

	span.LogFields(
		log.String(spanLogKeyCluster, cluster),
		log.String(spanLogKeyTable, table))

	db, err := m.sqlPrepare(ctx, cluster, table)
	if err != nil {
		return err
	}

	defer func() {
		dur := st.Duration()
		m.report.IncQuery(cluster, table, st.Duration())
		slog.Tracef(ctx, "%s cls:%s table:%s dur:%d", fun, cluster, table, dur)
	}()

	var tmptables []interface{}
	for _, item := range tables {
		tmptables = append(tmptables, item)
	}

	return query(db, tmptables)
}

func (m *Router) ormPrepare(ctx context.Context, cluster, table string) (db *GormDB, err error) {
	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.ormPrepare")
	defer span.Finish()

	instance := m.configer.GetInstance(ctx, cluster, table)
	in := m.instances.Get(ctx, generateKey(instance))
	if in == nil {
		err = fmt.Errorf("db instance not find: instance:%s", instance)
		return
	}

	dbsql, ok := in.(*Sql)
	if !ok {
		err = fmt.Errorf("db instance type error: instance:%s, dbtype:%s", instance, in.GetType())
		return
	}

	db = dbsql.getGormDB()
	return
}

func (m *Router) OrmExec(ctx context.Context, cluster string, query func(*GormDB, []interface{}) error, tables ...string) error {
	fun := "Router.OrmExec -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.OrmExec")
	defer span.Finish()

	st := stime.NewTimeStat()

	if len(tables) <= 0 {
		return fmt.Errorf("tables is empty")
	}

	table := tables[0]
	span.LogFields(
		log.String(spanLogKeyCluster, cluster),
		log.String(spanLogKeyTable, table))

	db, err := m.ormPrepare(ctx, cluster, table)
	if err != nil {
		return err
	}

	defer func() {
		dur := st.Duration()
		m.report.IncQuery(cluster, table, st.Duration())
		slog.Tracef(ctx, "%s cls:%s table:%s dur:%d", fun, cluster, table, dur)
	}()

	var tmptables []interface{}
	for _, item := range tables {
		tmptables = append(tmptables, item)
	}

	return query(db, tmptables)
}

func (m *Router) MongoExecEventual(ctx context.Context, cluster, table string, query func(*mgo.Collection) error) error {
	return m.mongoExec(ctx, eventual, cluster, table, query)
}

func (m *Router) MongoExecMonotonic(ctx context.Context, cluster, table string, query func(*mgo.Collection) error) error {
	return m.mongoExec(ctx, monotonic, cluster, table, query)
}

func (m *Router) MongoExecStrong(ctx context.Context, cluster, table string, query func(*mgo.Collection) error) error {
	return m.mongoExec(ctx, strong, cluster, table, query)
}

func (m *Router) mongoPrepare(ctx context.Context, consistency mode, cluster, table string) (sess *mgo.Session, err error) {
	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.mongoPrepare")
	defer span.Finish()

	instance := m.configer.GetInstance(ctx, cluster, table)
	in := m.instances.Get(ctx, generateKey(instance))
	if in == nil {
		err = fmt.Errorf("db instance not find: cluster:%s table:%s", cluster, table)
		return
	}

	db, ok := in.(*dbMongo)
	if !ok {
		err = fmt.Errorf("db instance type error: cluster:%s table:%s type:%s", cluster, table, in.GetType())
		return
	}

	ss, err := db.getSession(consistency)
	if err != nil {
		return
	}

	if ss == nil {
		err = fmt.Errorf("db instance session empty: cluster:%s table:%s type:%s", cluster, table, in.GetType())
		return
	}

	sess = ss.Copy()
	return
}

func (m *Router) mongoExec(ctx context.Context, consistency mode, cluster, table string, query func(*mgo.Collection) error) error {
	fun := "Router.mongoExec -->"

	span, ctx := opentracing.StartSpanFromContext(ctx, "dbrouter.mongoExec")
	defer span.Finish()
	span.LogFields(
		log.String(spanLogKeyCluster, cluster),
		log.String(spanLogKeyTable, table))

	st := stime.NewTimeStat()

	sess, err := m.mongoPrepare(ctx, consistency, cluster, table)
	if err != nil {
		return err
	}
	defer sess.Close()
	coll := sess.DB("").C(table)

	defer func() {
		dur := st.Duration()
		m.report.IncQuery(cluster, table, st.Duration())
		slog.Tracef(ctx, "%s const:%d cls:%s table:%s dur:%d", fun, consistency, cluster, table, dur)
	}()

	return query(coll)
}
<|endoftext|>


<commit_before>package connector

import (
	"sync"

	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server/kvstore"
	"github.com/smancke/guble/server/router"
)

type Manager interface {
	Load() error
	List() []Subscriber
	Filter(map[string]string) []Subscriber
	Find(string) Subscriber
	Exists(string) bool
	Create(protocol.Path, router.RouteParams) (Subscriber, error)
	Add(Subscriber) error
	Update(Subscriber) error
	Remove(Subscriber) error
}

type manager struct {
	sync.RWMutex
	schema      string
	kvstore     kvstore.KVStore
	subscribers map[string]Subscriber
}

func NewManager(schema string, kvstore kvstore.KVStore) Manager {
	return &manager{
		schema:      schema,
		kvstore:     kvstore,
		subscribers: make(map[string]Subscriber, 0),
	}
}

func (m *manager) Load() error {
	// try to load s from kvstore
	entries := m.kvstore.Iterate(m.schema, "")
	for e := range entries {
		subscriber, err := NewSubscriberFromJSON([]byte(e[1]))
		if err != nil {
			return err
		}
		m.subscribers[subscriber.Key()] = subscriber
	}
	return nil
}

func (m *manager) Find(key string) Subscriber {
	m.RLock()
	defer m.RUnlock()

	if s, exists := m.subscribers[key]; exists {
		return s
	}
	return nil
}

func (m *manager) Create(topic protocol.Path, params router.RouteParams) (Subscriber, error) {
	key := GenerateKey(string(topic), params)
	//TODO MARIAN  remove this logs   when 503 is done.
	logger.WithField("key", key).Debug("Create generated key")
	if m.Exists(key) {
		logger.WithField("key", key).Debug("Create key exists already")
		return nil, ErrSubscriberExists
	}

	s := NewSubscriber(topic, params, 0)

	logger.WithField("subscriber", s).Debug("Create  newSubscriber created")
	err := m.Add(s)
	if err != nil {
		logger.WithField("error", err.Error()).Debug("Create Manager Add failed")
		return nil, err
	}
	logger.Debug("Create  finished")
	return s, nil
}

func (m *manager) List() []Subscriber {
	m.RLock()
	defer m.RUnlock()

	l := make([]Subscriber, 0, len(m.subscribers))
	for _, s := range m.subscribers {
		l = append(l, s)
	}
	return l
}

func (m *manager) Filter(filters map[string]string) (subscribers []Subscriber) {
	m.RLock()
	defer m.RUnlock()

	for _, s := range m.subscribers {
		if s.Filter(filters) {
			subscribers = append(subscribers, s)
		}
	}
	return
}

func (m *manager) Add(s Subscriber) error {
	logger.WithField("subscriber", s).WithField("lock", m.RWMutex).Debug("Add subscriber before locking")
	err := m.checkSubscribersExist(s)
	if err != nil {
		return err
	}

	if err := m.updateStore(s); err != nil {
		return err
	}

	m.putSubscriberInMap(s)
	logger.WithField("subscriber", s).Debug("Add subscriber after updating store")
	return nil
}

func (m *manager) Update(s Subscriber) error {
	err := m.checkSubscriberDoesNotExists(s)
	if err != nil {
		return err
	}

	m.putSubscriberInMap(s)
	return m.updateStore(s)
}

func (m *manager) checkSubscriberDoesNotExists(s Subscriber) error {
	m.RLock()
	defer m.RUnlock()

	if _, found := m.subscribers[s.Key()]; !found {
		return ErrSubscriberDoesNotExist
	}

	return nil
}

func (m *manager) checkSubscribersExist(s Subscriber) error {
	m.RLock()
	defer m.RUnlock()

	if _, found := m.subscribers[s.Key()]; found {
		return ErrSubscriberExists
	}

	return nil
}

func (m *manager) putSubscriberInMap(s Subscriber) {
	m.Lock()
	defer m.Unlock()
	m.subscribers[s.Key()] = s
}

func (m *manager) deleteSubscriberFromMap(s Subscriber) {
	m.Lock()
	defer m.Unlock()
	delete(m.subscribers, s.Key())
}

func (m *manager) Exists(key string) bool {
	m.RLock()
	defer m.RUnlock()

	_, found := m.subscribers[key]
	return found
}

func (m *manager) Remove(s Subscriber) error {
	m.cancelSubscribers(s)

	err := m.checkSubscriberDoesNotExists(s)
	if err != nil {
		return err
	}

	m.deleteSubscriberFromMap(s)

	return m.removeStore(s)
}

func (m *manager) cancelSubscribers(s Subscriber) {
	m.Lock()
	defer m.Unlock()

	s.Cancel()
}

func (m *manager) updateStore(s Subscriber) error {
	data, err := s.Encode()
	if err != nil {
		return err
	}
	//TODO MARIAN also remove this logs.
	logger.WithField("subscriber", s).Debug("UpdateStore")
	return m.kvstore.Put(m.schema, s.Key(), data)
}

func (m *manager) removeStore(s Subscriber) error {
	//TODO MARIAN also remove this logs.
	logger.WithField("subscriber", s).Debug("RemoveStore")
	return m.kvstore.Delete(m.schema, s.Key())
}
<commit_msg>remove some methods<commit_after>package connector

import (
	"sync"

	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server/kvstore"
	"github.com/smancke/guble/server/router"
)

type Manager interface {
	Load() error
	List() []Subscriber
	Filter(map[string]string) []Subscriber
	Find(string) Subscriber
	Exists(string) bool
	Create(protocol.Path, router.RouteParams) (Subscriber, error)
	Add(Subscriber) error
	Update(Subscriber) error
	Remove(Subscriber) error
}

type manager struct {
	sync.RWMutex
	schema      string
	kvstore     kvstore.KVStore
	subscribers map[string]Subscriber
}

func NewManager(schema string, kvstore kvstore.KVStore) Manager {
	return &manager{
		schema:      schema,
		kvstore:     kvstore,
		subscribers: make(map[string]Subscriber, 0),
	}
}

func (m *manager) Load() error {
	// try to load s from kvstore
	entries := m.kvstore.Iterate(m.schema, "")
	for e := range entries {
		subscriber, err := NewSubscriberFromJSON([]byte(e[1]))
		if err != nil {
			return err
		}
		m.subscribers[subscriber.Key()] = subscriber
	}
	return nil
}

func (m *manager) Find(key string) Subscriber {
	m.RLock()
	defer m.RUnlock()

	if s, exists := m.subscribers[key]; exists {
		return s
	}
	return nil
}

func (m *manager) Create(topic protocol.Path, params router.RouteParams) (Subscriber, error) {
	key := GenerateKey(string(topic), params)
	//TODO MARIAN  remove this logs   when 503 is done.
	logger.WithField("key", key).Debug("Create generated key")
	if m.Exists(key) {
		logger.WithField("key", key).Debug("Create key exists already")
		return nil, ErrSubscriberExists
	}

	s := NewSubscriber(topic, params, 0)

	logger.WithField("subscriber", s).Debug("Create  newSubscriber created")
	err := m.Add(s)
	if err != nil {
		logger.WithField("error", err.Error()).Debug("Create Manager Add failed")
		return nil, err
	}
	logger.Debug("Create  finished")
	return s, nil
}

func (m *manager) List() []Subscriber {
	m.RLock()
	defer m.RUnlock()

	l := make([]Subscriber, 0, len(m.subscribers))
	for _, s := range m.subscribers {
		l = append(l, s)
	}
	return l
}

func (m *manager) Filter(filters map[string]string) (subscribers []Subscriber) {
	m.RLock()
	defer m.RUnlock()

	for _, s := range m.subscribers {
		if s.Filter(filters) {
			subscribers = append(subscribers, s)
		}
	}
	return
}

func (m *manager) Add(s Subscriber) error {
	logger.WithField("subscriber", s).WithField("lock", m.RWMutex).Debug("Add subscriber started")

	if m.Exists(s.Key()) {
		return ErrSubscriberExists
	}

	if err := m.updateStore(s); err != nil {
		return err
	}

	m.putSubscriber(s)
	logger.WithField("subscriber", s).Debug("Add subscriber finished")
	return nil
}

func (m *manager) Update(s Subscriber) error {
	logger.WithField("subscriber", s).Debug("Update subscriber started")
	if !m.Exists(s.Key()) {
		return ErrSubscriberDoesNotExist
	}

	err := m.updateStore(s)
	if err != nil {
		return err
	}

	m.putSubscriber(s)
	logger.WithField("subscriber", s).Debug("Update subscriber finished")
	return nil
}

func (m *manager) putSubscriber(s Subscriber) {
	m.Lock()
	defer m.Unlock()
	m.subscribers[s.Key()] = s
}

func (m *manager) deleteSubscriber(s Subscriber) {
	m.Lock()
	defer m.Unlock()
	delete(m.subscribers, s.Key())
}

func (m *manager) Exists(key string) bool {
	m.RLock()
	defer m.RUnlock()

	_, found := m.subscribers[key]
	return found
}

func (m *manager) Remove(s Subscriber) error {
	logger.WithField("subscriber", s).Debug("Remove subscriber started")
	m.cancelSubscriber(s)

	if !m.Exists(s.Key()) {
		return ErrSubscriberDoesNotExist
	}

	err := m.removeStore(s)
	if err != nil {
		return err
	}

	m.deleteSubscriber(s)
	logger.WithField("subscriber", s).Debug("Remove subscriber finished")
	return nil
}

func (m *manager) cancelSubscriber(s Subscriber) {
	m.Lock()
	defer m.Unlock()

	s.Cancel()
}

func (m *manager) updateStore(s Subscriber) error {
	data, err := s.Encode()
	if err != nil {
		return err
	}
	//TODO MARIAN also remove this logs.
	logger.WithField("subscriber", s).Debug("UpdateStore")
	return m.kvstore.Put(m.schema, s.Key(), data)
}

func (m *manager) removeStore(s Subscriber) error {
	//TODO MARIAN also remove this logs.
	logger.WithField("subscriber", s).Debug("RemoveStore")
	return m.kvstore.Delete(m.schema, s.Key())
}
<|endoftext|>


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

package main

import (
	"fmt"
	"log"
	"math"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	"github.com/hajimehoshi/ebiten/exp/audio"
)

const (
	screenWidth  = 320
	screenHeight = 240
)

var frames = 0

const (
	freqA  = 440.0
	freqAS = 466.2
	freqB  = 493.9
	freqC  = 523.3
	freqCS = 554.4
	freqD  = 587.3
	freqDS = 622.3
	freqE  = 659.3
	freqF  = 698.5
	freqFS = 740.0
	freqG  = 784.0
	freqGS = 830.6
)

const score = `CCGGAAGR FFEEDDCR GGFFEEDR GGFFEEDR CCGGAAGR FFEEDDCR`

var scoreIndex = 0

func square(out []int16, volume float64, freq float64, sequence float64) {
	if freq == 0 {
		for i := 0; i < len(out); i++ {
			out[i] = 0
		}
		return
	}
	length := int(float64(audio.SampleRate) / freq)
	if length == 0 {
		panic("invalid freq")
	}
	for i := 0; i < len(out); i++ {
		a := int16(volume * math.MaxInt16)
		if i%length < int(float64(length)*sequence) {
			a = -a
		}
		out[i] = a
	}
}

func toBytes(l, r []int16) []byte {
	if len(l) != len(r) {
		panic("len(l) must equal to len(r)")
	}
	b := make([]byte, len(l)*4)
	for i, _ := range l {
		b[4*i] = byte(l[i])
		b[4*i+1] = byte(l[i] >> 8)
		b[4*i+2] = byte(r[i])
		b[4*i+3] = byte(r[i] >> 8)
	}
	return b
}

func addNote() {
	size := audio.SampleRate / 60
	notes := []float64{freqC, freqD, freqE, freqF, freqG, freqA * 2, freqB * 2}

	defer func() {
		scoreIndex++
		scoreIndex %= len(score)
	}()
	l := make([]int16, size*30)
	r := make([]int16, size*30)
	note := score[scoreIndex]
	for note == ' ' {
		scoreIndex++
		scoreIndex %= len(score)
		note = score[scoreIndex]
	}
	freq := 0.0
	switch {
	case note == 'R':
		freq = 0
	case note <= 'B':
		freq = notes[int(note)+len(notes)-int('C')]
	default:
		freq = notes[note-'C']
	}
	vol := 1.0 / 16.0
	square(l, vol, freq, 0.25)
	square(r, vol, freq, 0.25)
	audio.Queue(0, toBytes(l, r))
}

func update(screen *ebiten.Image) error {
	defer func() {
		frames++
	}()
	if frames%30 == 0 {
		addNote()
	}
	ebitenutil.DebugPrint(screen, fmt.Sprintf("FPS: %0.2f", ebiten.CurrentFPS()))
	return nil
}

func main() {
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Audio (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>audio: Use a new channel: We don't guarantee the exact condition when IsPlaying returns true/false<commit_after>// Copyright 2015 Hajime Hoshi
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"log"
	"math"

	"github.com/hajimehoshi/ebiten"
	"github.com/hajimehoshi/ebiten/ebitenutil"
	"github.com/hajimehoshi/ebiten/exp/audio"
)

const (
	screenWidth  = 320
	screenHeight = 240
)

var frames = 0

const (
	freqA  = 440.0
	freqAS = 466.2
	freqB  = 493.9
	freqC  = 523.3
	freqCS = 554.4
	freqD  = 587.3
	freqDS = 622.3
	freqE  = 659.3
	freqF  = 698.5
	freqFS = 740.0
	freqG  = 784.0
	freqGS = 830.6
)

const score = `CCGGAAGR FFEEDDCR GGFFEEDR GGFFEEDR CCGGAAGR FFEEDDCR`

var scoreIndex = 0

func square(out []int16, volume float64, freq float64, sequence float64) {
	if freq == 0 {
		for i := 0; i < len(out); i++ {
			out[i] = 0
		}
		return
	}
	length := int(float64(audio.SampleRate) / freq)
	if length == 0 {
		panic("invalid freq")
	}
	for i := 0; i < len(out); i++ {
		a := int16(volume * math.MaxInt16)
		if i%length < int(float64(length)*sequence) {
			a = -a
		}
		out[i] = a
	}
}

func toBytes(l, r []int16) []byte {
	if len(l) != len(r) {
		panic("len(l) must equal to len(r)")
	}
	b := make([]byte, len(l)*4)
	for i, _ := range l {
		b[4*i] = byte(l[i])
		b[4*i+1] = byte(l[i] >> 8)
		b[4*i+2] = byte(r[i])
		b[4*i+3] = byte(r[i] >> 8)
	}
	return b
}

func addNote() {
	size := audio.SampleRate / 60
	notes := []float64{freqC, freqD, freqE, freqF, freqG, freqA * 2, freqB * 2}

	defer func() {
		scoreIndex++
		scoreIndex %= len(score)
	}()
	l := make([]int16, size*30)
	r := make([]int16, size*30)
	note := score[scoreIndex]
	for note == ' ' {
		scoreIndex++
		scoreIndex %= len(score)
		note = score[scoreIndex]
	}
	freq := 0.0
	switch {
	case note == 'R':
		freq = 0
	case note <= 'B':
		freq = notes[int(note)+len(notes)-int('C')]
	default:
		freq = notes[note-'C']
	}
	vol := 1.0 / 16.0
	square(l, vol, freq, 0.25)
	square(r, vol, freq, 0.25)
	audio.Queue(-1, toBytes(l, r))
}

func update(screen *ebiten.Image) error {
	defer func() {
		frames++
	}()
	if frames%30 == 0 {
		addNote()
	}
	ebitenutil.DebugPrint(screen, fmt.Sprintf("FPS: %0.2f", ebiten.CurrentFPS()))
	return nil
}

func main() {
	if err := ebiten.Run(update, screenWidth, screenHeight, 2, "Audio (Ebiten Demo)"); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package derpiapi

import (
	"encoding/json"
	//	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

//Image contains data we got from API that we are using to filter and fetch said image next
type Image struct {
	Imgid          int    `json:"id_number"`
	URL            string `json:"image"`
	Filename       string
	Score          int    `json:"score"`
	OriginalFormat string `json:"original_format"`
	// Hashval     string `json:"sha512_hash"`
}

//Search returns to us array of searched images...
type Search struct {
	Images []Image `json:"search"`
}

type ImageCh chan Image

//infotochannel gets unmarchalled JSON info and plugs it into channel so it would be processed in other places
func infotochannel(dat Image, imgchan ImageCh) {
	dat.Filename = strconv.Itoa(dat.Imgid) + "." + dat.OriginalFormat
	dat.URL = "https:" + dat.URL
	if dat.OriginalFormat == "svg" {
		i := strings.LastIndex(dat.URL, ".")
		if i != -1 {
			dat.URL = dat.URL[:i] + ".svg" //Was afraid to extract things I needed from the date field, so extracting them from URL.
		}
	}
	imgchan <- dat
}

//ParseImg gets image ID and, fetches information about this image from Derpibooru and puts it into the channel.
func (imgchan ImageCh) ParseImg(imgid int, KEY string, elog *log.Logger) {

	source := "https://derpiboo.ru/images/" + strconv.Itoa(imgid) + ".json"
	if KEY != "" {
		source = source + "?key=" + KEY
	}

	log.Println("Getting image info at:", source)

	resp, err := http.Get(source) //Getting our nice http response. Needs checking for 404 and other responses that are... less expected
	if err != nil {
		elog.Println(err)
		return
	}

	defer resp.Body.Close() //and not forgetting to close it when it's done

	var dat Image

	body, err := ioutil.ReadAll(resp.Body) //stolen from official documentation
	if err != nil {
		elog.Println(err)
		return
	}

	if err := json.Unmarshal(body, &dat); //transforming json into native map

	err != nil {
		elog.Println(err)
		return
	}

	infotochannel(dat, imgchan)

	close(imgchan) //closing channel, we are done here

	return
}

//DlImg downloads image on disk, given image data
func (imgchan ImageCh) DlImg(done chan bool, elog *log.Logger, IMGDIR string) {

	log.Println("Worker started; reading channel") //nice notification that we are not forgotten

	for {

		imgdata, more := <-imgchan

		if !more { //checking that there is an image in channel
			done <- true //well, there is no images in channel, it means we got them all, so synchronization is kicking in and ending the process
			break        //Just in case, so it would not stupidly die when program finishes - it will die smartly
		}

		if imgdata.Filename == "" {
			elog.Println("Empty filename. Oops?") //something somewhere had gone wrong, not a cause to die, going to the next image
			continue
		}

		log.Println("Saving as", imgdata.Filename)

		func() { // To not hold all the files open when there is no need. All pointers to files are in the scope of this function.

			output, err := os.Create(IMGDIR + string(os.PathSeparator) + imgdata.Filename) //And now, THE FILE!
			if err != err {
				elog.Println("Error when creating file for image" + strconv.Itoa(imgdata.Imgid))
				elog.Println(err) //Either we got no permisson or no space, end of line
				return
			}
			defer output.Close() //Not forgetting to deal with it after completing download

			start := time.Now()

			response, err := http.Get(imgdata.URL)
			if err != nil {
				elog.Println("Error when getting image", imgdata.Imgid)
				elog.Println(err)
				return
			}
			defer response.Body.Close() //Same, we shall not listen to the void when we finished getting image

			size, err := io.Copy(output, response.Body) //	Writing things we got from Derpibooru into the file and into hasher
			if err != nil {
				elog.Println("Unable to write image on disk, id ", imgdata.Imgid)
				elog.Println(err)
				return
			}
			timed := time.Since(start).Seconds()
			log.Printf("Downloaded %d bytes in %.2fs, speed %s/s\n", size, timed, fmtbytes(float64(size)/timed))
		}()

		//fmt.Println("\n", hex.EncodeToString(hash.Sum(nil)), "\n", imgdata.hash )

	}
}

func (imgchan ImageCh) ParseTag(tag string, KEY string, STARTPAGE int, STOPPAGE int, elog *log.Logger) {

	source := "https://derpiboo.ru/search.json?" //yay hardwiring url strings!

	if KEY != "" {
		source = source + "key=" + KEY + "&"
	}

	log.Println("Searching as", source+"q="+tag)
	var working = true
	i := STARTPAGE
	for working {
		func() { //I suspect that all those returns could be dealt with in some way. But lazy.
			log.Println("Searching page", i)
			resp, err := http.Get(source + "q=" + tag + "&page=" + strconv.Itoa(i)) //Getting our nice http response. Needs checking for 404 and other responses that are... less expected
			defer resp.Body.Close()                                                 //and not forgetting to close it when it's done. And before we panic and die horribly.
			if err != nil {
				elog.Println("Error while getting search page", i)
				elog.Println(err)
				return
			}

			var dats Search //Because we got array incoming instead of single object, we using an slive of maps!

			//fmt.Println(resp)

			body, err := ioutil.ReadAll(resp.Body) //stolen from official documentation
			if err != nil {
				elog.Println("Error while reading search page", i)
				elog.Println(err)
				return
			}

			//fmt.Println(body)

			err = json.Unmarshal(body, &dats) //transforming json into native slice of maps

			if err != nil {
				elog.Println("Error while parsing search page", i)
				elog.Println(err)
				return

			}

			if len(dats.Images) == 0 {
				log.Println("Pages are over") //Does not mean that process is over.
				working = false
				return
			} //exit due to finishing all pages

			for _, dat := range dats.Images {
				infotochannel(dat, imgchan)
			}
			if STOPPAGE != 0 && i > STOPPAGE { //stop page is included, but if not set? Work to the end of tag
				working = false
				return
			}
			i++

		}()
	}

	close(imgchan)
}
<commit_msg>Added hash check to make sure stuff downloaded correctly<commit_after>package derpiapi

import (
	"encoding/json"
	//	"fmt"
	"crypto/sha512"
	"encoding/hex"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

//Image contains data we got from API that we are using to filter and fetch said image next
type Image struct {
	Imgid          int    `json:"id_number"`
	URL            string `json:"image"`
	Filename       string
	Score          int    `json:"score"`
	OriginalFormat string `json:"original_format"`
	Hashval        string `json:"sha512_hash"`
}

//Search returns to us array of searched images...
type Search struct {
	Images []Image `json:"search"`
}

type ImageCh chan Image

//infotochannel gets unmarchalled JSON info and plugs it into channel so it would be processed in other places
func infotochannel(dat Image, imgchan ImageCh) {
	dat.Filename = strconv.Itoa(dat.Imgid) + "." + dat.OriginalFormat
	dat.URL = "https:" + dat.URL
	if dat.OriginalFormat == "svg" {
		i := strings.LastIndex(dat.URL, ".")
		if i != -1 {
			dat.URL = dat.URL[:i] + ".svg" //Was afraid to extract things I needed from the date field, so extracting them from URL.
		}
	}
	imgchan <- dat
}

//ParseImg gets image ID and, fetches information about this image from Derpibooru and puts it into the channel.
func (imgchan ImageCh) ParseImg(imgid int, KEY string, elog *log.Logger) {

	source := "https://derpiboo.ru/images/" + strconv.Itoa(imgid) + ".json"
	if KEY != "" {
		source = source + "?key=" + KEY
	}

	log.Println("Getting image info at:", source)

	resp, err := http.Get(source) //Getting our nice http response. Needs checking for 404 and other responses that are... less expected
	if err != nil {
		elog.Println(err)
		return
	}

	defer resp.Body.Close() //and not forgetting to close it when it's done

	var dat Image

	body, err := ioutil.ReadAll(resp.Body) //stolen from official documentation
	if err != nil {
		elog.Println(err)
		return
	}

	if err := json.Unmarshal(body, &dat); //transforming json into native map

	err != nil {
		elog.Println(err)
		return
	}

	infotochannel(dat, imgchan)

	close(imgchan) //closing channel, we are done here

	return
}

//DlImg downloads image on disk, given image data
func (imgchan ImageCh) DlImg(done chan bool, elog *log.Logger, IMGDIR string) {

	log.Println("Worker started; reading channel") //nice notification that we are not forgotten

	hasher := sha512.New()

	for {

		imgdata, more := <-imgchan

		if !more { //checking that there is an image in channel
			done <- true //well, there is no images in channel, it means we got them all, so synchronization is kicking in and ending the process
			break        //Just in case, so it would not stupidly die when program finishes - it will die smartly
		}

		if imgdata.Filename == "" {
			elog.Println("Empty filename. Oops?") //something somewhere had gone wrong, not a cause to die, going to the next image
			continue
		}

		log.Println("Saving as", imgdata.Filename)

		func() { // To not hold all the files open when there is no need. All pointers to files are in the scope of this function.

			output, err := os.Create(IMGDIR + string(os.PathSeparator) + imgdata.Filename) //And now, THE FILE!
			if err != err {
				elog.Println("Error when creating file for image" + strconv.Itoa(imgdata.Imgid))
				elog.Println(err) //Either we got no permisson or no space, end of line
				return
			}
			defer output.Close() //Not forgetting to deal with it after completing download

			start := time.Now()

			response, err := http.Get(imgdata.URL)
			if err != nil {
				elog.Println("Error when getting image", imgdata.Imgid)
				elog.Println(err)
				return
			}
			defer response.Body.Close() //Same, we shall not listen to the void when we finished getting image

			if _, err := io.Copy(hasher, response.Body); err != nil {
				elog.Println(err)
				return
			}
			hash := hex.EncodeToString(hasher.Sum(nil))

			if hash != imgdata.Hashval {
			elog.Println("Hash mismatch, got ", hash, " instead of ", imgdata.Hashval)
			}

			size, err := io.Copy(output, response.Body) //	Writing things we got from Derpibooru into the file and into hasher
			if err != nil {
				elog.Println("Unable to write image on disk, id ", imgdata.Imgid)
				elog.Println(err)
				return
			}
			timed := time.Since(start).Seconds()
			log.Printf("Downloaded %d bytes in %.2fs, speed %s/s\n", size, timed, fmtbytes(float64(size)/timed))
		}()

		//fmt.Println("\n", hex.EncodeToString(hash.Sum(nil)), "\n", imgdata.hash )

	}
}

func (imgchan ImageCh) ParseTag(tag string, KEY string, STARTPAGE int, STOPPAGE int, elog *log.Logger) {

	source := "https://derpiboo.ru/search.json?" //yay hardwiring url strings!

	if KEY != "" {
		source = source + "key=" + KEY + "&"
	}

	log.Println("Searching as", source+"q="+tag)
	var working = true
	i := STARTPAGE
	for working {
		func() { //I suspect that all those returns could be dealt with in some way. But lazy.
			log.Println("Searching page", i)
			resp, err := http.Get(source + "q=" + tag + "&page=" + strconv.Itoa(i)) //Getting our nice http response. Needs checking for 404 and other responses that are... less expected
			defer resp.Body.Close()                                                 //and not forgetting to close it when it's done. And before we panic and die horribly.
			if err != nil {
				elog.Println("Error while getting search page", i)
				elog.Println(err)
				return
			}

			var dats Search //Because we got array incoming instead of single object, we using an slive of maps!

			//fmt.Println(resp)

			body, err := ioutil.ReadAll(resp.Body) //stolen from official documentation
			if err != nil {
				elog.Println("Error while reading search page", i)
				elog.Println(err)
				return
			}

			//fmt.Println(body)

			err = json.Unmarshal(body, &dats) //transforming json into native slice of maps

			if err != nil {
				elog.Println("Error while parsing search page", i)
				elog.Println(err)
				return

			}

			if len(dats.Images) == 0 {
				log.Println("Pages are over") //Does not mean that process is over.
				working = false
				return
			} //exit due to finishing all pages

			for _, dat := range dats.Images {
				infotochannel(dat, imgchan)
			}
			if STOPPAGE != 0 && i > STOPPAGE { //stop page is included, but if not set? Work to the end of tag
				working = false
				return
			}
			i++

		}()
	}

	close(imgchan)
}
<|endoftext|>


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

package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/HewlettPackard/lustre_exporter/sources"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
)

var (
	scrapeDurations = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Namespace: sources.Namespace,
			Subsystem: "exporter",
			Name:      "scrape_duration_seconds",
			Help:      "lustre_exporter: Duration of a scrape job.",
		},
		[]string{"source", "result"},
	)
)

//LustreSource is a list of all sources that the user would like to collect.
type LustreSource struct {
	sourceList map[string]sources.LustreSource
}

//Describe implements the prometheus.Describe interface
func (l LustreSource) Describe(ch chan<- *prometheus.Desc) {
	scrapeDurations.Describe(ch)
}

//Collect implements the prometheus.Collect interface
func (l LustreSource) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(l.sourceList))
	for name, c := range l.sourceList {
		go func(name string, s sources.LustreSource) {
			collectFromSource(name, c, ch)
			wg.Done()
		}(name, c)
	}
	wg.Wait()
	scrapeDurations.Collect(ch)
}

func collectFromSource(name string, s sources.LustreSource, ch chan<- prometheus.Metric) {
	result := "success"
	begin := time.Now()
	err := s.Update(ch)
	duration := time.Since(begin)
	if err != nil {
		log.Errorf("ERROR: %q source failed after %f seconds: %s", name, duration.Seconds(), err)
		result = "error"
	} else {
		log.Debugf("OK: %q source succeeded after %f seconds: %s", name, duration.Seconds(), err)
	}
	scrapeDurations.WithLabelValues(name, result).Observe(duration.Seconds())
}

func loadSources(list string) (map[string]sources.LustreSource, error) {
	sourceList := map[string]sources.LustreSource{}
	for _, name := range strings.Split(list, ",") {
		fn, ok := sources.Factories[name]
		if !ok {
			return nil, fmt.Errorf("source %q not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		sourceList[name] = c
	}
	return sourceList, nil
}

func init() {
	prometheus.MustRegister(version.NewCollector("lustre_exporter"))
}

func main() {
	var (
		showVersion   = flag.Bool("version", false, "Print version information.")
		listenAddress = flag.String("web.listen-address", ":9169", "Address to use to expose Lustre metrics.")
		metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path to use to expose Lustre metrics.")
	)
	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("lustre_exporter"))
		os.Exit(0)
	}

	log.Infoln("Starting lustre_exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	//expand to include more sources eventually (CLI, other?)
	enabledSources := "procfs"

	sourceList, err := loadSources(enabledSources)
	if err != nil {
		log.Fatalf("Couldn't load sources: %q", err)
	}

	log.Infof("Enabled sources:")
	for s := range sourceList {
		log.Infof(" - %s", s)
	}

	prometheus.MustRegister(LustreSource{sourceList: sourceList})
	handler := promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{ErrorLog: log.NewErrorLogger()})

	http.Handle(*metricsPath, prometheus.InstrumentHandler("prometheus", handler))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Lustre Exporter</title>

Lustre Exporter

<a href="` + *metricsPath + `">Metrics

			`))
	})

	log.Infoln("Listening on", *listenAddress)
	err = http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Fix range variable c captured by func literal<commit_after>// (C) Copyright 2017 Hewlett Packard Enterprise Development LP
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/HewlettPackard/lustre_exporter/sources"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/version"
)

var (
	scrapeDurations = prometheus.NewSummaryVec(
		prometheus.SummaryOpts{
			Namespace: sources.Namespace,
			Subsystem: "exporter",
			Name:      "scrape_duration_seconds",
			Help:      "lustre_exporter: Duration of a scrape job.",
		},
		[]string{"source", "result"},
	)
)

//LustreSource is a list of all sources that the user would like to collect.
type LustreSource struct {
	sourceList map[string]sources.LustreSource
}

//Describe implements the prometheus.Describe interface
func (l LustreSource) Describe(ch chan<- *prometheus.Desc) {
	scrapeDurations.Describe(ch)
}

//Collect implements the prometheus.Collect interface
func (l LustreSource) Collect(ch chan<- prometheus.Metric) {
	wg := sync.WaitGroup{}
	wg.Add(len(l.sourceList))
	for name, c := range l.sourceList {
		go func(name string, c sources.LustreSource) {
			collectFromSource(name, c, ch)
			wg.Done()
		}(name, c)
	}
	wg.Wait()
	scrapeDurations.Collect(ch)
}

func collectFromSource(name string, s sources.LustreSource, ch chan<- prometheus.Metric) {
	result := "success"
	begin := time.Now()
	err := s.Update(ch)
	duration := time.Since(begin)
	if err != nil {
		log.Errorf("ERROR: %q source failed after %f seconds: %s", name, duration.Seconds(), err)
		result = "error"
	} else {
		log.Debugf("OK: %q source succeeded after %f seconds: %s", name, duration.Seconds(), err)
	}
	scrapeDurations.WithLabelValues(name, result).Observe(duration.Seconds())
}

func loadSources(list string) (map[string]sources.LustreSource, error) {
	sourceList := map[string]sources.LustreSource{}
	for _, name := range strings.Split(list, ",") {
		fn, ok := sources.Factories[name]
		if !ok {
			return nil, fmt.Errorf("source %q not available", name)
		}
		c, err := fn()
		if err != nil {
			return nil, err
		}
		sourceList[name] = c
	}
	return sourceList, nil
}

func init() {
	prometheus.MustRegister(version.NewCollector("lustre_exporter"))
}

func main() {
	var (
		showVersion   = flag.Bool("version", false, "Print version information.")
		listenAddress = flag.String("web.listen-address", ":9169", "Address to use to expose Lustre metrics.")
		metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path to use to expose Lustre metrics.")
	)
	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("lustre_exporter"))
		os.Exit(0)
	}

	log.Infoln("Starting lustre_exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())

	//expand to include more sources eventually (CLI, other?)
	enabledSources := "procfs"

	sourceList, err := loadSources(enabledSources)
	if err != nil {
		log.Fatalf("Couldn't load sources: %q", err)
	}

	log.Infof("Enabled sources:")
	for s := range sourceList {
		log.Infof(" - %s", s)
	}

	prometheus.MustRegister(LustreSource{sourceList: sourceList})
	handler := promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{ErrorLog: log.NewErrorLogger()})

	http.Handle(*metricsPath, prometheus.InstrumentHandler("prometheus", handler))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Lustre Exporter</title>

Lustre Exporter

<a href="` + *metricsPath + `">Metrics

			`))
	})

	log.Infoln("Listening on", *listenAddress)
	err = http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


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

package commands

import (
	"errors"
	"fmt"

	"github.com/juju/cmd"

	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
)

// NewAddKeysCommand is used to add a new ssh key to a model.
func NewAddKeysCommand() cmd.Command {
	return modelcmd.Wrap(&addKeysCommand{})
}

var addKeysDoc = `
Add new authorized ssh keys to allow the holder of those keys to log on to Juju nodes or machines.
`

// addKeysCommand is used to add a new authorized ssh key for a user.
type addKeysCommand struct {
	SSHKeysBase
	user    string
	sshKeys []string
}

// Info implements Command.Info.
func (c *addKeysCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "add-ssh-key",
		Args:    "<ssh key> ...",
		Doc:     addKeysDoc,
		Purpose: "add new authorized ssh key to a Juju model",
		Aliases: []string{"add-ssh-keys"},
	}
}

// Init implements Command.Init.
func (c *addKeysCommand) Init(args []string) error {
	switch len(args) {
	case 0:
		return errors.New("no ssh key specified")
	default:
		c.sshKeys = args
	}
	return nil
}

// Run implements Command.Run.
func (c *addKeysCommand) Run(context *cmd.Context) error {
	client, err := c.NewKeyManagerClient()
	if err != nil {
		return err
	}
	defer client.Close()
	// TODO(alexisb) - currently keys are global which is not ideal.
	// keymanager needs to be updated to allow keys per user
	c.user = "admin"
	results, err := client.AddKeys(c.user, c.sshKeys...)
	if err != nil {
		return block.ProcessBlockedError(err, block.BlockChange)
	}
	for i, result := range results {
		if result.Error != nil {
			fmt.Fprintf(context.Stderr, "cannot add key %q: %v\n", c.sshKeys[i], result.Error)
		}
	}
	return nil
}
<commit_msg>Update add-ssh-key usage<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package commands

import (
	"errors"
	"fmt"

	"github.com/juju/cmd"

	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
)

var usageAddSSHKeySummary = `
Adds a public SSH key to a model.`[1:]

var usageAddSSHKeyDetails = `
Juju maintains a per-model cache of public SSH keys which it copies to
each unit (including units already deployed). By default this includes the
key of the user who created the model (assuming it is stored in the
default location ~/.ssh/). Additional keys may be added with this command,
quoting the entire public key as an argument.

Examples:
    juju add-ssh-key "ssh-rsa qYfS5LieM79HIOr535ret6xy
    AAAAB3NzaC1yc2EAAAADAQA6fgBAAABAQCygc6Rc9XgHdhQqTJ
    Wsoj+I3xGrOtk21xYtKijnhkGqItAHmrE5+VH6PY1rVIUXhpTg
    pSkJsHLmhE29OhIpt6yr8vQSOChqYfS5LieM79HIOJEgJEzIqC
    52rCYXLvr/BVkd6yr4IoM1vpb/n6u9o8v1a0VUGfc/J6tQAcPR
    ExzjZUVsfjj8HdLtcFq4JLYC41miiJtHw4b3qYu7qm3vh4eCiK
    1LqLncXnBCJfjj0pADXaL5OQ9dmD3aCbi8KFyOEs3UumPosgmh
    VCAfjjHObWHwNQ/ZU2KrX1/lv/+lBChx2tJliqQpyYMiA3nrtS
    jfqQgZfjVF5vz8LESQbGc6+vLcXZ9KQpuYDt joe@ubuntu"

For ease of use it is possible to use shell substitution to pass the key
to the command:

    juju add-ssh-key $(cat ~/mykey.pub)

See also:
    list-ssh-key
    remove-ssh-key
    import-ssh-key`[1:]

// NewAddKeysCommand is used to add a new ssh key to a model.
func NewAddKeysCommand() cmd.Command {
	return modelcmd.Wrap(&addKeysCommand{})
}

// addKeysCommand is used to add a new authorized ssh key for a user.
type addKeysCommand struct {
	SSHKeysBase
	user    string
	sshKeys []string
}

// Info implements Command.Info.
func (c *addKeysCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "add-ssh-key",
		Args:    "<ssh key> ...",
		Purpose: usageAddSSHKeySummary,
		Doc:     usageAddSSHKeyDetails,
		Aliases: []string{"add-ssh-keys"},
	}
}

// Init implements Command.Init.
func (c *addKeysCommand) Init(args []string) error {
	switch len(args) {
	case 0:
		return errors.New("no ssh key specified")
	default:
		c.sshKeys = args
	}
	return nil
}

// Run implements Command.Run.
func (c *addKeysCommand) Run(context *cmd.Context) error {
	client, err := c.NewKeyManagerClient()
	if err != nil {
		return err
	}
	defer client.Close()
	// TODO(alexisb) - currently keys are global which is not ideal.
	// keymanager needs to be updated to allow keys per user
	c.user = "admin"
	results, err := client.AddKeys(c.user, c.sshKeys...)
	if err != nil {
		return block.ProcessBlockedError(err, block.BlockChange)
	}
	for i, result := range results {
		if result.Error != nil {
			fmt.Fprintf(context.Stderr, "cannot add key %q: %v\n", c.sshKeys[i], result.Error)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/kops/cmd/kops/util"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/featureflag"
	"k8s.io/kops/upup/pkg/fi/cloudup"
	"k8s.io/kops/upup/pkg/kutil"
	"k8s.io/kops/util/pkg/tables"
)

// Command Object for a Rolling Update.
type RollingUpdateOptions struct {
	Yes       bool
	Force     bool
	CloudOnly bool

	// The following two variables are when kops is validating a cluster
	// during a rolling update.

	// FailOnDrainError fail rolling-update if drain errors.
	FailOnDrainError bool

	// FailOnValidate fail the cluster rolling-update when the cluster
	// does not validate, after a validation period.
	FailOnValidate bool

	DrainInterval time.Duration

	ValidateRetries int

	MasterInterval  time.Duration
	NodeInterval    time.Duration
	BastionInterval time.Duration

	ClusterName string

	// InstanceGroups is the list of instance groups to rolling-update;
	// if not specified all instance groups will be updated
	InstanceGroups []string
}

func (o *RollingUpdateOptions) InitDefaults() {
	o.Yes = false
	o.Force = false
	o.CloudOnly = false
	o.FailOnDrainError = false
	o.FailOnValidate = true

	o.MasterInterval = 5 * time.Minute
	o.NodeInterval = 2 * time.Minute
	o.BastionInterval = 5 * time.Minute

	o.ValidateRetries = 8

	o.DrainInterval = 90 * time.Second

}

func NewCmdRollingUpdateCluster(f *util.Factory, out io.Writer) *cobra.Command {

	var options RollingUpdateOptions
	options.InitDefaults()

	cmd := &cobra.Command{
		Use:   "cluster",
		Short: "Rolling update a cluster",
		Long: `Rolling update a cluster instance groups.

This command updates a kubernetes cluster to match the cloud, and kops specifications.

To perform rolling update, you need to update the cloud resources first with "kops update cluster"

Note: terraform users will need to "kops update cluster --target=terraform" then "terraform plan" then "terraform apply"
prior to running "kops rolling-update cluster"

Use KOPS_FEATURE_FLAGS="+DrainAndValidateRollingUpdate" to use beta code that drains the nodes
and validates the cluster.  New flags for Drain and Validation operations will be shown when
the environment variable is set.`,
	}

	cmd.Flags().BoolVar(&options.Yes, "yes", options.Yes, "perform rolling update without confirmation")
	cmd.Flags().BoolVar(&options.Force, "force", options.Force, "Force rolling update, even if no changes")
	cmd.Flags().BoolVar(&options.CloudOnly, "cloudonly", options.CloudOnly, "Perform rolling update without confirming progress with k8s")

	cmd.Flags().DurationVar(&options.MasterInterval, "master-interval", options.MasterInterval, "Time to wait between restarting masters")
	cmd.Flags().DurationVar(&options.NodeInterval, "node-interval", options.NodeInterval, "Time to wait between restarting nodes")
	cmd.Flags().DurationVar(&options.BastionInterval, "bastion-interval", options.BastionInterval, "Time to wait between restarting bastions")
	cmd.Flags().StringSliceVar(&options.InstanceGroups, "instance-group", options.InstanceGroups, "List of instance groups to update (defaults to all if not specified)")

	if featureflag.DrainAndValidateRollingUpdate.Enabled() {
		cmd.Flags().BoolVar(&options.FailOnDrainError, "fail-on-drain-error", true, "The rolling-update will fail if draining a node fails.")
		cmd.Flags().BoolVar(&options.FailOnValidate, "fail-on-validate-error", true, "The rolling-update will fail if the cluster fails to validate.")
		cmd.Flags().IntVar(&options.ValidateRetries, "validate-retries", options.ValidateRetries, "The number of times that a node will be validated.  Between validation kops sleeps the master-interval/2 or node-interval/2 duration.")
		cmd.Flags().DurationVar(&options.DrainInterval, "drain-interval", options.DrainInterval, "The duration that a rolling-update will wait after the node is drained.")
	}

	cmd.Run = func(cmd *cobra.Command, args []string) {
		err := rootCommand.ProcessArgs(args)
		if err != nil {
			exitWithError(err)
			return
		}

		clusterName := rootCommand.ClusterName()
		if clusterName == "" {
			exitWithError(fmt.Errorf("--name is required"))
			return
		}

		options.ClusterName = clusterName

		err = RunRollingUpdateCluster(f, os.Stdout, &options)
		if err != nil {
			exitWithError(err)
			return
		}

	}

	return cmd
}

func RunRollingUpdateCluster(f *util.Factory, out io.Writer, options *RollingUpdateOptions) error {

	clientset, err := f.Clientset()
	if err != nil {
		return err
	}

	cluster, err := GetCluster(f, options.ClusterName)
	if err != nil {
		return err
	}

	contextName := cluster.ObjectMeta.Name
	config, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(
		clientcmd.NewDefaultClientConfigLoadingRules(),
		&clientcmd.ConfigOverrides{CurrentContext: contextName}).ClientConfig()
	if err != nil {
		return fmt.Errorf("cannot load kubecfg settings for %q: %v", contextName, err)
	}

	if options.ValidateRetries <= 0 {
		return fmt.Errorf("validate-retries flag cannot be 0 or smaller")
	}

	var nodes []v1.Node
	var k8sClient kubernetes.Interface
	if !options.CloudOnly {
		k8sClient, err = kubernetes.NewForConfig(config)
		if err != nil {
			return fmt.Errorf("cannot build kube client for %q: %v", contextName, err)
		}

		nodeList, err := k8sClient.CoreV1().Nodes().List(metav1.ListOptions{})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to reach the kubernetes API.\n")
			fmt.Fprintf(os.Stderr, "Use --cloudonly to do a rolling-update without confirming progress with the k8s API\n\n")
			return fmt.Errorf("error listing nodes in cluster: %v", err)
		}

		if nodeList != nil {
			nodes = nodeList.Items
		}
	}

	list, err := clientset.InstanceGroups(cluster.ObjectMeta.Name).List(metav1.ListOptions{})
	if err != nil {
		return err
	}

	var instanceGroups []*api.InstanceGroup
	for i := range list.Items {
		instanceGroups = append(instanceGroups, &list.Items[i])
	}

	warnUnmatched := true

	if len(options.InstanceGroups) != 0 {
		var filtered []*api.InstanceGroup

		for _, instanceGroupName := range options.InstanceGroups {
			var found *api.InstanceGroup
			for _, ig := range instanceGroups {
				if ig.ObjectMeta.Name == instanceGroupName {
					found = ig
					break
				}
			}
			if found == nil {
				return fmt.Errorf("InstanceGroup %q not found", instanceGroupName)
			}

			filtered = append(filtered, found)
		}

		instanceGroups = filtered

		// Don't warn if we find more ASGs than IGs
		warnUnmatched = false
	}

	cloud, err := cloudup.BuildCloud(cluster)
	if err != nil {
		return err
	}

	groups, err := kutil.FindCloudInstanceGroups(cloud, cluster, instanceGroups, warnUnmatched, nodes)
	if err != nil {
		return err
	}

	{
		t := &tables.Table{}
		t.AddColumn("NAME", func(r *kutil.CloudInstanceGroup) string {
			return r.InstanceGroup.ObjectMeta.Name
		})
		t.AddColumn("STATUS", func(r *kutil.CloudInstanceGroup) string {
			return r.Status
		})
		t.AddColumn("NEEDUPDATE", func(r *kutil.CloudInstanceGroup) string {
			return strconv.Itoa(len(r.NeedUpdate))
		})
		t.AddColumn("READY", func(r *kutil.CloudInstanceGroup) string {
			return strconv.Itoa(len(r.Ready))
		})
		t.AddColumn("MIN", func(r *kutil.CloudInstanceGroup) string {
			return strconv.Itoa(r.MinSize())
		})
		t.AddColumn("MAX", func(r *kutil.CloudInstanceGroup) string {
			return strconv.Itoa(r.MaxSize())
		})
		t.AddColumn("NODES", func(r *kutil.CloudInstanceGroup) string {
			var nodes []*v1.Node
			for _, i := range r.Ready {
				if i.Node != nil {
					nodes = append(nodes, i.Node)
				}
			}
			for _, i := range r.NeedUpdate {
				if i.Node != nil {
					nodes = append(nodes, i.Node)
				}
			}
			return strconv.Itoa(len(nodes))
		})
		var l []*kutil.CloudInstanceGroup
		for _, v := range groups {
			l = append(l, v)
		}

		columns := []string{"NAME", "STATUS", "NEEDUPDATE", "READY", "MIN", "MAX"}
		if !options.CloudOnly {
			columns = append(columns, "NODES")
		}
		err := t.Render(l, out, columns...)
		if err != nil {
			return err
		}
	}

	needUpdate := false
	for _, group := range groups {
		if len(group.NeedUpdate) != 0 {
			needUpdate = true
		}
	}

	if !needUpdate && !options.Force {
		fmt.Printf("\nNo rolling-update required.\n")
		return nil
	}

	if !options.Yes {
		fmt.Printf("\nMust specify --yes to rolling-update.\n")
		return nil
	}

	if featureflag.DrainAndValidateRollingUpdate.Enabled() {
		glog.V(2).Infof("New rolling update with drain and validate enabled.")
	}
	d := &kutil.RollingUpdateCluster{
		MasterInterval:   options.MasterInterval,
		NodeInterval:     options.NodeInterval,
		Force:            options.Force,
		Cloud:            cloud,
		K8sClient:        k8sClient,
		ClientConfig:     kutil.NewClientConfig(config, "kube-system"),
		FailOnDrainError: options.FailOnDrainError,
		FailOnValidate:   options.FailOnValidate,
		CloudOnly:        options.CloudOnly,
		ClusterName:      options.ClusterName,
		ValidateRetries:  options.ValidateRetries,
		DrainInterval:    options.DrainInterval,
	}
	return d.RollingUpdate(groups, list)
}
<commit_msg>Specified that the commands need to be ran all in the same dir<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"fmt"
	"io"
	"os"
	"strconv"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/kops/cmd/kops/util"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/featureflag"
	"k8s.io/kops/upup/pkg/fi/cloudup"
	"k8s.io/kops/upup/pkg/kutil"
	"k8s.io/kops/util/pkg/tables"
)

// Command Object for a Rolling Update.
type RollingUpdateOptions struct {
	Yes       bool
	Force     bool
	CloudOnly bool

	// The following two variables are when kops is validating a cluster
	// during a rolling update.

	// FailOnDrainError fail rolling-update if drain errors.
	FailOnDrainError bool

	// FailOnValidate fail the cluster rolling-update when the cluster
	// does not validate, after a validation period.
	FailOnValidate bool

	DrainInterval time.Duration

	ValidateRetries int

	MasterInterval  time.Duration
	NodeInterval    time.Duration
	BastionInterval time.Duration

	ClusterName string

	// InstanceGroups is the list of instance groups to rolling-update;
	// if not specified all instance groups will be updated
	InstanceGroups []string
}

func (o *RollingUpdateOptions) InitDefaults() {
	o.Yes = false
	o.Force = false
	o.CloudOnly = false
	o.FailOnDrainError = false
	o.FailOnValidate = true

	o.MasterInterval = 5 * time.Minute
	o.NodeInterval = 2 * time.Minute
	o.BastionInterval = 5 * time.Minute

	o.ValidateRetries = 8

	o.DrainInterval = 90 * time.Second

}

func NewCmdRollingUpdateCluster(f *util.Factory, out io.Writer) *cobra.Command {

	var options RollingUpdateOptions
	options.InitDefaults()

	cmd := &cobra.Command{
		Use:   "cluster",
		Short: "Rolling update a cluster",
		Long: `Rolling update a cluster instance groups.

This command updates a kubernetes cluster to match the cloud, and kops specifications.

To perform rolling update, you need to update the cloud resources first with "kops update cluster"

Note: terraform users will need run the following commands all from the same directory `kops update cluster --target=terraform` then `terraform plan` then `terraform apply`
prior to running `kops rolling-update cluster`

Use KOPS_FEATURE_FLAGS="+DrainAndValidateRollingUpdate" to use beta code that drains the nodes
and validates the cluster.  New flags for Drain and Validation operations will be shown when
the environment variable is set.`,
	}

	cmd.Flags().BoolVar(&options.Yes, "yes", options.Yes, "perform rolling update without confirmation")
	cmd.Flags().BoolVar(&options.Force, "force", options.Force, "Force rolling update, even if no changes")
	cmd.Flags().BoolVar(&options.CloudOnly, "cloudonly", options.CloudOnly, "Perform rolling update without confirming progress with k8s")

	cmd.Flags().DurationVar(&options.MasterInterval, "master-interval", options.MasterInterval, "Time to wait between restarting masters")
	cmd.Flags().DurationVar(&options.NodeInterval, "node-interval", options.NodeInterval, "Time to wait between restarting nodes")
	cmd.Flags().DurationVar(&options.BastionInterval, "bastion-interval", options.BastionInterval, "Time to wait between restarting bastions")
	cmd.Flags().StringSliceVar(&options.InstanceGroups, "instance-group", options.InstanceGroups, "List of instance groups to update (defaults to all if not specified)")

	if featureflag.DrainAndValidateRollingUpdate.Enabled() {
		cmd.Flags().BoolVar(&options.FailOnDrainError, "fail-on-drain-error", true, "The rolling-update will fail if draining a node fails.")
		cmd.Flags().BoolVar(&options.FailOnValidate, "fail-on-validate-error", true, "The rolling-update will fail if the cluster fails to validate.")
		cmd.Flags().IntVar(&options.ValidateRetries, "validate-retries", options.ValidateRetries, "The number of times that a node will be validated.  Between validation kops sleeps the master-interval/2 or node-interval/2 duration.")
		cmd.Flags().DurationVar(&options.DrainInterval, "drain-interval", options.DrainInterval, "The duration that a rolling-update will wait after the node is drained.")
	}

	cmd.Run = func(cmd *cobra.Command, args []string) {
		err := rootCommand.ProcessArgs(args)
		if err != nil {
			exitWithError(err)
			return
		}

		clusterName := rootCommand.ClusterName()
		if clusterName == "" {
			exitWithError(fmt.Errorf("--name is required"))
			return
		}

		options.ClusterName = clusterName

		err = RunRollingUpdateCluster(f, os.Stdout, &options)
		if err != nil {
			exitWithError(err)
			return
		}

	}

	return cmd
}

func RunRollingUpdateCluster(f *util.Factory, out io.Writer, options *RollingUpdateOptions) error {

	clientset, err := f.Clientset()
	if err != nil {
		return err
	}

	cluster, err := GetCluster(f, options.ClusterName)
	if err != nil {
		return err
	}

	contextName := cluster.ObjectMeta.Name
	config, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(
		clientcmd.NewDefaultClientConfigLoadingRules(),
		&clientcmd.ConfigOverrides{CurrentContext: contextName}).ClientConfig()
	if err != nil {
		return fmt.Errorf("cannot load kubecfg settings for %q: %v", contextName, err)
	}

	if options.ValidateRetries <= 0 {
		return fmt.Errorf("validate-retries flag cannot be 0 or smaller")
	}

	var nodes []v1.Node
	var k8sClient kubernetes.Interface
	if !options.CloudOnly {
		k8sClient, err = kubernetes.NewForConfig(config)
		if err != nil {
			return fmt.Errorf("cannot build kube client for %q: %v", contextName, err)
		}

		nodeList, err := k8sClient.CoreV1().Nodes().List(metav1.ListOptions{})
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to reach the kubernetes API.\n")
			fmt.Fprintf(os.Stderr, "Use --cloudonly to do a rolling-update without confirming progress with the k8s API\n\n")
			return fmt.Errorf("error listing nodes in cluster: %v", err)
		}

		if nodeList != nil {
			nodes = nodeList.Items
		}
	}

	list, err := clientset.InstanceGroups(cluster.ObjectMeta.Name).List(metav1.ListOptions{})
	if err != nil {
		return err
	}

	var instanceGroups []*api.InstanceGroup
	for i := range list.Items {
		instanceGroups = append(instanceGroups, &list.Items[i])
	}

	warnUnmatched := true

	if len(options.InstanceGroups) != 0 {
		var filtered []*api.InstanceGroup

		for _, instanceGroupName := range options.InstanceGroups {
			var found *api.InstanceGroup
			for _, ig := range instanceGroups {
				if ig.ObjectMeta.Name == instanceGroupName {
					found = ig
					break
				}
			}
			if found == nil {
				return fmt.Errorf("InstanceGroup %q not found", instanceGroupName)
			}

			filtered = append(filtered, found)
		}

		instanceGroups = filtered

		// Don't warn if we find more ASGs than IGs
		warnUnmatched = false
	}

	cloud, err := cloudup.BuildCloud(cluster)
	if err != nil {
		return err
	}

	groups, err := kutil.FindCloudInstanceGroups(cloud, cluster, instanceGroups, warnUnmatched, nodes)
	if err != nil {
		return err
	}

	{
		t := &tables.Table{}
		t.AddColumn("NAME", func(r *kutil.CloudInstanceGroup) string {
			return r.InstanceGroup.ObjectMeta.Name
		})
		t.AddColumn("STATUS", func(r *kutil.CloudInstanceGroup) string {
			return r.Status
		})
		t.AddColumn("NEEDUPDATE", func(r *kutil.CloudInstanceGroup) string {
			return strconv.Itoa(len(r.NeedUpdate))
		})
		t.AddColumn("READY", func(r *kutil.CloudInstanceGroup) string {
			return strconv.Itoa(len(r.Ready))
		})
		t.AddColumn("MIN", func(r *kutil.CloudInstanceGroup) string {
			return strconv.Itoa(r.MinSize())
		})
		t.AddColumn("MAX", func(r *kutil.CloudInstanceGroup) string {
			return strconv.Itoa(r.MaxSize())
		})
		t.AddColumn("NODES", func(r *kutil.CloudInstanceGroup) string {
			var nodes []*v1.Node
			for _, i := range r.Ready {
				if i.Node != nil {
					nodes = append(nodes, i.Node)
				}
			}
			for _, i := range r.NeedUpdate {
				if i.Node != nil {
					nodes = append(nodes, i.Node)
				}
			}
			return strconv.Itoa(len(nodes))
		})
		var l []*kutil.CloudInstanceGroup
		for _, v := range groups {
			l = append(l, v)
		}

		columns := []string{"NAME", "STATUS", "NEEDUPDATE", "READY", "MIN", "MAX"}
		if !options.CloudOnly {
			columns = append(columns, "NODES")
		}
		err := t.Render(l, out, columns...)
		if err != nil {
			return err
		}
	}

	needUpdate := false
	for _, group := range groups {
		if len(group.NeedUpdate) != 0 {
			needUpdate = true
		}
	}

	if !needUpdate && !options.Force {
		fmt.Printf("\nNo rolling-update required.\n")
		return nil
	}

	if !options.Yes {
		fmt.Printf("\nMust specify --yes to rolling-update.\n")
		return nil
	}

	if featureflag.DrainAndValidateRollingUpdate.Enabled() {
		glog.V(2).Infof("New rolling update with drain and validate enabled.")
	}
	d := &kutil.RollingUpdateCluster{
		MasterInterval:   options.MasterInterval,
		NodeInterval:     options.NodeInterval,
		Force:            options.Force,
		Cloud:            cloud,
		K8sClient:        k8sClient,
		ClientConfig:     kutil.NewClientConfig(config, "kube-system"),
		FailOnDrainError: options.FailOnDrainError,
		FailOnValidate:   options.FailOnValidate,
		CloudOnly:        options.CloudOnly,
		ClusterName:      options.ClusterName,
		ValidateRetries:  options.ValidateRetries,
		DrainInterval:    options.DrainInterval,
	}
	return d.RollingUpdate(groups, list)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"fmt"

	"github.com/spf13/cobra"

	"k8s.io/minikube/pkg/minikube/constants"
	"k8s.io/minikube/pkg/minikube/notify"
	"k8s.io/minikube/pkg/version"
)

var updateCheckCmd = &cobra.Command{
	Use:   "update-check",
	Short: "Print current and latest version number",
	Long:  `Print current and latest version number`,
	PersistentPreRun: func(cmd *cobra.Command, args []string) {
		// Explicitly disable update checking for the version command
		enableUpdateNotification = false
	},
	Run: func(command *cobra.Command, args []string) {
		url := constants.GithubMinikubeReleasesURL
		r, err := notify.GetAllVersionsFromURL(url)
		if err != nil {
			fmt.Println("Error fetching latest version from internet")
			return
		}

		fmt.Println("CurrentVersion:", version.GetVersion())
		fmt.Println("LatestVersion:", r[0].Name)
	},
}

func init() {
	RootCmd.AddCommand(updateCheckCmd)
}
<commit_msg>Better error handling<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"

	"k8s.io/minikube/pkg/minikube/constants"
	"k8s.io/minikube/pkg/minikube/notify"
	"k8s.io/minikube/pkg/version"
)

var updateCheckCmd = &cobra.Command{
	Use:   "update-check",
	Short: "Print current and latest version number",
	Long:  `Print current and latest version number`,
	PersistentPreRun: func(cmd *cobra.Command, args []string) {
		// Explicitly disable update checking for the version command
		enableUpdateNotification = false
	},
	Run: func(command *cobra.Command, args []string) {
		url := constants.GithubMinikubeReleasesURL
		r, err := notify.GetAllVersionsFromURL(url)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error fetching latest version from internet")
			os.Exit(1)
		}

		if len(r) < 1 {
			fmt.Fprintf(os.Stderr, "Got empty version list from server")
			os.Exit(2)
		}

		fmt.Println("CurrentVersion:", version.GetVersion())
		fmt.Println("LatestVersion:", r[0].Name)
	},
}

func init() {
	RootCmd.AddCommand(updateCheckCmd)
}
<|endoftext|>


<commit_before>/**
 * Provides an API for triggering the execution of Watches.
 */

package main

import (
	// Utilities.
	"fmt"
	"net/http"

	// Gin.
	gin "gopkg.in/gin-gonic/gin.v1"

	// Internal dependencies.
	sdk "github.com/krystalcode/go-mantis-shrimp/actions/sdk"
	util "github.com/krystalcode/go-mantis-shrimp/util"
	common "github.com/krystalcode/go-mantis-shrimp/watches/common"
	config "github.com/krystalcode/go-mantis-shrimp/watches/config"
	storage "github.com/krystalcode/go-mantis-shrimp/watches/storage"
	wrapper "github.com/krystalcode/go-mantis-shrimp/watches/wrapper"
)

/**
 * Constants.
 */

// WatchAPIConfigFile holds the default path to the file containing the
// configuration for the Watch API.
const WatchAPIConfigFile = "/etc/mantis-shrimp/watch_api.config.json"

/**
 * Main program entry.
 */
func main() {
	// Load configuration.
	// @I Support providing configuration file for Watch API via cli options
	// @I Validate Watch API configuration when loading from JSON file
	var watchAPIConfig config.Config
	err := util.ReadJSONFile(WatchAPIConfigFile, &watchAPIConfig)
	if err != nil {
		panic(err)
	}

	// Load Watches provided in the config, if we run on ephemeral storage mode.
	loadEphemeralWatches(&watchAPIConfig)

	router := gin.Default()

	// Make storage available to the controllers.
	router.Use(Storage(watchAPIConfig.Storage))

	// Version 1 of the Watch API.
	v1 := router.Group("/v1")
	{
		// Create a new Watch.
		v1.POST("/", v1Create)

		// Trigger execution of the Watch via its ID.
		v1.POST("/:ids/trigger", v1Trigger)
	}

	/**
	 * @I Make the trigger API port configurable
	 */
	router.Run(":8888")
}

/**
 * Endpoint functions.
 */

// v1Create provides an endpoint that creates a new Watch based on the JSON object
// given in the request.
func v1Create(c *gin.Context) {
	/**
	 * @I Implement authentication of the caller
	 * @I Validate parameters per Watch type
	 * @I Ensure the caller has the permissions to create Watches
	 * @I Log errors and send a 500 response instead of panicking
	 * @I Implement creating and triggering a Watch in a single request
	 */

	// The parameters are provided as a JSON object in the request. Bind it to an
	// object of the corresponding type.
	var wrapper wrapper.WatchWrapper
	err := c.BindJSON(&wrapper)
	if err != nil {
		panic(err)
	}
	// @I Return 400 Bad Request if we are given no Watch type in a Create request
	if wrapper.Watch == nil {
		c.JSON(
			http.StatusBadRequest,
			gin.H{
				"status": http.StatusBadRequest,
			},
		)
		return
	}

	// Get the Watch as an object of the appropriate type.
	watch := wrapper.Watch

	// Store the Watch.
	storage := c.MustGet("storage").(storage.Storage)
	id, err := storage.Create(&watch)
	if err != nil {
		panic(err)
	}

	// All good.
	c.JSON(
		http.StatusOK,
		gin.H{
			"status": http.StatusOK,
			"id":     *id,
		},
	)
}

// v1Trigger provides an endpoint that triggers execution of the Action given in
// the request by its ID, by making a call to the Action API.
func v1Trigger(c *gin.Context) {
	/**
	 * @I Implement authentication of the caller
	 * @I Does the id need any escaping?
	 * @I Ensure the caller has the permissions to trigger evaluation of a Watch
	 * @I Investigate whether we need our own response status codes
	 */

	// The "ids" parameter is required. We allow for multiple comma-separated
	// string IDs, so we need to convert them to an array of integer IDs.
	// We want to make sure that the caller makes the request they want to without
	// mistakes, so we do not trigger any Watches if there is any error, even in
	// one of the IDs.
	// @I Refactor converting a comma-separated list of string IDs to an array of
	//    integer IDs into a utility function
	sIDs := c.Param("ids")
	aIDsInt, err := util.StringToIntegers(sIDs, ",")
	if err != nil {
		c.JSON(
			http.StatusNotFound,
			gin.H{
				"status": http.StatusNotFound,
			},
		)
		return
	}

	// Get the Watches with the requested IDs from storage.
	storage := c.MustGet("storage").(storage.Storage)

	var watches []*common.Watch
	for iID, _ := range aIDsInt {
		watch, err := storage.Get(iID)
		if err != nil {
			// Return a Not Found response if there is no Watch with such ID.
			if watch == nil {
				c.JSON(
					http.StatusNotFound,
					gin.H{
						"status": http.StatusNotFound,
					},
				)
			}
			return
		}

		// We could trigger the Watch at this point, however we prefer to check
		// that all Watches exist first.
		watches = append(watches, watch)
	}

	// Trigger execution of the Watches.
	// We only need to acknowledge that the Watches were triggered; we don't have to
	// for the execution to finish as this can take time.
	watchAPIConfig := c.MustGet("config").(config.Config)
	sdkConfig := sdk.Config{
		watchAPIConfig.ActionAPI.BaseURL,
		watchAPIConfig.ActionAPI.Version,
	}
	for _, pointer := range watches {
		go func() {
			watch := *pointer
			actionsIds := watch.Do()
			if len(actionsIds) == 0 {
				return
			}

			// @I Trigger all Watch Actions in one request
			for _, actionID := range actionsIds {
				go func() {
					err := sdk.TriggerByID(actionID, sdkConfig)
					if err != nil {
						// @I Investigate log management strategy for all services
						fmt.Println(err)
					}
				}()
			}
		}()
	}

	// All good.
	c.JSON(
		http.StatusOK,
		gin.H{
			"status": http.StatusOK,
		},
	)
}

/**
 * Middleware.
 */

// Storage is a Gin middleware that makes available the Storage engine to the
// endpoint controllers.
func Storage(config map[string]interface{}) gin.HandlerFunc {
	return func(c *gin.Context) {
		storage, err := storage.Create(config)
		if err != nil {
			panic(err)
		}
		c.Set("storage", storage)
		c.Next()
	}
}

// Config is a Gin middleware that makes available the Watch API configuration
// to the endpoint controllers.
func Config(watchAPIConfig *config.Config) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Set("config", *watchAPIConfig)
		c.Next()
	}
}

/**
 * Functions/types for internal use.
 */

// loadEphmeralWatches checks if the storage engine is configured to run in
// "ephemeral" mode, and if so, it loads into it any Watches contained in the
// configuration file.
func loadEphemeralWatches(watchAPIConfig *config.Config) {
	// @I Load init Watches directly in Redis via a script so that services
	//    don't have to be restarted together
	mode, ok := watchAPIConfig.Storage["mode"]
	if !ok || mode.(string) != "ephemeral" || watchAPIConfig.WatchWrappers == nil {
		return
	}

	storage, err := storage.Create(watchAPIConfig.Storage)
	if err != nil {
		panic(err)
	}

	for _, wrapper := range watchAPIConfig.WatchWrappers {
		_, err := storage.Create(&wrapper.Watch)
		if err != nil {
			panic(err)
		}
	}
}
<commit_msg>WatchAPI Simplified for statement<commit_after>/**
 * Provides an API for triggering the execution of Watches.
 */

package main

import (
	// Utilities.
	"fmt"
	"net/http"

	// Gin.
	gin "gopkg.in/gin-gonic/gin.v1"

	// Internal dependencies.
	sdk "github.com/krystalcode/go-mantis-shrimp/actions/sdk"
	util "github.com/krystalcode/go-mantis-shrimp/util"
	common "github.com/krystalcode/go-mantis-shrimp/watches/common"
	config "github.com/krystalcode/go-mantis-shrimp/watches/config"
	storage "github.com/krystalcode/go-mantis-shrimp/watches/storage"
	wrapper "github.com/krystalcode/go-mantis-shrimp/watches/wrapper"
)

/**
 * Constants.
 */

// WatchAPIConfigFile holds the default path to the file containing the
// configuration for the Watch API.
const WatchAPIConfigFile = "/etc/mantis-shrimp/watch_api.config.json"

/**
 * Main program entry.
 */
func main() {
	// Load configuration.
	// @I Support providing configuration file for Watch API via cli options
	// @I Validate Watch API configuration when loading from JSON file
	var watchAPIConfig config.Config
	err := util.ReadJSONFile(WatchAPIConfigFile, &watchAPIConfig)
	if err != nil {
		panic(err)
	}

	// Load Watches provided in the config, if we run on ephemeral storage mode.
	loadEphemeralWatches(&watchAPIConfig)

	router := gin.Default()

	// Make storage available to the controllers.
	router.Use(Storage(watchAPIConfig.Storage))

	// Version 1 of the Watch API.
	v1 := router.Group("/v1")
	{
		// Create a new Watch.
		v1.POST("/", v1Create)

		// Trigger execution of the Watch via its ID.
		v1.POST("/:ids/trigger", v1Trigger)
	}

	/**
	 * @I Make the trigger API port configurable
	 */
	router.Run(":8888")
}

/**
 * Endpoint functions.
 */

// v1Create provides an endpoint that creates a new Watch based on the JSON object
// given in the request.
func v1Create(c *gin.Context) {
	/**
	 * @I Implement authentication of the caller
	 * @I Validate parameters per Watch type
	 * @I Ensure the caller has the permissions to create Watches
	 * @I Log errors and send a 500 response instead of panicking
	 * @I Implement creating and triggering a Watch in a single request
	 */

	// The parameters are provided as a JSON object in the request. Bind it to an
	// object of the corresponding type.
	var wrapper wrapper.WatchWrapper
	err := c.BindJSON(&wrapper)
	if err != nil {
		panic(err)
	}
	// @I Return 400 Bad Request if we are given no Watch type in a Create request
	if wrapper.Watch == nil {
		c.JSON(
			http.StatusBadRequest,
			gin.H{
				"status": http.StatusBadRequest,
			},
		)
		return
	}

	// Get the Watch as an object of the appropriate type.
	watch := wrapper.Watch

	// Store the Watch.
	storage := c.MustGet("storage").(storage.Storage)
	id, err := storage.Create(&watch)
	if err != nil {
		panic(err)
	}

	// All good.
	c.JSON(
		http.StatusOK,
		gin.H{
			"status": http.StatusOK,
			"id":     *id,
		},
	)
}

// v1Trigger provides an endpoint that triggers execution of the Action given in
// the request by its ID, by making a call to the Action API.
func v1Trigger(c *gin.Context) {
	/**
	 * @I Implement authentication of the caller
	 * @I Does the id need any escaping?
	 * @I Ensure the caller has the permissions to trigger evaluation of a Watch
	 * @I Investigate whether we need our own response status codes
	 */

	// The "ids" parameter is required. We allow for multiple comma-separated
	// string IDs, so we need to convert them to an array of integer IDs.
	// We want to make sure that the caller makes the request they want to without
	// mistakes, so we do not trigger any Watches if there is any error, even in
	// one of the IDs.
	// @I Refactor converting a comma-separated list of string IDs to an array of
	//    integer IDs into a utility function
	sIDs := c.Param("ids")
	aIDsInt, err := util.StringToIntegers(sIDs, ",")
	if err != nil {
		c.JSON(
			http.StatusNotFound,
			gin.H{
				"status": http.StatusNotFound,
			},
		)
		return
	}

	// Get the Watches with the requested IDs from storage.
	storage := c.MustGet("storage").(storage.Storage)

	var watches []*common.Watch
	for iID := range aIDsInt {
		watch, err := storage.Get(iID)
		if err != nil {
			// Return a Not Found response if there is no Watch with such ID.
			if watch == nil {
				c.JSON(
					http.StatusNotFound,
					gin.H{
						"status": http.StatusNotFound,
					},
				)
			}
			return
		}

		// We could trigger the Watch at this point, however we prefer to check
		// that all Watches exist first.
		watches = append(watches, watch)
	}

	// Trigger execution of the Watches.
	// We only need to acknowledge that the Watches were triggered; we don't have to
	// for the execution to finish as this can take time.
	watchAPIConfig := c.MustGet("config").(config.Config)
	sdkConfig := sdk.Config{
		watchAPIConfig.ActionAPI.BaseURL,
		watchAPIConfig.ActionAPI.Version,
	}
	for _, pointer := range watches {
		go func() {
			watch := *pointer
			actionsIds := watch.Do()
			if len(actionsIds) == 0 {
				return
			}

			// @I Trigger all Watch Actions in one request
			for _, actionID := range actionsIds {
				go func() {
					err := sdk.TriggerByID(actionID, sdkConfig)
					if err != nil {
						// @I Investigate log management strategy for all services
						fmt.Println(err)
					}
				}()
			}
		}()
	}

	// All good.
	c.JSON(
		http.StatusOK,
		gin.H{
			"status": http.StatusOK,
		},
	)
}

/**
 * Middleware.
 */

// Storage is a Gin middleware that makes available the Storage engine to the
// endpoint controllers.
func Storage(config map[string]interface{}) gin.HandlerFunc {
	return func(c *gin.Context) {
		storage, err := storage.Create(config)
		if err != nil {
			panic(err)
		}
		c.Set("storage", storage)
		c.Next()
	}
}

// Config is a Gin middleware that makes available the Watch API configuration
// to the endpoint controllers.
func Config(watchAPIConfig *config.Config) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Set("config", *watchAPIConfig)
		c.Next()
	}
}

/**
 * Functions/types for internal use.
 */

// loadEphmeralWatches checks if the storage engine is configured to run in
// "ephemeral" mode, and if so, it loads into it any Watches contained in the
// configuration file.
func loadEphemeralWatches(watchAPIConfig *config.Config) {
	// @I Load init Watches directly in Redis via a script so that services
	//    don't have to be restarted together
	mode, ok := watchAPIConfig.Storage["mode"]
	if !ok || mode.(string) != "ephemeral" || watchAPIConfig.WatchWrappers == nil {
		return
	}

	storage, err := storage.Create(watchAPIConfig.Storage)
	if err != nil {
		panic(err)
	}

	for _, wrapper := range watchAPIConfig.WatchWrappers {
		_, err := storage.Create(&wrapper.Watch)
		if err != nil {
			panic(err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"testing"
	"time"

	"labix.org/v2/mgo"

	"github.com/stretchr/testify/assert"

	bsonScanner "github.com/Clever/oplog-replay/bson"
)

// simpleDocStruct is for inserting simple Mongo documents
type simpleDocStruct struct {
	key string
}

func TestComposingWithOplogReplay(t *testing.T) {
	// This test assumes that we're connecting to Mongo replica set. Otherwise, the oplog won't
	// be generated.

	unixTime := int(time.Now().Unix())

	time.Sleep(time.Duration(2) * time.Second)

	session, err := mgo.Dial("localhost")
	assert.Nil(t, err)
	db := session.DB("myTestDb")
	c := db.C("myCollection")
	assert.Nil(t, c.Insert(&simpleDocStruct{key: "key"}))

	time.Sleep(time.Duration(2) * time.Second)

	assert.Nil(t, c.Insert(&simpleDocStruct{key: "key2"}))
	// Dump at the time we started operations. Should get both operations
	dumpAtTime(t, unixTime, 2)
	// Three seconds later we should get only one.
	dumpAtTime(t, unixTime+3, 1)
}

func dumpAtTime(t *testing.T, unixTime, expectedResults int) {
	tempDir, err := ioutil.TempDir("/tmp", "oplogDumpTest")
	assert.Nil(t, err)
	defer os.RemoveAll(tempDir)
	assert.Nil(t, runDump(tempDir, "localhost", "", unixTime))

	file, err := os.Open(tempDir + "/local/oplog.rs.bson")
	assert.Nil(t, err)
	scanner := bsonScanner.New(file)
	count := 0
	for scanner.Scan() {
		count = count + 1
	}
	assert.Equal(t, expectedResults, count)
}

func TestCopyBsonFile(t *testing.T) {
	tempDir, err := ioutil.TempDir("/tmp", "oplogDumpTest")
	assert.Nil(t, err)
	defer os.RemoveAll(tempDir)

	// Create a directory structure that mirrors the oplog one the code expects
	err = os.Mkdir(tempDir+"/local", 0744)
	assert.Nil(t, err)
	file, err := os.Create(tempDir + "/local/oplog.rs.bson")
	assert.Nil(t, err)
	assert.Nil(t, ioutil.WriteFile(file.Name(), []byte("test-bson-file"), 0644))

	// Create a file to copy to
	toFile, err := ioutil.TempFile(tempDir, "bsonCopyTest")
	assert.Nil(t, err)
	assert.Nil(t, copyBsonFile(tempDir, toFile.Name()))

	// Check that the data matches
	fileData, err := ioutil.ReadFile(toFile.Name())
	assert.Nil(t, err)
	assert.Equal(t, "test-bson-file", string(fileData))
}
<commit_msg>Add test for filtration by collection<commit_after>package main

import (
	"io/ioutil"
	"os"
	"testing"
	"time"

	"labix.org/v2/mgo"

	"github.com/stretchr/testify/assert"

	bsonScanner "github.com/Clever/oplog-replay/bson"
)

// simpleDocStruct is for inserting simple Mongo documents
type simpleDocStruct struct {
	key string
}

func getPaddedTime() (n int) {
	time.Sleep(time.Duration(1) * time.Second)
	n = int(time.Now().Unix())
	time.Sleep(time.Duration(2) * time.Second)
	return
}

func TestComposingWithOplogReplay(t *testing.T) {
	// This test assumes that we're connecting to Mongo replica set. Otherwise, the oplog won't
	// be generated.

	unixTime := getPaddedTime()
	session, err := mgo.Dial("localhost")
	assert.Nil(t, err)
	db := session.DB("myTestDb")
	c := db.C("myCollection")
	assert.Nil(t, c.Insert(&simpleDocStruct{key: "key"}))

	time.Sleep(time.Duration(2) * time.Second)

	assert.Nil(t, c.Insert(&simpleDocStruct{key: "key2"}))
	// Dump at the time we started operations. Should get both operations
	dumpAtTime(t, unixTime, 2, "")
	// Three seconds later we should get only one.
	dumpAtTime(t, unixTime+3, 1, "")
}

func TestCollectionFiltering(t *testing.T) {
	// This test assumes that we're connecting to Mongo replica set. Otherwise, the oplog won't
	// be generated.

	unixTime := getPaddedTime()
	session, err := mgo.Dial("localhost")
	assert.Nil(t, err)
	db := session.DB("myTestDb")
	c1, c2, c3 := db.C("myCollection"), db.C("yourCollection"), db.C("ourCollection")

	assert.Nil(t, c1.Insert(&simpleDocStruct{key: "key"}))
	assert.Nil(t, c2.Insert(&simpleDocStruct{key: "key2"}))
	assert.Nil(t, c3.Insert(&simpleDocStruct{key: "key3"}))

	dumpAtTime(t, unixTime, 3, "")
	dumpAtTime(t, unixTime, 1, "myTestDb.myCollection")
	dumpAtTime(t, unixTime, 1, "myTestDb.yourCollection")
	dumpAtTime(t, unixTime, 2, "{$ne : \"myTestDb.myCollection\"}")
}

func dumpAtTime(t *testing.T, unixTime, expectedResults int, collection string) {
	tempDir, err := ioutil.TempDir("/tmp", "oplogDumpTest")
	assert.Nil(t, err)
	defer os.RemoveAll(tempDir)
	assert.Nil(t, runDump(tempDir, "localhost", collection, unixTime))

	file, err := os.Open(tempDir + "/local/oplog.rs.bson")
	assert.Nil(t, err)
	scanner := bsonScanner.New(file)
	count := 0
	for scanner.Scan() {
		count = count + 1
	}
	assert.Equal(t, expectedResults, count)
}

func TestCopyBsonFile(t *testing.T) {
	tempDir, err := ioutil.TempDir("/tmp", "oplogDumpTest")
	assert.Nil(t, err)
	defer os.RemoveAll(tempDir)

	// Create a directory structure that mirrors the oplog one the code expects
	err = os.Mkdir(tempDir+"/local", 0744)
	assert.Nil(t, err)
	file, err := os.Create(tempDir + "/local/oplog.rs.bson")
	assert.Nil(t, err)
	assert.Nil(t, ioutil.WriteFile(file.Name(), []byte("test-bson-file"), 0644))

	// Create a file to copy to
	toFile, err := ioutil.TempFile(tempDir, "bsonCopyTest")
	assert.Nil(t, err)
	assert.Nil(t, copyBsonFile(tempDir, toFile.Name()))

	// Check that the data matches
	fileData, err := ioutil.ReadFile(toFile.Name())
	assert.Nil(t, err)
	assert.Equal(t, "test-bson-file", string(fileData))
}
<|endoftext|>