<commit_before>/*
   Copyright (c) 2014-2015, Percona LLC and/or its affiliates. All rights reserved.

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

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

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

package qan

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/percona/cloud-protocol/proto"
	"github.com/percona/percona-agent/instance"
	"github.com/percona/percona-agent/mrms"
	"github.com/percona/percona-agent/mysql"
	"github.com/percona/percona-agent/pct"
	"github.com/percona/percona-agent/ticker"
)

// An AnalyzerInstnace is an Analyzer ran by a Manager, one per MySQL instance
// as configured.
type AnalyzerInstance struct {
	config      Config
	mysqlConn   mysql.Connector
	restartChan <-chan bool
	tickChan    chan time.Time
	analyzer    Analyzer
}

// A Manager runs AnalyzerInstances, one per MySQL instance as configured.
type Manager struct {
	logger          *pct.Logger
	clock           ticker.Manager
	im              *instance.Repo
	mrm             mrms.Monitor
	mysqlFactory    mysql.ConnectionFactory
	analyzerFactory AnalyzerFactory
	// --
	mux       *sync.RWMutex
	running   bool
	analyzers map[uint]AnalyzerInstance
	status    *pct.Status
}

func NewManager(
	logger *pct.Logger,
	clock ticker.Manager,
	im *instance.Repo,
	mrm mrms.Monitor,
	mysqlFactory mysql.ConnectionFactory,
	analyzerFactory AnalyzerFactory,
) *Manager {
	m := &Manager{
		logger:          logger,
		clock:           clock,
		im:              im,
		mrm:             mrm,
		mysqlFactory:    mysqlFactory,
		analyzerFactory: analyzerFactory,
		// --
		mux:       &sync.RWMutex{},
		analyzers: make(map[uint]AnalyzerInstance),
		status:    pct.NewStatus([]string{"qan"}),
	}
	return m
}

/////////////////////////////////////////////////////////////////////////////
// Interface
/////////////////////////////////////////////////////////////////////////////

func (m *Manager) Start() error {
	m.logger.Debug("Start:call")
	defer m.logger.Debug("Start:return")

	m.mux.Lock()
	defer m.mux.Unlock()

	if m.running {
		return pct.ServiceIsRunningError{Service: "qan"}
	}

	// Mangaer ("qan" in status) runs indepdent from qan-parser.
	m.status.Update("qan", "Starting")
	defer func() {
		m.running = true
		m.logger.Info("Started")
		m.status.Update("qan", "Running")
	}()

	// Load qan config from disk.
	// todo-1.1: get and start all qan-*.conf
	config := Config{}
	if err := pct.Basedir.ReadConfig("qan", &config); err != nil {
		if os.IsNotExist(err) {
			m.logger.Info("Not enabled")
			return nil
		}
		m.logger.Error("Read qan config:", err)
		return nil
	}

	// Start the slow log or perf schema analyzer. If it fails that's ok for
	// the qan manager itself (i.e. don't fail this func) because user can fix
	// or reconfigure this analyzer instannce later and have qan manager try
	// again to start it.
	// todo: this fails if agent starts before MySQL is running because MRMS
	//       fails to connect to MySQL in mrms/monitor/instance.NewMysqlInstance();
	//       it should succeed and retry until MySQL is online.
	if err := m.startAnalyzer(config); err != nil {
		m.logger.Error("Cannot start Query Analytics: %s. Verify that MySQL is running, " +
			"then try again.")
		return nil
	}

	return nil // success
}

func (m *Manager) Stop() error {
	m.logger.Debug("Stop:call")
	defer m.logger.Debug("Stop:return")

	m.mux.Lock()
	defer m.mux.Unlock()
	if !m.running {
		return nil
	}

	for instanceId := range m.analyzers {
		if err := m.stopAnalyzer(instanceId); err != nil {
			m.logger.Error(err)
		}
	}

	m.running = false
	m.logger.Info("Stopped")
	m.status.Update("qan", "Stopped")
	return nil
}

func (m *Manager) Status() map[string]string {
	m.mux.RLock()
	defer m.mux.RUnlock()
	status := m.status.All()
	for _, a := range m.analyzers {
		for k, v := range a.analyzer.Status() {
			status[k] = v
		}
	}
	return status
}

func (m *Manager) Handle(cmd *proto.Cmd) *proto.Reply {
	m.status.UpdateRe("qan", "Handling", cmd)
	defer m.status.Update("qan", "Running")

	switch cmd.Cmd {
	case "StartService":
		m.mux.Lock()
		defer m.mux.Unlock()
		if !m.running {
			return cmd.Reply(nil, pct.ServiceIsNotRunningError{Service: "qan"})
		}
		config := Config{}
		if err := json.Unmarshal(cmd.Data, &config); err != nil {
			return cmd.Reply(nil, err)
		}
		if err := m.startAnalyzer(config); err != nil {
			return cmd.Reply(nil, err)
		}
		// Write qan.conf to disk so agent runs qan on restart.

		if err := pct.Basedir.WriteConfig("qan", config); err != nil {
			return cmd.Reply(nil, err)
		}
		return cmd.Reply(nil) // success
	case "StopService":
		m.mux.Lock()
		defer m.mux.Unlock()
		if !m.running {
			return cmd.Reply(nil, pct.ServiceIsNotRunningError{Service: "qan"})
		}
		errs := []error{}
		for instanceId := range m.analyzers {
			if err := m.stopAnalyzer(instanceId); err != nil {
				errs = append(errs, err)
			}
		}
		// Remove qan.conf from disk so agent doesn't run qan on restart.
		if err := pct.Basedir.RemoveConfig("qan"); err != nil {
			errs = append(errs, err)
		}
		return cmd.Reply(nil, errs...)
	case "GetConfig":
		config, errs := m.GetConfig()
		return cmd.Reply(config, errs...)
	default:
		// SetConfig does not work by design.  To re-configure QAN,
		// stop it then start it again with the new config.
		return cmd.Reply(nil, pct.UnknownCmdError{Cmd: cmd.Cmd})
	}
}

func (m *Manager) GetConfig() ([]proto.AgentConfig, []error) {
	m.logger.Debug("GetConfig:call")
	defer m.logger.Debug("GetConfig:return")

	m.mux.RLock()
	defer m.mux.RUnlock()

	// Configs are always returned as array of AgentConfig resources.
	configs := []proto.AgentConfig{}
	for _, a := range m.analyzers {
		bytes, err := json.Marshal(a.config)
		if err != nil {
			m.logger.Warn(err)
			continue
		}
		configs = append(configs, proto.AgentConfig{
			InternalService: "qan",
			// no external service
			Config:  string(bytes),
			Running: true,
		})
	}
	return configs, nil
}

func ValidateConfig(config *Config) error {
	if config.CollectFrom == "" {
		// Before perf schema, CollectFrom didn't exist, so existing default QAN configs
		// don't have it.  To be backwards-compatible, no CollectFrom == slowlog.
		config.CollectFrom = "slowlog"
	}
	if config.CollectFrom != "slowlog" && config.CollectFrom != "perfschema" {
		return fmt.Errorf("Invalid CollectFrom: '%s'.  Expected 'perfschema' or 'slowlog'.", config.CollectFrom)
	}
	if config.Start == nil || len(config.Start) == 0 {
		return errors.New("qan.Config.Start array is empty")
	}
	if config.Stop == nil || len(config.Stop) == 0 {
		return errors.New("qan.Config.Stop array is empty")
	}
	if config.MaxWorkers < 1 {
		return errors.New("MaxWorkers must be > 0")
	}
	if config.MaxWorkers > 4 {
		return errors.New("MaxWorkers must be < 4")
	}
	if config.Interval == 0 {
		return errors.New("Interval must be > 0")
	}
	if config.Interval > 3600 {
		return errors.New("Interval must be <= 3600 (1 hour)")
	}
	if config.WorkerRunTime == 0 {
		return errors.New("WorkerRuntime must be > 0")
	}
	if config.WorkerRunTime > 1200 {
		return errors.New("WorkerRuntime must be <= 1200 (20 minutes)")
	}
	return nil
}

/////////////////////////////////////////////////////////////////////////////
// Implementation
/////////////////////////////////////////////////////////////////////////////

func (m *Manager) startAnalyzer(config Config) error {
	/*
		XXX Assume caller has locked m.mux.
	*/

	m.logger.Debug("startAnalyzer:call")
	defer m.logger.Debug("startAnalyzer:return")

	// Validate the config. This func may modify the config.
	if err := ValidateConfig(&config); err != nil {
		return fmt.Errorf("Invalid qan.Config: %s", err)
	}

	// Check if an analyzer for this MySQL instance already exists.
	if a, ok := m.analyzers[config.InstanceId]; ok {
		return pct.ServiceIsRunningError{Service: a.analyzer.String()}

	}

	// Get the MySQL DSN and create a MySQL connection.
	mysqlInstance := proto.MySQLInstance{}
	if err := m.im.Get(config.Service, config.InstanceId, &mysqlInstance); err != nil {
		return fmt.Errorf("Cannot get MySQL instance from repo: %s", err)
	}
	mysqlConn := m.mysqlFactory.Make(mysqlInstance.DSN)

	// Add the MySQL DSN to the MySQL restart monitor. If MySQL restarts,
	// the analyzer will stop its worker and re-configure MySQL.
	restartChan, err := m.mrm.Add(mysqlConn.DSN())
	if err != nil {
		return fmt.Errorf("Cannot add MySQL instance to restart monitor: %s", err)
	}

	// Make a chan on which the clock will tick at even intervals:
	// clock -> tickChan -> iter -> analyzer -> worker
	tickChan := make(chan time.Time, 1)
	m.clock.Add(tickChan, config.Interval, true)

	// Create and start a new analyzer. This should return immediately.
	// The analyzer will configure MySQL, start its iter, then run it worker
	// for each interval.
	analyzer := m.analyzerFactory.Make(
		config,
		"qan-analyzer", // todo-1.1: append instance name
		mysqlConn,
		restartChan,
		tickChan,
	)
	if err := analyzer.Start(); err != nil {
		return fmt.Errorf("Cannot start analyzer: %s", err)
	}

	// Save the new analyzer and its associated parts.
	m.analyzers[config.InstanceId] = AnalyzerInstance{
		config:      config,
		mysqlConn:   mysqlConn,
		restartChan: restartChan,
		tickChan:    tickChan,
		analyzer:    analyzer,
	}

	return nil // success
}

func (m *Manager) stopAnalyzer(instanceId uint) error {
	/*
		XXX Assume caller has locked m.mux.
	*/

	m.logger.Debug("stopAnalyzer:call")
	defer m.logger.Debug("stopAnalyzer:return")

	a, ok := m.analyzers[instanceId]
	if !ok {
		m.logger.Debug("stopAnalyzer:na", instanceId)
		return nil
	}

	m.status.Update("qan", fmt.Sprintf("Stopping %s", a.analyzer))
	m.logger.Info(fmt.Sprintf("Stopping %s", a.analyzer))

	// Stop ticking on this tickChan. Other tools receiving ticks at the same
	// interval are not affected.
	m.clock.Remove(a.tickChan)

	// Stop watching this MySQL instance. Other tools watching this MySQL
	// instance are not affected.
	m.mrm.Remove(a.mysqlConn.DSN(), a.restartChan)

	// Stop the analyzer. It stops its iter and worker and un-configures MySQL.
	if err := a.analyzer.Stop(); err != nil {
		return err
	}

	// Stop managing this analyzer.
	delete(m.analyzers, instanceId)

	// todo-1.1: remove the analyzer's config file?

	return nil // success
}
<commit_msg>Fix miss Errorf() param in qan.Manager err.<commit_after>/*
   Copyright (c) 2014-2015, Percona LLC and/or its affiliates. All rights reserved.

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

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

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

package qan

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/percona/cloud-protocol/proto"
	"github.com/percona/percona-agent/instance"
	"github.com/percona/percona-agent/mrms"
	"github.com/percona/percona-agent/mysql"
	"github.com/percona/percona-agent/pct"
	"github.com/percona/percona-agent/ticker"
)

// An AnalyzerInstnace is an Analyzer ran by a Manager, one per MySQL instance
// as configured.
type AnalyzerInstance struct {
	config      Config
	mysqlConn   mysql.Connector
	restartChan <-chan bool
	tickChan    chan time.Time
	analyzer    Analyzer
}

// A Manager runs AnalyzerInstances, one per MySQL instance as configured.
type Manager struct {
	logger          *pct.Logger
	clock           ticker.Manager
	im              *instance.Repo
	mrm             mrms.Monitor
	mysqlFactory    mysql.ConnectionFactory
	analyzerFactory AnalyzerFactory
	// --
	mux       *sync.RWMutex
	running   bool
	analyzers map[uint]AnalyzerInstance
	status    *pct.Status
}

func NewManager(
	logger *pct.Logger,
	clock ticker.Manager,
	im *instance.Repo,
	mrm mrms.Monitor,
	mysqlFactory mysql.ConnectionFactory,
	analyzerFactory AnalyzerFactory,
) *Manager {
	m := &Manager{
		logger:          logger,
		clock:           clock,
		im:              im,
		mrm:             mrm,
		mysqlFactory:    mysqlFactory,
		analyzerFactory: analyzerFactory,
		// --
		mux:       &sync.RWMutex{},
		analyzers: make(map[uint]AnalyzerInstance),
		status:    pct.NewStatus([]string{"qan"}),
	}
	return m
}

/////////////////////////////////////////////////////////////////////////////
// Interface
/////////////////////////////////////////////////////////////////////////////

func (m *Manager) Start() error {
	m.logger.Debug("Start:call")
	defer m.logger.Debug("Start:return")

	m.mux.Lock()
	defer m.mux.Unlock()

	if m.running {
		return pct.ServiceIsRunningError{Service: "qan"}
	}

	// Mangaer ("qan" in status) runs indepdent from qan-parser.
	m.status.Update("qan", "Starting")
	defer func() {
		m.running = true
		m.logger.Info("Started")
		m.status.Update("qan", "Running")
	}()

	// Load qan config from disk.
	// todo-1.1: get and start all qan-*.conf
	config := Config{}
	if err := pct.Basedir.ReadConfig("qan", &config); err != nil {
		if os.IsNotExist(err) {
			m.logger.Info("Not enabled")
			return nil
		}
		m.logger.Error("Read qan config:", err)
		return nil
	}

	// Start the slow log or perf schema analyzer. If it fails that's ok for
	// the qan manager itself (i.e. don't fail this func) because user can fix
	// or reconfigure this analyzer instannce later and have qan manager try
	// again to start it.
	// todo: this fails if agent starts before MySQL is running because MRMS
	//       fails to connect to MySQL in mrms/monitor/instance.NewMysqlInstance();
	//       it should succeed and retry until MySQL is online.
	if err := m.startAnalyzer(config); err != nil {
		m.logger.Error("Cannot start Query Analytics: %s. Verify that MySQL is running, "+
			"then try again.", err)
		return nil
	}

	return nil // success
}

func (m *Manager) Stop() error {
	m.logger.Debug("Stop:call")
	defer m.logger.Debug("Stop:return")

	m.mux.Lock()
	defer m.mux.Unlock()
	if !m.running {
		return nil
	}

	for instanceId := range m.analyzers {
		if err := m.stopAnalyzer(instanceId); err != nil {
			m.logger.Error(err)
		}
	}

	m.running = false
	m.logger.Info("Stopped")
	m.status.Update("qan", "Stopped")
	return nil
}

func (m *Manager) Status() map[string]string {
	m.mux.RLock()
	defer m.mux.RUnlock()
	status := m.status.All()
	for _, a := range m.analyzers {
		for k, v := range a.analyzer.Status() {
			status[k] = v
		}
	}
	return status
}

func (m *Manager) Handle(cmd *proto.Cmd) *proto.Reply {
	m.status.UpdateRe("qan", "Handling", cmd)
	defer m.status.Update("qan", "Running")

	switch cmd.Cmd {
	case "StartService":
		m.mux.Lock()
		defer m.mux.Unlock()
		if !m.running {
			return cmd.Reply(nil, pct.ServiceIsNotRunningError{Service: "qan"})
		}
		config := Config{}
		if err := json.Unmarshal(cmd.Data, &config); err != nil {
			return cmd.Reply(nil, err)
		}
		if err := m.startAnalyzer(config); err != nil {
			return cmd.Reply(nil, err)
		}
		// Write qan.conf to disk so agent runs qan on restart.

		if err := pct.Basedir.WriteConfig("qan", config); err != nil {
			return cmd.Reply(nil, err)
		}
		return cmd.Reply(nil) // success
	case "StopService":
		m.mux.Lock()
		defer m.mux.Unlock()
		if !m.running {
			return cmd.Reply(nil, pct.ServiceIsNotRunningError{Service: "qan"})
		}
		errs := []error{}
		for instanceId := range m.analyzers {
			if err := m.stopAnalyzer(instanceId); err != nil {
				errs = append(errs, err)
			}
		}
		// Remove qan.conf from disk so agent doesn't run qan on restart.
		if err := pct.Basedir.RemoveConfig("qan"); err != nil {
			errs = append(errs, err)
		}
		return cmd.Reply(nil, errs...)
	case "GetConfig":
		config, errs := m.GetConfig()
		return cmd.Reply(config, errs...)
	default:
		// SetConfig does not work by design.  To re-configure QAN,
		// stop it then start it again with the new config.
		return cmd.Reply(nil, pct.UnknownCmdError{Cmd: cmd.Cmd})
	}
}

func (m *Manager) GetConfig() ([]proto.AgentConfig, []error) {
	m.logger.Debug("GetConfig:call")
	defer m.logger.Debug("GetConfig:return")

	m.mux.RLock()
	defer m.mux.RUnlock()

	// Configs are always returned as array of AgentConfig resources.
	configs := []proto.AgentConfig{}
	for _, a := range m.analyzers {
		bytes, err := json.Marshal(a.config)
		if err != nil {
			m.logger.Warn(err)
			continue
		}
		configs = append(configs, proto.AgentConfig{
			InternalService: "qan",
			// no external service
			Config:  string(bytes),
			Running: true,
		})
	}
	return configs, nil
}

func ValidateConfig(config *Config) error {
	if config.CollectFrom == "" {
		// Before perf schema, CollectFrom didn't exist, so existing default QAN configs
		// don't have it.  To be backwards-compatible, no CollectFrom == slowlog.
		config.CollectFrom = "slowlog"
	}
	if config.CollectFrom != "slowlog" && config.CollectFrom != "perfschema" {
		return fmt.Errorf("Invalid CollectFrom: '%s'.  Expected 'perfschema' or 'slowlog'.", config.CollectFrom)
	}
	if config.Start == nil || len(config.Start) == 0 {
		return errors.New("qan.Config.Start array is empty")
	}
	if config.Stop == nil || len(config.Stop) == 0 {
		return errors.New("qan.Config.Stop array is empty")
	}
	if config.MaxWorkers < 1 {
		return errors.New("MaxWorkers must be > 0")
	}
	if config.MaxWorkers > 4 {
		return errors.New("MaxWorkers must be < 4")
	}
	if config.Interval == 0 {
		return errors.New("Interval must be > 0")
	}
	if config.Interval > 3600 {
		return errors.New("Interval must be <= 3600 (1 hour)")
	}
	if config.WorkerRunTime == 0 {
		return errors.New("WorkerRuntime must be > 0")
	}
	if config.WorkerRunTime > 1200 {
		return errors.New("WorkerRuntime must be <= 1200 (20 minutes)")
	}
	return nil
}

/////////////////////////////////////////////////////////////////////////////
// Implementation
/////////////////////////////////////////////////////////////////////////////

func (m *Manager) startAnalyzer(config Config) error {
	/*
		XXX Assume caller has locked m.mux.
	*/

	m.logger.Debug("startAnalyzer:call")
	defer m.logger.Debug("startAnalyzer:return")

	// Validate the config. This func may modify the config.
	if err := ValidateConfig(&config); err != nil {
		return fmt.Errorf("Invalid qan.Config: %s", err)
	}

	// Check if an analyzer for this MySQL instance already exists.
	if a, ok := m.analyzers[config.InstanceId]; ok {
		return pct.ServiceIsRunningError{Service: a.analyzer.String()}

	}

	// Get the MySQL DSN and create a MySQL connection.
	mysqlInstance := proto.MySQLInstance{}
	if err := m.im.Get(config.Service, config.InstanceId, &mysqlInstance); err != nil {
		return fmt.Errorf("Cannot get MySQL instance from repo: %s", err)
	}
	mysqlConn := m.mysqlFactory.Make(mysqlInstance.DSN)

	// Add the MySQL DSN to the MySQL restart monitor. If MySQL restarts,
	// the analyzer will stop its worker and re-configure MySQL.
	restartChan, err := m.mrm.Add(mysqlConn.DSN())
	if err != nil {
		return fmt.Errorf("Cannot add MySQL instance to restart monitor: %s", err)
	}

	// Make a chan on which the clock will tick at even intervals:
	// clock -> tickChan -> iter -> analyzer -> worker
	tickChan := make(chan time.Time, 1)
	m.clock.Add(tickChan, config.Interval, true)

	// Create and start a new analyzer. This should return immediately.
	// The analyzer will configure MySQL, start its iter, then run it worker
	// for each interval.
	analyzer := m.analyzerFactory.Make(
		config,
		"qan-analyzer", // todo-1.1: append instance name
		mysqlConn,
		restartChan,
		tickChan,
	)
	if err := analyzer.Start(); err != nil {
		return fmt.Errorf("Cannot start analyzer: %s", err)
	}

	// Save the new analyzer and its associated parts.
	m.analyzers[config.InstanceId] = AnalyzerInstance{
		config:      config,
		mysqlConn:   mysqlConn,
		restartChan: restartChan,
		tickChan:    tickChan,
		analyzer:    analyzer,
	}

	return nil // success
}

func (m *Manager) stopAnalyzer(instanceId uint) error {
	/*
		XXX Assume caller has locked m.mux.
	*/

	m.logger.Debug("stopAnalyzer:call")
	defer m.logger.Debug("stopAnalyzer:return")

	a, ok := m.analyzers[instanceId]
	if !ok {
		m.logger.Debug("stopAnalyzer:na", instanceId)
		return nil
	}

	m.status.Update("qan", fmt.Sprintf("Stopping %s", a.analyzer))
	m.logger.Info(fmt.Sprintf("Stopping %s", a.analyzer))

	// Stop ticking on this tickChan. Other tools receiving ticks at the same
	// interval are not affected.
	m.clock.Remove(a.tickChan)

	// Stop watching this MySQL instance. Other tools watching this MySQL
	// instance are not affected.
	m.mrm.Remove(a.mysqlConn.DSN(), a.restartChan)

	// Stop the analyzer. It stops its iter and worker and un-configures MySQL.
	if err := a.analyzer.Stop(); err != nil {
		return err
	}

	// Stop managing this analyzer.
	delete(m.analyzers, instanceId)

	// todo-1.1: remove the analyzer's config file?

	return nil // success
}
<|endoftext|>


<commit_before>package gorethink

import (
	"crypto/tls"
	"sync"
	"time"

	p "gopkg.in/gorethink/gorethink.v3/ql2"
)

// A Session represents a connection to a RethinkDB cluster and should be used
// when executing queries.
type Session struct {
	hosts []Host
	opts  *ConnectOpts

	mu      sync.RWMutex
	cluster *Cluster
	closed  bool
}

// ConnectOpts is used to specify optional arguments when connecting to a cluster.
type ConnectOpts struct {
	// Address holds the address of the server initially used when creating the
	// session. Only used if Addresses is empty
	Address string `gorethink:"address,omitempty"`
	// Addresses holds the addresses of the servers initially used when creating
	// the session.
	Addresses []string `gorethink:"addresses,omitempty"`
	// Database is the default database name used when executing queries, this
	// value is only used if the query does not contain any DB term
	Database string `gorethink:"database,omitempty"`
	// Username holds the username used for authentication, if blank (and the v1
	// handshake protocol is being used) then the admin user is used
	Username string `gorethink:"username,omitempty"`
	// Password holds the password used for authentication (only used when using
	// the v1 handshake protocol)
	Password string `gorethink:"password,omitempty"`
	// AuthKey is used for authentication when using the v0.4 handshake protocol
	// This field is no deprecated
	AuthKey string `gorethink:"authkey,omitempty"`
	// Timeout is the time the driver waits when creating new connections, to
	// configure the timeout used when executing queries use WriteTimeout and
	// ReadTimeout
	Timeout time.Duration `gorethink:"timeout,omitempty"`
	// WriteTimeout is the amount of time the driver will wait when sending the
	// query to the server
	WriteTimeout time.Duration `gorethink:"write_timeout,omitempty"`
	// ReadTimeout is the amount of time the driver will wait for a response from
	// the server when executing queries.
	ReadTimeout time.Duration `gorethink:"read_timeout,omitempty"`
	// KeepAlivePeriod is the keep alive period used by the connection, by default
	// this is 30s. It is not possible to disable keep alive messages
	KeepAlivePeriod time.Duration `gorethink:"keep_alive_timeout,omitempty"`
	// TLSConfig holds the TLS configuration and can be used when connecting
	// to a RethinkDB server protected by SSL
	TLSConfig *tls.Config `gorethink:"tlsconfig,omitempty"`
	// HandshakeVersion is used to specify which handshake version should be
	// used, this currently defaults to v1 which is used by RethinkDB 2.3 and
	// later. If you are using an older version then you can set the handshake
	// version to 0.4
	HandshakeVersion HandshakeVersion `gorethink:"handshake_version,omitempty"`
	// UseJSONNumber indicates whether the cursors running in this session should
	// use json.Number instead of float64 while unmarshaling documents with
	// interface{}. The default is `false`.
	UseJSONNumber bool
	// NumRetries is the number of times a query is retried if a connection
	// error is detected, queries are not retried if RethinkDB returns a
	// runtime error.
	NumRetries int

	// InitialCap is used by the internal connection pool and is used to
	// configure how many connections are created for each host when the
	// session is created. If zero then no connections are created until
	// the first query is executed.
	InitialCap int `gorethink:"initial_cap,omitempty"`
	// MaxOpen is used by the internal connection pool and is used to configure
	// the maximum number of connections held in the pool. If all available
	// connections are being used then the driver will open new connections as
	// needed however they will not be returned to the pool. By default the
	// maximum number of connections is 2
	MaxOpen int `gorethink:"max_open,omitempty"`

	// Below options are for cluster discovery, please note there is a high
	// probability of these changing as the API is still being worked on.

	// DiscoverHosts is used to enable host discovery, when true the driver
	// will attempt to discover any new nodes added to the cluster and then
	// start sending queries to these new nodes.
	DiscoverHosts bool `gorethink:"discover_hosts,omitempty"`
	// HostDecayDuration is used by the go-hostpool package to calculate a weighted
	// score when selecting a host. By default a value of 5 minutes is used.
	HostDecayDuration time.Duration

	// Deprecated: This function is no longer used due to changes in the
	// way hosts are selected.
	NodeRefreshInterval time.Duration `gorethink:"node_refresh_interval,omitempty"`
	// Deprecated: Use InitialCap instead
	MaxIdle int `gorethink:"max_idle,omitempty"`
}

func (o ConnectOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Connect creates a new database session. To view the available connection
// options see ConnectOpts.
//
// By default maxIdle and maxOpen are set to 1: passing values greater
// than the default (e.g. MaxIdle: "10", MaxOpen: "20") will provide a
// pool of re-usable connections.
//
// Basic connection example:
//
// 	session, err := r.Connect(r.ConnectOpts{
// 		Host: "localhost:28015",
// 		Database: "test",
// 		AuthKey:  "14daak1cad13dj",
// 	})
//
// Cluster connection example:
//
// 	session, err := r.Connect(r.ConnectOpts{
// 		Hosts: []string{"localhost:28015", "localhost:28016"},
// 		Database: "test",
// 		AuthKey:  "14daak1cad13dj",
// 	})
func Connect(opts ConnectOpts) (*Session, error) {
	var addresses = opts.Addresses
	if len(addresses) == 0 {
		addresses = []string{opts.Address}
	}

	hosts := make([]Host, len(addresses))
	for i, address := range addresses {
		hostname, port := splitAddress(address)
		hosts[i] = NewHost(hostname, port)
	}
	if len(hosts) <= 0 {
		return nil, ErrNoHosts
	}

	// Connect
	s := &Session{
		hosts: hosts,
		opts:  &opts,
	}

	err := s.Reconnect()
	if err != nil {
		// note: s.Reconnect() will initialize cluster information which
		// will cause the .IsConnected() method to be caught in a loop
		return &Session{
			hosts: hosts,
			opts:  &opts,
		}, err
	}

	return s, nil
}

// CloseOpts allows calls to the Close function to be configured.
type CloseOpts struct {
	NoReplyWait bool `gorethink:"noreplyWait,omitempty"`
}

func (o CloseOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// IsConnected returns true if session has a valid connection.
func (s *Session) IsConnected() bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.cluster == nil || s.closed {
		return false
	}
	return s.cluster.IsConnected()
}

// Reconnect closes and re-opens a session.
func (s *Session) Reconnect(optArgs ...CloseOpts) error {
	var err error

	if err = s.Close(optArgs...); err != nil {
		return err
	}

	s.mu.Lock()
	s.cluster, err = NewCluster(s.hosts, s.opts)
	if err != nil {
		s.mu.Unlock()
		return err
	}

	s.closed = false
	s.mu.Unlock()

	return nil
}

// Close closes the session
func (s *Session) Close(optArgs ...CloseOpts) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return nil
	}

	if len(optArgs) >= 1 {
		if optArgs[0].NoReplyWait {
			s.mu.Unlock()
			s.NoReplyWait()
			s.mu.Lock()
		}
	}

	if s.cluster != nil {
		s.cluster.Close()
	}
	s.cluster = nil
	s.closed = true

	return nil
}

// SetInitialPoolCap sets the initial capacity of the connection pool.
func (s *Session) SetInitialPoolCap(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.InitialCap = n
	s.cluster.SetInitialPoolCap(n)
}

// SetMaxIdleConns sets the maximum number of connections in the idle
// connection pool.
func (s *Session) SetMaxIdleConns(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.MaxIdle = n
	s.cluster.SetMaxIdleConns(n)
}

// SetMaxOpenConns sets the maximum number of open connections to the database.
func (s *Session) SetMaxOpenConns(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.MaxOpen = n
	s.cluster.SetMaxOpenConns(n)
}

// NoReplyWait ensures that previous queries with the noreply flag have been
// processed by the server. Note that this guarantee only applies to queries
// run on the given connection
func (s *Session) NoReplyWait() error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return ErrConnectionClosed
	}

	return s.cluster.Exec(Query{
		Type: p.Query_NOREPLY_WAIT,
	})
}

// Use changes the default database used
func (s *Session) Use(database string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.Database = database
}

// Database returns the selected database set by Use
func (s *Session) Database() string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.opts.Database
}

// Query executes a ReQL query using the session to connect to the database
func (s *Session) Query(q Query) (*Cursor, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrConnectionClosed
	}

	return s.cluster.Query(q)
}

// Exec executes a ReQL query using the session to connect to the database
func (s *Session) Exec(q Query) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return ErrConnectionClosed
	}

	return s.cluster.Exec(q)
}

// Server returns the server name and server UUID being used by a connection.
func (s *Session) Server() (ServerResponse, error) {
	return s.cluster.Server()
}

// SetHosts resets the hosts used when connecting to the RethinkDB cluster
func (s *Session) SetHosts(hosts []Host) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.hosts = hosts
}

func (s *Session) newQuery(t Term, opts map[string]interface{}) (Query, error) {
	return newQuery(t, opts, s.opts)
}
<commit_msg>Session.IsConnected: use RLock instead of Lock to improve performance<commit_after>package gorethink

import (
	"crypto/tls"
	"sync"
	"time"

	p "gopkg.in/gorethink/gorethink.v3/ql2"
)

// A Session represents a connection to a RethinkDB cluster and should be used
// when executing queries.
type Session struct {
	hosts []Host
	opts  *ConnectOpts

	mu      sync.RWMutex
	cluster *Cluster
	closed  bool
}

// ConnectOpts is used to specify optional arguments when connecting to a cluster.
type ConnectOpts struct {
	// Address holds the address of the server initially used when creating the
	// session. Only used if Addresses is empty
	Address string `gorethink:"address,omitempty"`
	// Addresses holds the addresses of the servers initially used when creating
	// the session.
	Addresses []string `gorethink:"addresses,omitempty"`
	// Database is the default database name used when executing queries, this
	// value is only used if the query does not contain any DB term
	Database string `gorethink:"database,omitempty"`
	// Username holds the username used for authentication, if blank (and the v1
	// handshake protocol is being used) then the admin user is used
	Username string `gorethink:"username,omitempty"`
	// Password holds the password used for authentication (only used when using
	// the v1 handshake protocol)
	Password string `gorethink:"password,omitempty"`
	// AuthKey is used for authentication when using the v0.4 handshake protocol
	// This field is no deprecated
	AuthKey string `gorethink:"authkey,omitempty"`
	// Timeout is the time the driver waits when creating new connections, to
	// configure the timeout used when executing queries use WriteTimeout and
	// ReadTimeout
	Timeout time.Duration `gorethink:"timeout,omitempty"`
	// WriteTimeout is the amount of time the driver will wait when sending the
	// query to the server
	WriteTimeout time.Duration `gorethink:"write_timeout,omitempty"`
	// ReadTimeout is the amount of time the driver will wait for a response from
	// the server when executing queries.
	ReadTimeout time.Duration `gorethink:"read_timeout,omitempty"`
	// KeepAlivePeriod is the keep alive period used by the connection, by default
	// this is 30s. It is not possible to disable keep alive messages
	KeepAlivePeriod time.Duration `gorethink:"keep_alive_timeout,omitempty"`
	// TLSConfig holds the TLS configuration and can be used when connecting
	// to a RethinkDB server protected by SSL
	TLSConfig *tls.Config `gorethink:"tlsconfig,omitempty"`
	// HandshakeVersion is used to specify which handshake version should be
	// used, this currently defaults to v1 which is used by RethinkDB 2.3 and
	// later. If you are using an older version then you can set the handshake
	// version to 0.4
	HandshakeVersion HandshakeVersion `gorethink:"handshake_version,omitempty"`
	// UseJSONNumber indicates whether the cursors running in this session should
	// use json.Number instead of float64 while unmarshaling documents with
	// interface{}. The default is `false`.
	UseJSONNumber bool
	// NumRetries is the number of times a query is retried if a connection
	// error is detected, queries are not retried if RethinkDB returns a
	// runtime error.
	NumRetries int

	// InitialCap is used by the internal connection pool and is used to
	// configure how many connections are created for each host when the
	// session is created. If zero then no connections are created until
	// the first query is executed.
	InitialCap int `gorethink:"initial_cap,omitempty"`
	// MaxOpen is used by the internal connection pool and is used to configure
	// the maximum number of connections held in the pool. If all available
	// connections are being used then the driver will open new connections as
	// needed however they will not be returned to the pool. By default the
	// maximum number of connections is 2
	MaxOpen int `gorethink:"max_open,omitempty"`

	// Below options are for cluster discovery, please note there is a high
	// probability of these changing as the API is still being worked on.

	// DiscoverHosts is used to enable host discovery, when true the driver
	// will attempt to discover any new nodes added to the cluster and then
	// start sending queries to these new nodes.
	DiscoverHosts bool `gorethink:"discover_hosts,omitempty"`
	// HostDecayDuration is used by the go-hostpool package to calculate a weighted
	// score when selecting a host. By default a value of 5 minutes is used.
	HostDecayDuration time.Duration

	// Deprecated: This function is no longer used due to changes in the
	// way hosts are selected.
	NodeRefreshInterval time.Duration `gorethink:"node_refresh_interval,omitempty"`
	// Deprecated: Use InitialCap instead
	MaxIdle int `gorethink:"max_idle,omitempty"`
}

func (o ConnectOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Connect creates a new database session. To view the available connection
// options see ConnectOpts.
//
// By default maxIdle and maxOpen are set to 1: passing values greater
// than the default (e.g. MaxIdle: "10", MaxOpen: "20") will provide a
// pool of re-usable connections.
//
// Basic connection example:
//
// 	session, err := r.Connect(r.ConnectOpts{
// 		Host: "localhost:28015",
// 		Database: "test",
// 		AuthKey:  "14daak1cad13dj",
// 	})
//
// Cluster connection example:
//
// 	session, err := r.Connect(r.ConnectOpts{
// 		Hosts: []string{"localhost:28015", "localhost:28016"},
// 		Database: "test",
// 		AuthKey:  "14daak1cad13dj",
// 	})
func Connect(opts ConnectOpts) (*Session, error) {
	var addresses = opts.Addresses
	if len(addresses) == 0 {
		addresses = []string{opts.Address}
	}

	hosts := make([]Host, len(addresses))
	for i, address := range addresses {
		hostname, port := splitAddress(address)
		hosts[i] = NewHost(hostname, port)
	}
	if len(hosts) <= 0 {
		return nil, ErrNoHosts
	}

	// Connect
	s := &Session{
		hosts: hosts,
		opts:  &opts,
	}

	err := s.Reconnect()
	if err != nil {
		// note: s.Reconnect() will initialize cluster information which
		// will cause the .IsConnected() method to be caught in a loop
		return &Session{
			hosts: hosts,
			opts:  &opts,
		}, err
	}

	return s, nil
}

// CloseOpts allows calls to the Close function to be configured.
type CloseOpts struct {
	NoReplyWait bool `gorethink:"noreplyWait,omitempty"`
}

func (o CloseOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// IsConnected returns true if session has a valid connection.
func (s *Session) IsConnected() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.cluster == nil || s.closed {
		return false
	}
	return s.cluster.IsConnected()
}

// Reconnect closes and re-opens a session.
func (s *Session) Reconnect(optArgs ...CloseOpts) error {
	var err error

	if err = s.Close(optArgs...); err != nil {
		return err
	}

	s.mu.Lock()
	s.cluster, err = NewCluster(s.hosts, s.opts)
	if err != nil {
		s.mu.Unlock()
		return err
	}

	s.closed = false
	s.mu.Unlock()

	return nil
}

// Close closes the session
func (s *Session) Close(optArgs ...CloseOpts) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return nil
	}

	if len(optArgs) >= 1 {
		if optArgs[0].NoReplyWait {
			s.mu.Unlock()
			s.NoReplyWait()
			s.mu.Lock()
		}
	}

	if s.cluster != nil {
		s.cluster.Close()
	}
	s.cluster = nil
	s.closed = true

	return nil
}

// SetInitialPoolCap sets the initial capacity of the connection pool.
func (s *Session) SetInitialPoolCap(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.InitialCap = n
	s.cluster.SetInitialPoolCap(n)
}

// SetMaxIdleConns sets the maximum number of connections in the idle
// connection pool.
func (s *Session) SetMaxIdleConns(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.MaxIdle = n
	s.cluster.SetMaxIdleConns(n)
}

// SetMaxOpenConns sets the maximum number of open connections to the database.
func (s *Session) SetMaxOpenConns(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.MaxOpen = n
	s.cluster.SetMaxOpenConns(n)
}

// NoReplyWait ensures that previous queries with the noreply flag have been
// processed by the server. Note that this guarantee only applies to queries
// run on the given connection
func (s *Session) NoReplyWait() error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return ErrConnectionClosed
	}

	return s.cluster.Exec(Query{
		Type: p.Query_NOREPLY_WAIT,
	})
}

// Use changes the default database used
func (s *Session) Use(database string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.Database = database
}

// Database returns the selected database set by Use
func (s *Session) Database() string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.opts.Database
}

// Query executes a ReQL query using the session to connect to the database
func (s *Session) Query(q Query) (*Cursor, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrConnectionClosed
	}

	return s.cluster.Query(q)
}

// Exec executes a ReQL query using the session to connect to the database
func (s *Session) Exec(q Query) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return ErrConnectionClosed
	}

	return s.cluster.Exec(q)
}

// Server returns the server name and server UUID being used by a connection.
func (s *Session) Server() (ServerResponse, error) {
	return s.cluster.Server()
}

// SetHosts resets the hosts used when connecting to the RethinkDB cluster
func (s *Session) SetHosts(hosts []Host) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.hosts = hosts
}

func (s *Session) newQuery(t Term, opts map[string]interface{}) (Query, error) {
	return newQuery(t, opts, s.opts)
}
<|endoftext|>


<commit_before>package cloudbrain

import (
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/hashicorp/go-multierror"
	"github.com/pborman/uuid"
	"github.com/pkg/errors"
	"github.com/travis-ci/cloud-brain/background"
	"github.com/travis-ci/cloud-brain/cbcontext"
	"github.com/travis-ci/cloud-brain/cloud"
	"github.com/travis-ci/cloud-brain/database"
	"golang.org/x/crypto/scrypt"
	"golang.org/x/net/context"
	"gopkg.in/urfave/cli.v2"
)

var (
	//VersionString gets set during `make`
	VersionString = "?"
	//RevisionString gets set during `make`
	RevisionString = "?"
	//RevisionURLString gets set during `make`
	RevisionURLString = "?"
	//GeneratedString gets set during `make`
	GeneratedString = "?"
	//CopyrightString gets set during `make`
	CopyrightString = "?"
)

func init() {
	cli.VersionPrinter = customVersionPrinter
	_ = os.Setenv("VERSION", VersionString)
	_ = os.Setenv("REVISION", RevisionString)
	_ = os.Setenv("GENERATED", GeneratedString)
}

func customVersionPrinter(c *cli.Context) {
	fmt.Printf("%v v=%v rev=%v d=%v\n", filepath.Base(c.App.Name),
		VersionString, RevisionString, GeneratedString)
}

// MaxCreateRetries is the number of times the "create" job will be retried.
const MaxCreateRetries = 10

// Core is used as a central manager for all Cloud Brain functionality. The HTTP
// API and the background workers are just frontends for the Core, and calls
// methods on Core for functionality.
type Core struct {
	db database.DB
	bb background.Backend

	cloudProvidersMutex sync.Mutex
	cloudProviders      map[string]cloud.Provider
}

// NewCore is used to create a new Core backed by the given database and
// background Backend.
func NewCore(db database.DB, bb background.Backend) *Core {
	return &Core{
		db: db,
		bb: bb,
	}
}

// GetInstance gets the instance information stored in the database for a given
// instance ID.
func (c *Core) GetInstance(ctx context.Context, id string) (*Instance, error) {
	instance, err := c.db.GetInstance(id)

	if err == database.ErrInstanceNotFound {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	return &Instance{
		ID:           instance.ID,
		ProviderName: instance.ProviderName,
		Image:        instance.Image,
		State:        instance.State,
		IPAddress:    instance.IPAddress,
	}, nil
}

// CreateInstanceAttributes contains attributes needed to start an instance
type CreateInstanceAttributes struct {
	ImageName    string
	InstanceType string
	PublicSSHKey string
}

//DeleteInstanceAttributes contains attributes needed to delete an instance
type DeleteInstanceAttributes struct {
	InstanceID string
}

// CreateInstance creates an instance in the database and queues off the cloud
// create job in the background.
func (c *Core) CreateInstance(ctx context.Context, providerName string, attr CreateInstanceAttributes) (*Instance, error) {
	id, err := c.db.CreateInstance(database.Instance{
		ProviderName: providerName,
		Image:        attr.ImageName,
		InstanceType: attr.InstanceType,
		PublicSSHKey: attr.PublicSSHKey,
		State:        "creating",
	})
	if err != nil {
		return nil, errors.Wrap(err, "error creating instance in database")
	}

	err = c.bb.Enqueue(background.Job{
		UUID:       uuid.New(),
		Context:    ctx,
		Payload:    []byte(id),
		Queue:      "create",
		MaxRetries: MaxCreateRetries,
	})
	if err != nil {
		// TODO(henrikhodne): Delete the record in the database?
		return nil, errors.Wrap(err, "error enqueueing 'create' job in the background")
	}

	return &Instance{
		ID:           id,
		ProviderName: providerName,
		Image:        attr.ImageName,
		State:        "creating",
	}, nil
}

// RemoveInstance creates an instance in the database and queues off the cloud
// create job in the background.
func (c *Core) RemoveInstance(ctx context.Context, attr DeleteInstanceAttributes) error {
	inst, err := c.db.GetInstance(attr.InstanceID)
	if err != nil {
		return err
	}

	_, err = c.db.RemoveInstance(inst)
	if err != nil {
		return errors.Wrap(err, "error deleting instance in database")
	}

	err = c.bb.Enqueue(background.Job{
		UUID:       uuid.New(),
		Context:    ctx,
		Payload:    []byte(attr.InstanceID),
		Queue:      "remove",
		MaxRetries: MaxCreateRetries,
	})
	if err != nil {
		return errors.Wrap(err, "error enqueueing 'create' job in the background")
	}

	return nil
}

// ProviderCreateInstance is used to schedule the creation of the instance with
// the given ID on the provider selected for that instance.
func (c *Core) ProviderCreateInstance(ctx context.Context, byteID []byte) error {
	id := string(byteID)

	cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"instance_id": id,
	}).Info("creating instance")

	dbInstance, err := c.db.GetInstance(id)
	if err != nil {
		return errors.Wrap(err, "error fetching instance from DB")
	}

	cloudProvider, err := c.cloudProvider(dbInstance.ProviderName)
	if err != nil {
		return errors.Wrapf(err, "couldn't find provider with given name: %v", dbInstance.ProviderName)
	}

	instance, err := cloudProvider.Create(id, cloud.CreateAttributes{
		ImageName:    dbInstance.Image,
		InstanceType: cloud.InstanceType(dbInstance.InstanceType),
		PublicSSHKey: dbInstance.PublicSSHKey,
	})
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
		}).Error("error creating instance")
		dbInstance.State = "errored"

		err = c.db.UpdateInstance(dbInstance)
		if err != nil {
			cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
				"err":         err,
				"instance_id": id,
				"provider_id": instance.ID,
			}).Error("couldn't update instance in DB")
			return err
		}

		return err
	}

	dbInstance.State = "starting"

	err = c.db.UpdateInstance(dbInstance)
	if err != nil {
		return errors.Wrap(err, "couldn't update instance in DB")
	}

	cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"instance_id": id,
		"provider_id": instance.ID,
	}).Info("created instance")

	return nil
}

// ProviderRefresh is used to synchronize the data on all the cloud providers
// with the data in our database.
func (c *Core) ProviderRefresh(ctx context.Context) error {
	c.refreshProviders()
	c.cloudProvidersMutex.Lock()
	defer c.cloudProvidersMutex.Unlock()

	var result error

	for providerName, cloudProvider := range c.cloudProviders {
		instances, err := cloudProvider.List()
		if err != nil {
			result = multierror.Append(result, err)
			continue
		}

		for _, instance := range instances {
			dbInstance, err := c.db.GetInstance(instance.ID)
			if err != nil {
				cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"provider_name": providerName,
					"provider_id":   instance.ID,
				}).Error("failed fetching instance from database")
				continue
			}

			dbInstance.IPAddress = instance.IPAddress
			dbInstance.State = string(instance.State)

			err = c.db.UpdateInstance(dbInstance)
			if err != nil {
				cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"provider":    providerName,
					"provider_id": instance.ID,
					"db_id":       dbInstance.ID,
				}).Error("failed to update instance in database")
			}
		}

		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"provider":       providerName,
			"instance_count": len(instances),
		}).Info("refreshed instances")
	}

	return result
}

// CheckToken is used to check whether a given tokenID+token is in the database.
// Returns (true, nil) iff the token is valid, (false, nil) if the token is
// invalid, and (false, err) if an error occurred while fetching the token.
func (c *Core) CheckToken(tokenID uint64, token string) (bool, error) {
	salt, hash, err := c.db.GetSaltAndHashForTokenID(tokenID)
	if err != nil {
		return false, err
	}

	decodedToken, err := hex.DecodeString(token)
	if err != nil {
		return false, err
	}

	generatedHash, err := scrypt.Key(decodedToken, salt, 16384, 8, 1, 32)
	if err != nil {
		return false, err
	}

	return subtle.ConstantTimeCompare(generatedHash, hash) == 1, nil
}

// cloudProvider is used to get the provider implementation for the cloud
// provider with a given name. Return an error if no cloud provider with the
// given name exists, or if an error occurred refreshing the configuration from
// the database.
func (c *Core) cloudProvider(name string) (cloud.Provider, error) {
	c.cloudProvidersMutex.Lock()
	cloudProvider, ok := c.cloudProviders[name]
	c.cloudProvidersMutex.Unlock()
	if !ok {
		err := c.refreshProviders()
		if err != nil {
			return nil, err
		}

		c.cloudProvidersMutex.Lock()
		cloudProvider, ok = c.cloudProviders[name]
		c.cloudProvidersMutex.Unlock()
		if !ok {
			// This really shouldn't happen, since the database should ensure
			// that a provider with a matching name exists.
			return nil, fmt.Errorf("couldn't find a provider with that name")
		}
	}

	return cloudProvider, nil
}

// refreshProviders is used to regenerate the c.cloudProviders map with the
// configurations stored in the database.
func (c *Core) refreshProviders() error {
	c.cloudProvidersMutex.Lock()
	defer c.cloudProvidersMutex.Unlock()
	dbCloudProviders, err := c.db.ListProviders()
	if err != nil {
		return err
	}

	cloudProviders := make(map[string]cloud.Provider)

	for _, dbCloudProvider := range dbCloudProviders {
		cloudProvider, err := cloud.NewProvider(dbCloudProvider.Type, dbCloudProvider.Config)
		if err != nil {
			return err
		}

		cloudProviders[dbCloudProvider.Name] = cloudProvider
	}

	c.cloudProviders = cloudProviders

	return nil
}

// Instance is a single compute instance.
type Instance struct {
	ID           string
	ProviderName string
	Image        string
	State        string
	IPAddress    string
}
<commit_msg>wrap rerors in cloudbrain core<commit_after>package cloudbrain

import (
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/hashicorp/go-multierror"
	"github.com/pborman/uuid"
	"github.com/pkg/errors"
	"github.com/travis-ci/cloud-brain/background"
	"github.com/travis-ci/cloud-brain/cbcontext"
	"github.com/travis-ci/cloud-brain/cloud"
	"github.com/travis-ci/cloud-brain/database"
	"golang.org/x/crypto/scrypt"
	"golang.org/x/net/context"
	"gopkg.in/urfave/cli.v2"
)

var (
	//VersionString gets set during `make`
	VersionString = "?"
	//RevisionString gets set during `make`
	RevisionString = "?"
	//RevisionURLString gets set during `make`
	RevisionURLString = "?"
	//GeneratedString gets set during `make`
	GeneratedString = "?"
	//CopyrightString gets set during `make`
	CopyrightString = "?"
)

func init() {
	cli.VersionPrinter = customVersionPrinter
	_ = os.Setenv("VERSION", VersionString)
	_ = os.Setenv("REVISION", RevisionString)
	_ = os.Setenv("GENERATED", GeneratedString)
}

func customVersionPrinter(c *cli.Context) {
	fmt.Printf("%v v=%v rev=%v d=%v\n", filepath.Base(c.App.Name),
		VersionString, RevisionString, GeneratedString)
}

// MaxCreateRetries is the number of times the "create" job will be retried.
const MaxCreateRetries = 10

// Core is used as a central manager for all Cloud Brain functionality. The HTTP
// API and the background workers are just frontends for the Core, and calls
// methods on Core for functionality.
type Core struct {
	db database.DB
	bb background.Backend

	cloudProvidersMutex sync.Mutex
	cloudProviders      map[string]cloud.Provider
}

// NewCore is used to create a new Core backed by the given database and
// background Backend.
func NewCore(db database.DB, bb background.Backend) *Core {
	return &Core{
		db: db,
		bb: bb,
	}
}

// GetInstance gets the instance information stored in the database for a given
// instance ID.
func (c *Core) GetInstance(ctx context.Context, id string) (*Instance, error) {
	instance, err := c.db.GetInstance(id)

	if err == database.ErrInstanceNotFound {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	return &Instance{
		ID:           instance.ID,
		ProviderName: instance.ProviderName,
		Image:        instance.Image,
		State:        instance.State,
		IPAddress:    instance.IPAddress,
	}, nil
}

// CreateInstanceAttributes contains attributes needed to start an instance
type CreateInstanceAttributes struct {
	ImageName    string
	InstanceType string
	PublicSSHKey string
}

//DeleteInstanceAttributes contains attributes needed to delete an instance
type DeleteInstanceAttributes struct {
	InstanceID string
}

// CreateInstance creates an instance in the database and queues off the cloud
// create job in the background.
func (c *Core) CreateInstance(ctx context.Context, providerName string, attr CreateInstanceAttributes) (*Instance, error) {
	id, err := c.db.CreateInstance(database.Instance{
		ProviderName: providerName,
		Image:        attr.ImageName,
		InstanceType: attr.InstanceType,
		PublicSSHKey: attr.PublicSSHKey,
		State:        "creating",
	})
	if err != nil {
		return nil, errors.Wrap(err, "error creating instance in database")
	}

	err = c.bb.Enqueue(background.Job{
		UUID:       uuid.New(),
		Context:    ctx,
		Payload:    []byte(id),
		Queue:      "create",
		MaxRetries: MaxCreateRetries,
	})
	if err != nil {
		// TODO(henrikhodne): Delete the record in the database?
		return nil, errors.Wrap(err, "error enqueueing 'create' job in the background")
	}

	return &Instance{
		ID:           id,
		ProviderName: providerName,
		Image:        attr.ImageName,
		State:        "creating",
	}, nil
}

// RemoveInstance creates an instance in the database and queues off the cloud
// create job in the background.
func (c *Core) RemoveInstance(ctx context.Context, attr DeleteInstanceAttributes) error {
	inst, err := c.db.GetInstance(attr.InstanceID)
	if err != nil {
		return errors.Wrap(err, "error fetching instance from DB")
	}

	_, err = c.db.RemoveInstance(inst)
	if err != nil {
		return errors.Wrap(err, "error deleting instance in database")
	}

	err = c.bb.Enqueue(background.Job{
		UUID:       uuid.New(),
		Context:    ctx,
		Payload:    []byte(attr.InstanceID),
		Queue:      "remove",
		MaxRetries: MaxCreateRetries,
	})
	if err != nil {
		return errors.Wrap(err, "error enqueueing 'remove' job in the background")
	}

	return nil
}

// ProviderCreateInstance is used to schedule the creation of the instance with
// the given ID on the provider selected for that instance.
func (c *Core) ProviderCreateInstance(ctx context.Context, byteID []byte) error {
	id := string(byteID)

	cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"instance_id": id,
	}).Info("creating instance")

	dbInstance, err := c.db.GetInstance(id)
	if err != nil {
		return errors.Wrap(err, "error fetching instance from DB")
	}

	cloudProvider, err := c.cloudProvider(dbInstance.ProviderName)
	if err != nil {
		return errors.Wrapf(err, "couldn't find provider with given name: %v", dbInstance.ProviderName)
	}

	instance, err := cloudProvider.Create(id, cloud.CreateAttributes{
		ImageName:    dbInstance.Image,
		InstanceType: cloud.InstanceType(dbInstance.InstanceType),
		PublicSSHKey: dbInstance.PublicSSHKey,
	})
	if err != nil {
		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"err":         err,
			"instance_id": id,
		}).Error("error creating instance")
		dbInstance.State = "errored"

		err = c.db.UpdateInstance(dbInstance)
		if err != nil {
			cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
				"err":         err,
				"instance_id": id,
				"provider_id": instance.ID,
			}).Error("couldn't update instance in DB")
			return err
		}

		return err
	}

	dbInstance.State = "starting"

	err = c.db.UpdateInstance(dbInstance)
	if err != nil {
		return errors.Wrap(err, "couldn't update instance in DB")
	}

	cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
		"instance_id": id,
		"provider_id": instance.ID,
	}).Info("created instance")

	return nil
}

// ProviderRefresh is used to synchronize the data on all the cloud providers
// with the data in our database.
func (c *Core) ProviderRefresh(ctx context.Context) error {
	c.refreshProviders()
	c.cloudProvidersMutex.Lock()
	defer c.cloudProvidersMutex.Unlock()

	var result error

	for providerName, cloudProvider := range c.cloudProviders {
		instances, err := cloudProvider.List()
		if err != nil {
			result = multierror.Append(result, err)
			continue
		}

		for _, instance := range instances {
			dbInstance, err := c.db.GetInstance(instance.ID)
			if err != nil {
				cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"provider_name": providerName,
					"provider_id":   instance.ID,
				}).Error("failed fetching instance from database")
				continue
			}

			dbInstance.IPAddress = instance.IPAddress
			dbInstance.State = string(instance.State)

			err = c.db.UpdateInstance(dbInstance)
			if err != nil {
				cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
					"provider":    providerName,
					"provider_id": instance.ID,
					"db_id":       dbInstance.ID,
				}).Error("failed to update instance in database")
			}
		}

		cbcontext.LoggerFromContext(ctx).WithFields(logrus.Fields{
			"provider":       providerName,
			"instance_count": len(instances),
		}).Info("refreshed instances")
	}

	return result
}

// CheckToken is used to check whether a given tokenID+token is in the database.
// Returns (true, nil) iff the token is valid, (false, nil) if the token is
// invalid, and (false, err) if an error occurred while fetching the token.
func (c *Core) CheckToken(tokenID uint64, token string) (bool, error) {
	salt, hash, err := c.db.GetSaltAndHashForTokenID(tokenID)
	if err != nil {
		return false, err
	}

	decodedToken, err := hex.DecodeString(token)
	if err != nil {
		return false, err
	}

	generatedHash, err := scrypt.Key(decodedToken, salt, 16384, 8, 1, 32)
	if err != nil {
		return false, err
	}

	return subtle.ConstantTimeCompare(generatedHash, hash) == 1, nil
}

// cloudProvider is used to get the provider implementation for the cloud
// provider with a given name. Return an error if no cloud provider with the
// given name exists, or if an error occurred refreshing the configuration from
// the database.
func (c *Core) cloudProvider(name string) (cloud.Provider, error) {
	c.cloudProvidersMutex.Lock()
	cloudProvider, ok := c.cloudProviders[name]
	c.cloudProvidersMutex.Unlock()
	if !ok {
		err := c.refreshProviders()
		if err != nil {
			return nil, err
		}

		c.cloudProvidersMutex.Lock()
		cloudProvider, ok = c.cloudProviders[name]
		c.cloudProvidersMutex.Unlock()
		if !ok {
			// This really shouldn't happen, since the database should ensure
			// that a provider with a matching name exists.
			return nil, fmt.Errorf("couldn't find a provider with that name")
		}
	}

	return cloudProvider, nil
}

// refreshProviders is used to regenerate the c.cloudProviders map with the
// configurations stored in the database.
func (c *Core) refreshProviders() error {
	c.cloudProvidersMutex.Lock()
	defer c.cloudProvidersMutex.Unlock()
	dbCloudProviders, err := c.db.ListProviders()
	if err != nil {
		return err
	}

	cloudProviders := make(map[string]cloud.Provider)

	for _, dbCloudProvider := range dbCloudProviders {
		cloudProvider, err := cloud.NewProvider(dbCloudProvider.Type, dbCloudProvider.Config)
		if err != nil {
			return err
		}

		cloudProviders[dbCloudProvider.Name] = cloudProvider
	}

	c.cloudProviders = cloudProviders

	return nil
}

// Instance is a single compute instance.
type Instance struct {
	ID           string
	ProviderName string
	Image        string
	State        string
	IPAddress    string
}
<|endoftext|>


<commit_before>package gorethink

import (
	"crypto/tls"
	"sync"
	"time"

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

// A Session represents a connection to a RethinkDB cluster and should be used
// when executing queries.
type Session struct {
	hosts []Host
	opts  *ConnectOpts

	mu      sync.RWMutex
	cluster *Cluster
	closed  bool
}

// ConnectOpts is used to specify optional arguments when connecting to a cluster.
type ConnectOpts struct {
	// Address holds the address of the server initially used when creating the
	// session. Only used if Addresses is empty
	Address string `gorethink:"address,omitempty"`
	// Addresses holds the addresses of the servers initially used when creating
	// the session.
	Addresses []string `gorethink:"addresses,omitempty"`
	// Database is the default database name used when executing queries, this
	// value is only used if the query does not contain any DB term
	Database string `gorethink:"database,omitempty"`
	// Username holds the username used for authentication, if blank (and the v1
	// handshake protocol is being used) then the admin user is used
	Username string `gorethink:"username,omitempty"`
	// Password holds the password used for authentication (only used when using
	// the v1 handshake protocol)
	Password string `gorethink:"password,omitempty"`
	// AuthKey is used for authentication when using the v0.4 handshake protocol
	// This field is no deprecated
	AuthKey string `gorethink:"authkey,omitempty"`
	// Timeout is the time the driver waits when creating new connections, to
	// configure the timeout used when executing queries use WriteTimeout and
	// ReadTimeout
	Timeout time.Duration `gorethink:"timeout,omitempty"`
	// WriteTimeout is the amount of time the driver will wait when sending the
	// query to the server
	WriteTimeout time.Duration `gorethink:"write_timeout,omitempty"`
	// ReadTimeout is the amount of time the driver will wait for a response from
	// the server when executing queries.
	ReadTimeout time.Duration `gorethink:"read_timeout,omitempty"`
	// KeepAlivePeriod is the keep alive period used by the connection, by default
	// this is 30s. It is not possible to disable keep alive messages
	KeepAlivePeriod time.Duration `gorethink:"keep_alive_timeout,omitempty"`
	// TLSConfig holds the TLS configuration and can be used when connecting
	// to a RethinkDB server protected by SSL
	TLSConfig *tls.Config `gorethink:"tlsconfig,omitempty"`
	// HandshakeVersion is used to specify which handshake version should be
	// used, this currently defaults to v1 which is used by RethinkDB 2.3 and
	// later. If you are using an older version then you can set the handshake
	// version to 0.4
	HandshakeVersion HandshakeVersion `gorethink:"handshake_version,omitempty"`
	// UseJSONNumber indicates whether the cursors running in this session should
	// use json.Number instead of float64 while unmarshaling documents with
	// interface{}. The default is `false`.
	UseJSONNumber bool
	// NumRetries is the number of times a query is retried if a connection
	// error is detected, queries are not retried if RethinkDB returns a
	// runtime error.
	NumRetries int

	// InitialCap is used by the internal connection pool and is used to
	// configure how many connections are created for each host when the
	// session is created. If zero then no connections are created until
	// the first query is executed.
	InitialCap int `gorethink:"initial_cap,omitempty"`
	// MaxOpen is used by the internal connection pool and is used to configure
	// the maximum number of connections held in the pool. If all available
	// connections are being used then the driver will open new connections as
	// needed however they will not be returned to the pool. By default the
	// maximum number of connections is 2
	MaxOpen int `gorethink:"max_open,omitempty"`

	// Below options are for cluster discovery, please note there is a high
	// probability of these changing as the API is still being worked on.

	// DiscoverHosts is used to enable host discovery, when true the driver
	// will attempt to discover any new nodes added to the cluster and then
	// start sending queries to these new nodes.
	DiscoverHosts bool `gorethink:"discover_hosts,omitempty"`
	// HostDecayDuration is used by the go-hostpool package to calculate a weighted
	// score when selecting a host. By default a value of 5 minutes is used.
	HostDecayDuration time.Duration

	// Deprecated: This function is no longer used due to changes in the
	// way hosts are selected.
	NodeRefreshInterval time.Duration `gorethink:"node_refresh_interval,omitempty"`
	// Deprecated: Use InitialCap instead
	MaxIdle int `gorethink:"max_idle,omitempty"`
}

func (o ConnectOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Connect creates a new database session. To view the available connection
// options see ConnectOpts.
//
// By default maxIdle and maxOpen are set to 1: passing values greater
// than the default (e.g. MaxIdle: "10", MaxOpen: "20") will provide a
// pool of re-usable connections.
//
// Basic connection example:
//
// 	session, err := r.Connect(r.ConnectOpts{
// 		Host: "localhost:28015",
// 		Database: "test",
// 		AuthKey:  "14daak1cad13dj",
// 	})
//
// Cluster connection example:
//
// 	session, err := r.Connect(r.ConnectOpts{
// 		Hosts: []string{"localhost:28015", "localhost:28016"},
// 		Database: "test",
// 		AuthKey:  "14daak1cad13dj",
// 	})
func Connect(opts ConnectOpts) (*Session, error) {
	var addresses = opts.Addresses
	if len(addresses) == 0 {
		addresses = []string{opts.Address}
	}

	hosts := make([]Host, len(addresses))
	for i, address := range addresses {
		hostname, port := splitAddress(address)
		hosts[i] = NewHost(hostname, port)
	}
	if len(hosts) <= 0 {
		return nil, ErrNoHosts
	}

	// Connect
	s := &Session{
		hosts: hosts,
		opts:  &opts,
	}

	err := s.Reconnect()
	if err != nil {
		// note: s.Reconnect() will initialize cluster information which
		// will cause the .IsConnected() method to be caught in a loop
		return &Session{
			hosts: hosts,
			opts: &opts,
		}, err
	}

	return s, nil
}

// CloseOpts allows calls to the Close function to be configured.
type CloseOpts struct {
	NoReplyWait bool `gorethink:"noreplyWait,omitempty"`
}

func (o CloseOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// IsConnected returns true if session has a valid connection.
func (s *Session) IsConnected() bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.cluster == nil || s.closed {
		return false
	}
	return s.cluster.IsConnected()
}

// Reconnect closes and re-opens a session.
func (s *Session) Reconnect(optArgs ...CloseOpts) error {
	var err error

	if err = s.Close(optArgs...); err != nil {
		return err
	}

	s.mu.Lock()
	s.cluster, err = NewCluster(s.hosts, s.opts)
	if err != nil {
		return err
	}

	s.closed = false
	s.mu.Unlock()

	return nil
}

// Close closes the session
func (s *Session) Close(optArgs ...CloseOpts) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return nil
	}

	if len(optArgs) >= 1 {
		if optArgs[0].NoReplyWait {
			s.mu.Unlock()
			s.NoReplyWait()
			s.mu.Lock()
		}
	}

	if s.cluster != nil {
		s.cluster.Close()
	}
	s.cluster = nil
	s.closed = true

	return nil
}

// SetInitialPoolCap sets the initial capacity of the connection pool.
func (s *Session) SetInitialPoolCap(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.InitialCap = n
	s.cluster.SetInitialPoolCap(n)
}

// SetMaxIdleConns sets the maximum number of connections in the idle
// connection pool.
func (s *Session) SetMaxIdleConns(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.MaxIdle = n
	s.cluster.SetMaxIdleConns(n)
}

// SetMaxOpenConns sets the maximum number of open connections to the database.
func (s *Session) SetMaxOpenConns(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.MaxOpen = n
	s.cluster.SetMaxOpenConns(n)
}

// NoReplyWait ensures that previous queries with the noreply flag have been
// processed by the server. Note that this guarantee only applies to queries
// run on the given connection
func (s *Session) NoReplyWait() error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return ErrConnectionClosed
	}

	return s.cluster.Exec(Query{
		Type: p.Query_NOREPLY_WAIT,
	})
}

// Use changes the default database used
func (s *Session) Use(database string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.Database = database
}

// Database returns the selected database set by Use
func (s *Session) Database() string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.opts.Database
}

// Query executes a ReQL query using the session to connect to the database
func (s *Session) Query(q Query) (*Cursor, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrConnectionClosed
	}

	return s.cluster.Query(q)
}

// Exec executes a ReQL query using the session to connect to the database
func (s *Session) Exec(q Query) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return ErrConnectionClosed
	}

	return s.cluster.Exec(q)
}

// Server returns the server name and server UUID being used by a connection.
func (s *Session) Server() (ServerResponse, error) {
	return s.cluster.Server()
}

// SetHosts resets the hosts used when connecting to the RethinkDB cluster
func (s *Session) SetHosts(hosts []Host) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.hosts = hosts
}

func (s *Session) newQuery(t Term, opts map[string]interface{}) (Query, error) {
	return newQuery(t, opts, s.opts)
}
<commit_msg>Fix for possible deadlock<commit_after>package gorethink

import (
	"crypto/tls"
	"sync"
	"time"

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

// A Session represents a connection to a RethinkDB cluster and should be used
// when executing queries.
type Session struct {
	hosts []Host
	opts  *ConnectOpts

	mu      sync.RWMutex
	cluster *Cluster
	closed  bool
}

// ConnectOpts is used to specify optional arguments when connecting to a cluster.
type ConnectOpts struct {
	// Address holds the address of the server initially used when creating the
	// session. Only used if Addresses is empty
	Address string `gorethink:"address,omitempty"`
	// Addresses holds the addresses of the servers initially used when creating
	// the session.
	Addresses []string `gorethink:"addresses,omitempty"`
	// Database is the default database name used when executing queries, this
	// value is only used if the query does not contain any DB term
	Database string `gorethink:"database,omitempty"`
	// Username holds the username used for authentication, if blank (and the v1
	// handshake protocol is being used) then the admin user is used
	Username string `gorethink:"username,omitempty"`
	// Password holds the password used for authentication (only used when using
	// the v1 handshake protocol)
	Password string `gorethink:"password,omitempty"`
	// AuthKey is used for authentication when using the v0.4 handshake protocol
	// This field is no deprecated
	AuthKey string `gorethink:"authkey,omitempty"`
	// Timeout is the time the driver waits when creating new connections, to
	// configure the timeout used when executing queries use WriteTimeout and
	// ReadTimeout
	Timeout time.Duration `gorethink:"timeout,omitempty"`
	// WriteTimeout is the amount of time the driver will wait when sending the
	// query to the server
	WriteTimeout time.Duration `gorethink:"write_timeout,omitempty"`
	// ReadTimeout is the amount of time the driver will wait for a response from
	// the server when executing queries.
	ReadTimeout time.Duration `gorethink:"read_timeout,omitempty"`
	// KeepAlivePeriod is the keep alive period used by the connection, by default
	// this is 30s. It is not possible to disable keep alive messages
	KeepAlivePeriod time.Duration `gorethink:"keep_alive_timeout,omitempty"`
	// TLSConfig holds the TLS configuration and can be used when connecting
	// to a RethinkDB server protected by SSL
	TLSConfig *tls.Config `gorethink:"tlsconfig,omitempty"`
	// HandshakeVersion is used to specify which handshake version should be
	// used, this currently defaults to v1 which is used by RethinkDB 2.3 and
	// later. If you are using an older version then you can set the handshake
	// version to 0.4
	HandshakeVersion HandshakeVersion `gorethink:"handshake_version,omitempty"`
	// UseJSONNumber indicates whether the cursors running in this session should
	// use json.Number instead of float64 while unmarshaling documents with
	// interface{}. The default is `false`.
	UseJSONNumber bool
	// NumRetries is the number of times a query is retried if a connection
	// error is detected, queries are not retried if RethinkDB returns a
	// runtime error.
	NumRetries int

	// InitialCap is used by the internal connection pool and is used to
	// configure how many connections are created for each host when the
	// session is created. If zero then no connections are created until
	// the first query is executed.
	InitialCap int `gorethink:"initial_cap,omitempty"`
	// MaxOpen is used by the internal connection pool and is used to configure
	// the maximum number of connections held in the pool. If all available
	// connections are being used then the driver will open new connections as
	// needed however they will not be returned to the pool. By default the
	// maximum number of connections is 2
	MaxOpen int `gorethink:"max_open,omitempty"`

	// Below options are for cluster discovery, please note there is a high
	// probability of these changing as the API is still being worked on.

	// DiscoverHosts is used to enable host discovery, when true the driver
	// will attempt to discover any new nodes added to the cluster and then
	// start sending queries to these new nodes.
	DiscoverHosts bool `gorethink:"discover_hosts,omitempty"`
	// HostDecayDuration is used by the go-hostpool package to calculate a weighted
	// score when selecting a host. By default a value of 5 minutes is used.
	HostDecayDuration time.Duration

	// Deprecated: This function is no longer used due to changes in the
	// way hosts are selected.
	NodeRefreshInterval time.Duration `gorethink:"node_refresh_interval,omitempty"`
	// Deprecated: Use InitialCap instead
	MaxIdle int `gorethink:"max_idle,omitempty"`
}

func (o ConnectOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Connect creates a new database session. To view the available connection
// options see ConnectOpts.
//
// By default maxIdle and maxOpen are set to 1: passing values greater
// than the default (e.g. MaxIdle: "10", MaxOpen: "20") will provide a
// pool of re-usable connections.
//
// Basic connection example:
//
// 	session, err := r.Connect(r.ConnectOpts{
// 		Host: "localhost:28015",
// 		Database: "test",
// 		AuthKey:  "14daak1cad13dj",
// 	})
//
// Cluster connection example:
//
// 	session, err := r.Connect(r.ConnectOpts{
// 		Hosts: []string{"localhost:28015", "localhost:28016"},
// 		Database: "test",
// 		AuthKey:  "14daak1cad13dj",
// 	})
func Connect(opts ConnectOpts) (*Session, error) {
	var addresses = opts.Addresses
	if len(addresses) == 0 {
		addresses = []string{opts.Address}
	}

	hosts := make([]Host, len(addresses))
	for i, address := range addresses {
		hostname, port := splitAddress(address)
		hosts[i] = NewHost(hostname, port)
	}
	if len(hosts) <= 0 {
		return nil, ErrNoHosts
	}

	// Connect
	s := &Session{
		hosts: hosts,
		opts:  &opts,
	}

	err := s.Reconnect()
	if err != nil {
		// note: s.Reconnect() will initialize cluster information which
		// will cause the .IsConnected() method to be caught in a loop
		return &Session{
			hosts: hosts,
			opts: &opts,
		}, err
	}

	return s, nil
}

// CloseOpts allows calls to the Close function to be configured.
type CloseOpts struct {
	NoReplyWait bool `gorethink:"noreplyWait,omitempty"`
}

func (o CloseOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// IsConnected returns true if session has a valid connection.
func (s *Session) IsConnected() bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.cluster == nil || s.closed {
		return false
	}
	return s.cluster.IsConnected()
}

// Reconnect closes and re-opens a session.
func (s *Session) Reconnect(optArgs ...CloseOpts) error {
	var err error

	if err = s.Close(optArgs...); err != nil {
		return err
	}

	s.mu.Lock()
	s.cluster, err = NewCluster(s.hosts, s.opts)
	if err != nil {
		s.mu.Unlock()
		return err
	}

	s.closed = false
	s.mu.Unlock()

	return nil
}

// Close closes the session
func (s *Session) Close(optArgs ...CloseOpts) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return nil
	}

	if len(optArgs) >= 1 {
		if optArgs[0].NoReplyWait {
			s.mu.Unlock()
			s.NoReplyWait()
			s.mu.Lock()
		}
	}

	if s.cluster != nil {
		s.cluster.Close()
	}
	s.cluster = nil
	s.closed = true

	return nil
}

// SetInitialPoolCap sets the initial capacity of the connection pool.
func (s *Session) SetInitialPoolCap(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.InitialCap = n
	s.cluster.SetInitialPoolCap(n)
}

// SetMaxIdleConns sets the maximum number of connections in the idle
// connection pool.
func (s *Session) SetMaxIdleConns(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.MaxIdle = n
	s.cluster.SetMaxIdleConns(n)
}

// SetMaxOpenConns sets the maximum number of open connections to the database.
func (s *Session) SetMaxOpenConns(n int) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.MaxOpen = n
	s.cluster.SetMaxOpenConns(n)
}

// NoReplyWait ensures that previous queries with the noreply flag have been
// processed by the server. Note that this guarantee only applies to queries
// run on the given connection
func (s *Session) NoReplyWait() error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return ErrConnectionClosed
	}

	return s.cluster.Exec(Query{
		Type: p.Query_NOREPLY_WAIT,
	})
}

// Use changes the default database used
func (s *Session) Use(database string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.opts.Database = database
}

// Database returns the selected database set by Use
func (s *Session) Database() string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.opts.Database
}

// Query executes a ReQL query using the session to connect to the database
func (s *Session) Query(q Query) (*Cursor, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return nil, ErrConnectionClosed
	}

	return s.cluster.Query(q)
}

// Exec executes a ReQL query using the session to connect to the database
func (s *Session) Exec(q Query) error {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.closed {
		return ErrConnectionClosed
	}

	return s.cluster.Exec(q)
}

// Server returns the server name and server UUID being used by a connection.
func (s *Session) Server() (ServerResponse, error) {
	return s.cluster.Server()
}

// SetHosts resets the hosts used when connecting to the RethinkDB cluster
func (s *Session) SetHosts(hosts []Host) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.hosts = hosts
}

func (s *Session) newQuery(t Term, opts map[string]interface{}) (Query, error) {
	return newQuery(t, opts, s.opts)
}
<|endoftext|>


<commit_before>package req

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"time"
)

var defaultCookieJar http.CookieJar

func init() {
	defaultCookieJar, _ = cookiejar.New(nil)
}

type setting struct {
	Timeout             time.Duration // total timeout
	DialTimeout         time.Duration
	ReadTimeout         time.Duration
	WriteTimeout        time.Duration
	TLSHandshakeTimeout time.Duration
	InsecureTLS         bool
	Jar                 http.CookieJar
	Proxy               func(*http.Request) (*url.URL, error)
	TLSClientConfig     *tls.Config
	Transport           *http.Transport
	Client              *http.Client
}

func (r *Request) prepareSetting() bool {
	if r == nil {
		return false
	}
	if r.setting == nil {
		r.setting = &setting{}
	}
	return true
}

// GetClient returns the *http.Client according to the setting.
func (r *Request) GetClient() *http.Client {
	if !r.prepareSetting() {
		return http.DefaultClient
	}
	s := r.setting
	if s.Client == nil {
		c := &http.Client{
			Transport: r.GetTransport(),
		}
		if s.Jar != nil {
			c.Jar = s.Jar
		}
		if s.Timeout > 0 {
			c.Timeout = s.Timeout
		}
		s.Client = c
	}
	return s.Client
}

func (r *Request) createTransport() *http.Transport {
	s := r.setting
	trans := &http.Transport{}
	trans.Dial = func(network, address string) (conn net.Conn, err error) {
		if s.DialTimeout > 0 {
			conn, err = net.DialTimeout(network, address, s.DialTimeout)
			if err != nil {
				return
			}
		} else {
			conn, err = net.Dial(network, address)
			if err != nil {
				return
			}
		}
		if s.ReadTimeout > 0 {
			conn.SetReadDeadline(time.Now().Add(s.ReadTimeout))
		}
		if s.WriteTimeout > 0 {
			conn.SetWriteDeadline(time.Now().Add(s.WriteTimeout))
		}
		return
	}
	if s.TLSClientConfig != nil {
		trans.TLSClientConfig = s.TLSClientConfig
	}
	if s.TLSHandshakeTimeout > 0 {
		trans.TLSHandshakeTimeout = s.TLSHandshakeTimeout
	}
	if s.InsecureTLS {
		if trans.TLSClientConfig != nil {
			trans.TLSClientConfig.InsecureSkipVerify = true
		} else {
			trans.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
		}
	}
	if s.Proxy != nil {
		trans.Proxy = s.Proxy
	} else {
		trans.Proxy = http.ProxyFromEnvironment
	}
	return trans
}

// GetTransport return the http.Transport according to the setting.
func (r *Request) GetTransport() *http.Transport {
	if !r.prepareSetting() {
		trans, _ := http.DefaultTransport.(*http.Transport)
		return trans
	}
	s := r.setting
	if s.Transport == nil {
		s.Transport = r.createTransport()
	}
	return s.Transport
}

// Client set the http.Client for the request
func (r *Request) Client(client *http.Client) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Client = client
	return r
}

// Transport set the http.Transport for request
func (r *Request) Transport(trans *http.Transport) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Transport = trans
	r.setting.Client = nil
	return r
}

// TLSClientConfig set the tls config to the request
func (r *Request) TLSClientConfig(config *tls.Config) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.TLSClientConfig = config
	r.setting.Transport = nil
	r.setting.Client = nil
	return r
}

// Proxy set the proxy for request
func (r *Request) Proxy(proxy func(*http.Request) (*url.URL, error)) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Proxy = proxy
	r.setting.Client = nil
	r.setting.Transport = nil
	return r
}

// Timeout set the total timeout for request,
// once timeout reached, request will be canceled.
func (r *Request) Timeout(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Timeout = d
	r.setting.Client = nil
	r.setting.Transport = nil
	return r
}

// TimeoutDial sets the timeout for dial connection.
func (r *Request) TimeoutDial(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.DialTimeout = d
	r.setting.Client = nil
	r.setting.Transport = nil
	return r
}

// TimeoutRead sets the timeout for read operation.
func (r *Request) TimeoutRead(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.ReadTimeout = d
	r.setting.Client = nil
	r.setting.Transport = nil
	return r
}

// TimeoutWrite sets the timeout for write operation.
func (r *Request) TimeoutWrite(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.WriteTimeout = d
	r.setting.Client = nil
	r.setting.Transport = nil
	return r
}

// TimeoutTLSHandshake specifies the maximum amount of time waiting to
// wait for a TLS handshake. Zero means no timeout.
func (r *Request) TimeoutTLSHandshake(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.TLSHandshakeTimeout = d
	r.setting.Client = nil
	r.setting.Transport = nil
	return r
}

// InsecureTLS allows to access the insecure https server.
func (r *Request) InsecureTLS(ins bool) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.InsecureTLS = ins
	r.setting.Client = nil
	r.setting.Transport = nil
	return r
}

// EnableCookie set the default CookieJar to the request if enable==true, otherwise set to nil.
func (r *Request) EnableCookie(enable bool) *Request {
	if !r.prepareSetting() {
		return nil
	}
	if enable {
		r.setting.Jar = defaultCookieJar
	} else {
		r.setting.Jar = nil
	}
	r.setting.Client = nil
	return r
}

// EnableCookieWithJar set the specified *http.CookieJar to the request.
func (r *Request) EnableCookieWithJar(jar http.CookieJar) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Jar = jar
	r.setting.Client = nil
	return r
}

// Merge clone some properties of another Request into this.
func (r *Request) Merge(rr *Request) *Request {
	if r == nil || rr == nil {
		return nil
	}
	if len(rr.params) > 0 { // merge params
		for name, value := range rr.params {
			if _, ok := r.params[name]; !ok {
				r.params[name] = value
			}
		}
	}
	if rr.req != nil { // merge internal http.Request
		if r.req == nil {
			r.req = basicRequest()
		}
		if r.req.Method == "" && rr.req.Method != "" {
			r.req.Method = rr.req.Method
		}
		if r.req.Host == "" && rr.req.Host != "" {
			r.req.Host = rr.req.Host
		}
		if r.req.Proto != rr.req.Proto {
			r.req.Proto = rr.req.Proto
			r.req.ProtoMajor = rr.req.ProtoMajor
			r.req.ProtoMinor = rr.req.ProtoMinor
		}
		for name, value := range rr.req.Header {
			if _, ok := r.req.Header[name]; !ok {
				r.req.Header[name] = value
			}
		}
	}
	if rr.setting != nil { // merge setting
		rr.GetClient() // ensure client has been created. prevent creating client every request.
		s := *rr.setting
		r.setting = &s
	}

	return r
}
<commit_msg>disable keep alive to default.<commit_after>package req

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"time"
)

var defaultCookieJar http.CookieJar

func init() {
	defaultCookieJar, _ = cookiejar.New(nil)
}

type setting struct {
	Timeout             time.Duration // total timeout
	DialTimeout         time.Duration
	ReadTimeout         time.Duration
	WriteTimeout        time.Duration
	TLSHandshakeTimeout time.Duration
	InsecureTLS         bool
	Jar                 http.CookieJar
	Proxy               func(*http.Request) (*url.URL, error)
	TLSClientConfig     *tls.Config
	Transport           *http.Transport
	Client              *http.Client
	EnableKeepAlive     bool
}

func (r *Request) prepareSetting() bool {
	if r == nil {
		return false
	}
	if r.setting == nil {
		r.setting = &setting{}
	}
	return true
}

// GetClient returns the *http.Client according to the setting.
func (r *Request) GetClient() *http.Client {
	if !r.prepareSetting() {
		return http.DefaultClient
	}
	s := r.setting
	if s.Client == nil {
		c := &http.Client{
			Transport: r.GetTransport(),
		}
		if s.Jar != nil {
			c.Jar = s.Jar
		}
		if s.Timeout > 0 {
			c.Timeout = s.Timeout
		}
		s.Client = c
	}
	return s.Client
}

func (r *Request) createTransport() *http.Transport {
	s := r.setting
	trans := &http.Transport{}
	if !s.EnableKeepAlive {
		trans.DisableKeepAlives = true
	}
	trans.Dial = func(network, address string) (conn net.Conn, err error) {
		if s.DialTimeout > 0 {
			conn, err = net.DialTimeout(network, address, s.DialTimeout)
			if err != nil {
				return
			}
		} else {
			conn, err = net.Dial(network, address)
			if err != nil {
				return
			}
		}
		if s.ReadTimeout > 0 {
			conn.SetReadDeadline(time.Now().Add(s.ReadTimeout))
		}
		if s.WriteTimeout > 0 {
			conn.SetWriteDeadline(time.Now().Add(s.WriteTimeout))
		}
		return
	}
	if s.TLSClientConfig != nil {
		trans.TLSClientConfig = s.TLSClientConfig
	}
	if s.TLSHandshakeTimeout > 0 {
		trans.TLSHandshakeTimeout = s.TLSHandshakeTimeout
	}
	if s.InsecureTLS {
		if trans.TLSClientConfig != nil {
			trans.TLSClientConfig.InsecureSkipVerify = true
		} else {
			trans.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
		}
	}
	if s.Proxy != nil {
		trans.Proxy = s.Proxy
	} else {
		trans.Proxy = http.ProxyFromEnvironment
	}
	return trans
}

// GetTransport return the http.Transport according to the setting.
func (r *Request) GetTransport() *http.Transport {
	if !r.prepareSetting() {
		trans, _ := http.DefaultTransport.(*http.Transport)
		return trans
	}
	s := r.setting
	if s.Transport == nil {
		s.Transport = r.createTransport()
	}
	return s.Transport
}

// Client set the http.Client for the request
func (r *Request) Client(client *http.Client) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Client = client
	return r
}

// Transport set the http.Transport for request
func (r *Request) Transport(trans *http.Transport) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Transport = trans
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// TLSClientConfig set the tls config to the request
func (r *Request) TLSClientConfig(config *tls.Config) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.TLSClientConfig = config
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// Proxy set the proxy for request
func (r *Request) Proxy(proxy func(*http.Request) (*url.URL, error)) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Proxy = proxy
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// Timeout set the total timeout for request,
// once timeout reached, request will be canceled.
func (r *Request) Timeout(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Timeout = d
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// TimeoutDial sets the timeout for dial connection.
func (r *Request) TimeoutDial(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.DialTimeout = d
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// TimeoutRead sets the timeout for read operation.
func (r *Request) TimeoutRead(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.ReadTimeout = d
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// TimeoutWrite sets the timeout for write operation.
func (r *Request) TimeoutWrite(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.WriteTimeout = d
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// TimeoutTLSHandshake specifies the maximum amount of time waiting to
// wait for a TLS handshake. Zero means no timeout.
func (r *Request) TimeoutTLSHandshake(d time.Duration) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.TLSHandshakeTimeout = d
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// InsecureTLS allows to access the insecure https server.
func (r *Request) InsecureTLS(ins bool) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.InsecureTLS = ins
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// EnableKeepAlive set the default CookieJar to the request if enable==true, otherwise set to nil.
func (r *Request) EnableKeepAlive(enable bool) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.EnableKeepAlive = enable
	if r.setting.Transport != nil {
		r.setting.Transport = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// EnableCookie set the default CookieJar to the request if enable==true, otherwise set to nil.
func (r *Request) EnableCookie(enable bool) *Request {
	if !r.prepareSetting() {
		return nil
	}
	if enable {
		r.setting.Jar = defaultCookieJar
	} else {
		r.setting.Jar = nil
	}
	if r.setting.Client != nil {
		r.setting.Client = nil
	}
	return r
}

// EnableCookieWithJar set the specified *http.CookieJar to the request.
func (r *Request) EnableCookieWithJar(jar http.CookieJar) *Request {
	if !r.prepareSetting() {
		return nil
	}
	r.setting.Jar = jar
	r.setting.Client = nil
	return r
}

// Merge clone some properties of another Request into this.
func (r *Request) Merge(rr *Request) *Request {
	if r == nil || rr == nil {
		return nil
	}
	if len(rr.params) > 0 { // merge params
		for name, value := range rr.params {
			if _, ok := r.params[name]; !ok {
				r.params[name] = value
			}
		}
	}
	if rr.req != nil { // merge internal http.Request
		if r.req == nil {
			r.req = basicRequest()
		}
		if r.req.Method == "" && rr.req.Method != "" {
			r.req.Method = rr.req.Method
		}
		if r.req.Host == "" && rr.req.Host != "" {
			r.req.Host = rr.req.Host
		}
		if r.req.Proto != rr.req.Proto {
			r.req.Proto = rr.req.Proto
			r.req.ProtoMajor = rr.req.ProtoMajor
			r.req.ProtoMinor = rr.req.ProtoMinor
		}
		for name, value := range rr.req.Header {
			if _, ok := r.req.Header[name]; !ok {
				r.req.Header[name] = value
			}
		}
	}
	if rr.setting != nil { // merge setting
		rr.GetClient() // ensure client has been created. prevent creating client every request.
		s := *rr.setting
		r.setting = &s
	}

	return r
}
<|endoftext|>


<commit_before><commit_msg>httpserver: More organized startup output (#2497)<commit_after><|endoftext|>


<commit_before>/*
This package is just a collection of test cases
*/
package main

import (
	"github.com/urfave/cli"
	"os"
	"sort"
	"strings"
	"github.com/keltia/ripe-atlas"
	"log"
)

var (
	// flags
	fWant4 bool
	fWant6 bool

	fAllProbes bool
	fAllMeasurements bool

	fAsn string
	fCountry string
	fFieldList string
	fFormat string
	fOptFields string
	fSortOrder string
	fMeasureType string

	fVerbose bool
	fWantAnchor bool

	cliCommands []cli.Command
)

type ByAlphabet []cli.Command

func (a ByAlphabet) Len() int           { return len(a) }
func (a ByAlphabet) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByAlphabet) Less(i, j int) bool { return a[i].Name < a[j].Name }

// checkGlobalFlags is the place to check global parameters
func checkGlobalFlags(o map[string]string) (map[string]string) {
	opts := o
	if fSortOrder != "" {
		opts["sort"] = fSortOrder
	}

	if fFieldList != "" {
		opts["fields"] = fFieldList
	}

	if fOptFields != "" {
		opts["optional_fields"] = fOptFields
	}

	if fFormat != "" && validateFormat(fFormat) {
		opts["format"] = fFormat
	}
	return opts
}

// validateFormat allows only supported formats
func validateFormat(fmt string) bool {
	f := strings.ToLower(fmt)
	if f == "json" || f == "xml" || f == "api" || f == "txt" || f == "jsonp" {
		return true
	}
	return false
}

// main is the starting point (and everything)
func main() {
	app := cli.NewApp()
	app.Name = "atlas"
	app.Usage = "RIPE Atlas cli interface"
	app.Author = "Ollivier Robert <roberto@keltia.net>"
	app.Version = "0.1.0"
	app.HideVersion = true

	// General flags
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name: "format,f",
			Usage: "specify output format",
			Destination: &fFormat,
		},
		cli.BoolFlag{
			Name: "verbose,v",
			Usage: "verbose mode",
			Destination: &fVerbose,
		},
		cli.StringFlag{
			Name: "fields,F",
			Usage: "specify which fields are wanted",
			Destination: &fFieldList,
		},
		cli.StringFlag{
			Name: "opt-fields,O",
			Usage: "specify which optional fields are wanted",
			Destination: &fOptFields,
		},
		cli.StringFlag{
			Name:        "sort,S",
			Usage:       "sort results",
			Value:       "id",
			Destination: &fSortOrder,
		},
	}

	conf, err := atlas.LoadConfig("ripe-atlas")
	if conf.ApiKey != "" && err == nil {
		atlas.SetAuth(conf.ApiKey)
		log.Printf("Found API key!")
	} else {
		log.Printf("No API key!")
	}
	sort.Sort(ByAlphabet(cliCommands))
	app.Commands = cliCommands
	app.Run(os.Args)
}
<commit_msg>golint: ApiKey is now APIKey for consistency.<commit_after>/*
This package is just a collection of test cases
*/
package main

import (
	"github.com/urfave/cli"
	"os"
	"sort"
	"strings"
	"github.com/keltia/ripe-atlas"
	"log"
)

var (
	// flags
	fWant4 bool
	fWant6 bool

	fAllProbes bool
	fAllMeasurements bool

	fAsn string
	fCountry string
	fFieldList string
	fFormat string
	fOptFields string
	fSortOrder string
	fMeasureType string

	fVerbose bool
	fWantAnchor bool

	cliCommands []cli.Command
)

type ByAlphabet []cli.Command

func (a ByAlphabet) Len() int           { return len(a) }
func (a ByAlphabet) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByAlphabet) Less(i, j int) bool { return a[i].Name < a[j].Name }

// checkGlobalFlags is the place to check global parameters
func checkGlobalFlags(o map[string]string) (map[string]string) {
	opts := o
	if fSortOrder != "" {
		opts["sort"] = fSortOrder
	}

	if fFieldList != "" {
		opts["fields"] = fFieldList
	}

	if fOptFields != "" {
		opts["optional_fields"] = fOptFields
	}

	if fFormat != "" && validateFormat(fFormat) {
		opts["format"] = fFormat
	}
	return opts
}

// validateFormat allows only supported formats
func validateFormat(fmt string) bool {
	f := strings.ToLower(fmt)
	if f == "json" || f == "xml" || f == "api" || f == "txt" || f == "jsonp" {
		return true
	}
	return false
}

// main is the starting point (and everything)
func main() {
	app := cli.NewApp()
	app.Name = "atlas"
	app.Usage = "RIPE Atlas cli interface"
	app.Author = "Ollivier Robert <roberto@keltia.net>"
	app.Version = "0.1.0"
	app.HideVersion = true

	// General flags
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name: "format,f",
			Usage: "specify output format",
			Destination: &fFormat,
		},
		cli.BoolFlag{
			Name: "verbose,v",
			Usage: "verbose mode",
			Destination: &fVerbose,
		},
		cli.StringFlag{
			Name: "fields,F",
			Usage: "specify which fields are wanted",
			Destination: &fFieldList,
		},
		cli.StringFlag{
			Name: "opt-fields,O",
			Usage: "specify which optional fields are wanted",
			Destination: &fOptFields,
		},
		cli.StringFlag{
			Name:        "sort,S",
			Usage:       "sort results",
			Value:       "id",
			Destination: &fSortOrder,
		},
	}

	conf, err := atlas.LoadConfig("ripe-atlas")
	if conf.APIKey != "" && err == nil {
		atlas.SetAuth(conf.APIKey)
		log.Printf("Found API key!")
	} else {
		log.Printf("No API key!")
	}
	sort.Sort(ByAlphabet(cliCommands))
	app.Commands = cliCommands
	app.Run(os.Args)
}
<|endoftext|>