<commit_before>package server

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"time"

	"github.com/Comcast/webpa-common/concurrent"
	"github.com/Comcast/webpa-common/health"
	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/xhttp"
	"github.com/Comcast/webpa-common/xmetrics"
	"github.com/go-kit/kit/log"
	"github.com/justinas/alice"
	stdprometheus "github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const (
	DefaultBuild = "development"
	DefaultFqdn  = "localhost"
)

var (
	// ErrorNoPrimaryAddress is the error returned when no primary address is specified in a WebPA instance
	ErrorNoPrimaryAddress = errors.New("No primary address configured")
)

// executor is an internal type used to start an HTTP server.  *http.Server implements
// this interface.  It can be mocked for testing.
type executor interface {
	ListenAndServe() error
	ListenAndServeTLS(certificateFile, keyFile string) error
}

// Secure exposes the optional certificate information to be used when starting an HTTP server.
type Secure interface {
	// Certificate returns the certificate information associated with this Secure instance.
	// BOTH the returned file paths must be non-empty if a TLS server is desired.
	Certificate() (certificateFile, keyFile string)
}

// ListenAndServe invokes the appropriate server method based on the secure information.
// If Secure.Certificate() returns both a certificateFile and a keyFile, e.ListenAndServeTLS()
// is called to start the server.  Otherwise, e.ListenAndServe() is used.
func ListenAndServe(logger log.Logger, s Secure, e executor) {
	certificateFile, keyFile := s.Certificate()
	if len(certificateFile) > 0 && len(keyFile) > 0 {

		go func() {
			logging.Error(logger).Log(
				logging.ErrorKey(), e.ListenAndServeTLS(certificateFile, keyFile),
			)
		}()
	} else {
		go func() {
			logging.Error(logger).Log(
				logging.ErrorKey(), e.ListenAndServe(),
			)
		}()
	}
}

// Basic describes a simple HTTP server.  Typically, this struct has its values
// injected via Viper.  See the New function in this package.
type Basic struct {
	Name               string
	Address            string
	CertificateFile    string
	KeyFile            string
	ClientCACertFile   string
	LogConnectionState bool
}

func (b *Basic) Certificate() (certificateFile, keyFile string) {
	return b.CertificateFile, b.KeyFile
}

// New creates an http.Server using this instance's configuration.  The given logger is required,
// but the handler may be nil.  If the handler is nil, http.DefaultServeMux is used, which matches
// the behavior of http.Server.
//
// This method returns nil if the configured address is empty, effectively disabling
// this server from startup.
func (b *Basic) New(logger log.Logger, handler http.Handler) *http.Server {
	if len(b.Address) == 0 {
		return nil
	}

	// Adding MTLS support using client CA cert pool
	var tlsConfig *tls.Config
	certificateFile, keyFile := b.Certificate()
	// Only when HTTPS i.e. cert & key present, check for client CA and set TLS config for MTLS
	if len(certificateFile) > 0 && len(keyFile) > 0 && len(b.ClientCACertFile) > 0 {

		caCert, err := ioutil.ReadFile(b.ClientCACertFile)
		if err != nil {
			logging.Error(logger).Log(logging.MessageKey(), "Error in reading ClientCACertFile ",
				logging.ErrorKey(), err)
		} else {
			caCertPool := x509.NewCertPool()
			caCertPool.AppendCertsFromPEM(caCert)
			tlsConfig = &tls.Config{
				ClientCAs:  caCertPool,
				ClientAuth: tls.RequireAndVerifyClientCert,
			}
			tlsConfig.BuildNameToCertificate()
		}
	}

	server := &http.Server{
		Addr:      b.Address,
		Handler:   handler,
		ErrorLog:  NewErrorLog(b.Name, logger),
		TLSConfig: tlsConfig,
	}

	if b.LogConnectionState {
		server.ConnState = NewConnectionStateLogger(b.Name, logger)
	}

	return server
}

// Metric is the configurable factory for a metrics server.
type Metric struct {
	Name               string
	Address            string
	CertificateFile    string
	KeyFile            string
	LogConnectionState bool
	HandlerOptions     promhttp.HandlerOpts
	MetricsOptions     xmetrics.Options
}

func (m *Metric) Certificate() (certificateFile, keyFile string) {
	return m.CertificateFile, m.KeyFile
}

func (m *Metric) NewRegistry(modules ...xmetrics.Module) (xmetrics.Registry, error) {
	// always append the builtin server metrics, which can be overridden in configuration
	modules = append(modules, Metrics)
	return xmetrics.NewRegistry(&m.MetricsOptions, modules...)
}

func (m *Metric) New(logger log.Logger, chain alice.Chain, gatherer stdprometheus.Gatherer) *http.Server {
	if len(m.Address) == 0 {
		return nil
	}

	var (
		mux     = http.NewServeMux()
		handler = chain.Then(promhttp.HandlerFor(gatherer, m.HandlerOptions))
	)

	mux.Handle("/metrics", handler)
	server := &http.Server{
		Addr:     m.Address,
		Handler:  mux,
		ErrorLog: NewErrorLog(m.Name, logger),
	}

	if m.LogConnectionState {
		server.ConnState = NewConnectionStateLogger(m.Name, logger)
	}

	return server
}

// Health represents a configurable factory for a Health server.  As with the Basic type,
// if the Address is not specified, health is considered to be disabled.
//
// Due to a limitation of Viper, this struct does not use an embedded Basic
// instance.  Rather, it duplicates the fields so that Viper can inject them.
type Health struct {
	Name               string
	Address            string
	CertificateFile    string
	KeyFile            string
	LogConnectionState bool
	LogInterval        time.Duration
	Options            []string
}

func (h *Health) Certificate() (certificateFile, keyFile string) {
	return h.CertificateFile, h.KeyFile
}

// NewHealth creates a Health instance from this instance's configuration.  If the Address
// field is not supplied, this method returns nil.
func (h *Health) NewHealth(logger log.Logger, options ...health.Option) *health.Health {
	if len(h.Address) == 0 {
		return nil
	}

	for _, value := range h.Options {
		options = append(options, health.Stat(value))
	}

	return health.New(
		h.LogInterval,
		logger,
		options...,
	)
}

// New creates an HTTP server instance for serving health statistics.  If the health parameter
// is nil, then h.NewHealth is used to create a Health instance.  Otherwise, the health parameter
// is returned as is.
//
// If the Address option is not supplied, the health module is considered to be disabled.  In that
// case, this method simply returns the health parameter as the monitor and a nil server instance.
func (h *Health) New(logger log.Logger, chain alice.Chain, health *health.Health) (*health.Health, *http.Server) {
	if len(h.Address) == 0 {
		// health is disabled
		return nil, nil
	}

	if health == nil {
		if health = h.NewHealth(logger); health == nil {
			// should never hit this case, since NewHealth performs the same
			// Address field check as this method.  but, just to be safe ...
			return nil, nil
		}
	}

	mux := http.NewServeMux()
	mux.Handle("/health", chain.Then(health))

	server := &http.Server{
		Addr:     h.Address,
		Handler:  mux,
		ErrorLog: NewErrorLog(h.Name, logger),
	}

	if h.LogConnectionState {
		server.ConnState = NewConnectionStateLogger(h.Name, logger)
	}

	return health, server
}

// WebPA represents a server component within the WebPA cluster.  It is used for both
// primary servers (e.g. petasos) and supporting, embedded servers such as pprof.
type WebPA struct {
	// ApplicationName is the short identifier for the enclosing application, e.g. "talaria".
	// This value is defaulted to what's passed in via Initialize, but can be changed via injection.
	ApplicationName string

	// Primary is the main server for this application, e.g. petasos.
	Primary Basic

	// Alternate is an alternate server which serves the primary application logic.
	// Used to have the same API served on more than one port and possibly more than
	// one protocol, e.g. HTTP and HTTPS.
	Alternate Basic

	// Health describes the health server for this application.  Note that if the Address
	// is empty, no health server is started.
	Health Health

	// Pprof describes the pprof server for this application.  Note that if the Address
	// is empty, no pprof server is started.
	Pprof Basic

	// Metric describes the metrics provider server for this application
	Metric Metric

	// Build is the build string for the current codebase
	Build string

	// Fqdn is the fully-qualified domain name of this server, typically injected as a fact
	Fqdn string

	// Log is the logging configuration for this application.
	Log *logging.Options
}

// build returns the injected build string if available, DefaultBuild otherwise
func (w *WebPA) build() string {
	if w != nil && len(w.Build) > 0 {
		return w.Build
	}

	return DefaultBuild
}

// fqdn returns the injected fully-qualified domain name if available, DefaultFqdn otherwise
func (w *WebPA) fqdn() string {
	if w != nil && len(w.Fqdn) > 0 {
		return w.Fqdn
	}

	return DefaultFqdn
}

// Prepare gets a WebPA server ready for execution.  This method does not return errors, but the returned
// Runnable may return an error.  The supplied logger will usually come from the New function, but the
// WebPA.Log object can be used to create a different logger if desired.
//
// The caller may pass an arbitrary Health instance.  If this parameter is nil, this method will attempt to
// create one using Health.NewHealth.  In either case, if Health.Address is not supplied, no health server
// will be instantiated.
//
// The caller may also pass a gatherer type. If it is not provided, the default provided by prometheus is used.
//
// The supplied http.Handler is used for the primary server.  If the alternate server has an address,
// it will also be used for that server.  The health server uses an internally create handler, while pprof and metrics
// servers use http.DefaultServeMux.  The health Monitor created from configuration is returned so that other
// infrastructure can make use of it.
func (w *WebPA) Prepare(logger log.Logger, health *health.Health, registry xmetrics.Registry, primaryHandler http.Handler) (health.Monitor, concurrent.Runnable) {
	// allow the health instance to be non-nil, in which case it will be used in favor of
	// the WebPA-configured instance.
	var (
		staticHeaders = xhttp.StaticHeaders(http.Header{
			fmt.Sprintf("X-%s-Build", w.ApplicationName):      {w.build()},
			fmt.Sprintf("X-%s-Server", w.ApplicationName):     {w.fqdn()},
			fmt.Sprintf("X-%s-Start-Time", w.ApplicationName): {time.Now().UTC().Format(time.RFC822)},
		})

		healthHandler, healthServer = w.Health.New(logger, alice.New(staticHeaders), health)
		infoLog                     = logging.Info(logger)
	)

	return healthHandler, concurrent.RunnableFunc(func(waitGroup *sync.WaitGroup, shutdown <-chan struct{}) error {
		if healthHandler != nil && healthServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Health.Name, "address", w.Health.Address)
			ListenAndServe(logger, &w.Health, healthServer)
			healthHandler.Run(waitGroup, shutdown)
		}

		if pprofServer := w.Pprof.New(logger, nil); pprofServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Pprof.Name, "address", w.Pprof.Address)
			ListenAndServe(logger, &w.Pprof, pprofServer)
		}

		primaryHandler = staticHeaders(w.decorateWithBasicMetrics(registry, primaryHandler))
		if primaryServer := w.Primary.New(logger, primaryHandler); primaryServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Primary.Name, "address", w.Primary.Address)
			ListenAndServe(logger, &w.Primary, primaryServer)
		} else {
			return ErrorNoPrimaryAddress
		}

		if alternateServer := w.Alternate.New(logger, primaryHandler); alternateServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Alternate.Name, "address", w.Alternate.Address)
			ListenAndServe(logger, &w.Alternate, alternateServer)
		}

		if metricsServer := w.Metric.New(logger, alice.New(staticHeaders), registry); metricsServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Metric.Name, "address", w.Metric.Address)
			ListenAndServe(logger, &w.Metric, metricsServer)
		}

		return nil
	})
}

//decorateWithBasicMetrics wraps a WebPA server handler with basic instrumentation metrics
func (w *WebPA) decorateWithBasicMetrics(p xmetrics.PrometheusProvider, next http.Handler) http.Handler {
	var (
		requestCounterVec    = p.NewCounterVec("api_requests_total")
		inFlightGauge        = p.NewGaugeVec("in_flight_requests").WithLabelValues()
		requestDurationVec   = p.NewHistogramVec("request_duration_seconds")
		requestSizeVec       = p.NewHistogramVec("request_size_bytes")
		responseSizeVec      = p.NewHistogramVec("response_size_bytes")
		timeToWriteHeaderVec = p.NewHistogramVec("time_writing_header_seconds")
	)

	//todo: Example documentation does something interesting with /pull vs. /push endpoints
	//https://godoc.org/github.com/prometheus/client_golang/prometheus/promhttp#InstrumentHandlerDuration
	//for now, let's keep it simple so /metrics only

	return promhttp.InstrumentHandlerInFlight(inFlightGauge,
		promhttp.InstrumentHandlerCounter(requestCounterVec,
			promhttp.InstrumentHandlerDuration(requestDurationVec,
				promhttp.InstrumentHandlerResponseSize(responseSizeVec,
					promhttp.InstrumentHandlerRequestSize(requestSizeVec,
						promhttp.InstrumentHandlerTimeToWriteHeader(timeToWriteHeaderVec, next))),
			),
		),
	)
}
<commit_msg>Extra server metadata headers<commit_after>package server

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"time"

	"github.com/Comcast/webpa-common/concurrent"
	"github.com/Comcast/webpa-common/health"
	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/xhttp"
	"github.com/Comcast/webpa-common/xmetrics"
	"github.com/go-kit/kit/log"
	"github.com/justinas/alice"
	stdprometheus "github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

const (
	DefaultBuild  = "development"
	DefaultServer = "localhost"
	DefaultRegion = "local"
	DefaultFlavor = "development"
)

var (
	// ErrorNoPrimaryAddress is the error returned when no primary address is specified in a WebPA instance
	ErrorNoPrimaryAddress = errors.New("No primary address configured")
)

// executor is an internal type used to start an HTTP server.  *http.Server implements
// this interface.  It can be mocked for testing.
type executor interface {
	ListenAndServe() error
	ListenAndServeTLS(certificateFile, keyFile string) error
}

// Secure exposes the optional certificate information to be used when starting an HTTP server.
type Secure interface {
	// Certificate returns the certificate information associated with this Secure instance.
	// BOTH the returned file paths must be non-empty if a TLS server is desired.
	Certificate() (certificateFile, keyFile string)
}

// ListenAndServe invokes the appropriate server method based on the secure information.
// If Secure.Certificate() returns both a certificateFile and a keyFile, e.ListenAndServeTLS()
// is called to start the server.  Otherwise, e.ListenAndServe() is used.
func ListenAndServe(logger log.Logger, s Secure, e executor) {
	certificateFile, keyFile := s.Certificate()
	if len(certificateFile) > 0 && len(keyFile) > 0 {

		go func() {
			logging.Error(logger).Log(
				logging.ErrorKey(), e.ListenAndServeTLS(certificateFile, keyFile),
			)
		}()
	} else {
		go func() {
			logging.Error(logger).Log(
				logging.ErrorKey(), e.ListenAndServe(),
			)
		}()
	}
}

// Basic describes a simple HTTP server.  Typically, this struct has its values
// injected via Viper.  See the New function in this package.
type Basic struct {
	Name               string
	Address            string
	CertificateFile    string
	KeyFile            string
	ClientCACertFile   string
	LogConnectionState bool
}

func (b *Basic) Certificate() (certificateFile, keyFile string) {
	return b.CertificateFile, b.KeyFile
}

// New creates an http.Server using this instance's configuration.  The given logger is required,
// but the handler may be nil.  If the handler is nil, http.DefaultServeMux is used, which matches
// the behavior of http.Server.
//
// This method returns nil if the configured address is empty, effectively disabling
// this server from startup.
func (b *Basic) New(logger log.Logger, handler http.Handler) *http.Server {
	if len(b.Address) == 0 {
		return nil
	}

	// Adding MTLS support using client CA cert pool
	var tlsConfig *tls.Config
	certificateFile, keyFile := b.Certificate()
	// Only when HTTPS i.e. cert & key present, check for client CA and set TLS config for MTLS
	if len(certificateFile) > 0 && len(keyFile) > 0 && len(b.ClientCACertFile) > 0 {

		caCert, err := ioutil.ReadFile(b.ClientCACertFile)
		if err != nil {
			logging.Error(logger).Log(logging.MessageKey(), "Error in reading ClientCACertFile ",
				logging.ErrorKey(), err)
		} else {
			caCertPool := x509.NewCertPool()
			caCertPool.AppendCertsFromPEM(caCert)
			tlsConfig = &tls.Config{
				ClientCAs:  caCertPool,
				ClientAuth: tls.RequireAndVerifyClientCert,
			}
			tlsConfig.BuildNameToCertificate()
		}
	}

	server := &http.Server{
		Addr:      b.Address,
		Handler:   handler,
		ErrorLog:  NewErrorLog(b.Name, logger),
		TLSConfig: tlsConfig,
	}

	if b.LogConnectionState {
		server.ConnState = NewConnectionStateLogger(b.Name, logger)
	}

	return server
}

// Metric is the configurable factory for a metrics server.
type Metric struct {
	Name               string
	Address            string
	CertificateFile    string
	KeyFile            string
	LogConnectionState bool
	HandlerOptions     promhttp.HandlerOpts
	MetricsOptions     xmetrics.Options
}

func (m *Metric) Certificate() (certificateFile, keyFile string) {
	return m.CertificateFile, m.KeyFile
}

func (m *Metric) NewRegistry(modules ...xmetrics.Module) (xmetrics.Registry, error) {
	// always append the builtin server metrics, which can be overridden in configuration
	modules = append(modules, Metrics)
	return xmetrics.NewRegistry(&m.MetricsOptions, modules...)
}

func (m *Metric) New(logger log.Logger, chain alice.Chain, gatherer stdprometheus.Gatherer) *http.Server {
	if len(m.Address) == 0 {
		return nil
	}

	var (
		mux     = http.NewServeMux()
		handler = chain.Then(promhttp.HandlerFor(gatherer, m.HandlerOptions))
	)

	mux.Handle("/metrics", handler)
	server := &http.Server{
		Addr:     m.Address,
		Handler:  mux,
		ErrorLog: NewErrorLog(m.Name, logger),
	}

	if m.LogConnectionState {
		server.ConnState = NewConnectionStateLogger(m.Name, logger)
	}

	return server
}

// Health represents a configurable factory for a Health server.  As with the Basic type,
// if the Address is not specified, health is considered to be disabled.
//
// Due to a limitation of Viper, this struct does not use an embedded Basic
// instance.  Rather, it duplicates the fields so that Viper can inject them.
type Health struct {
	Name               string
	Address            string
	CertificateFile    string
	KeyFile            string
	LogConnectionState bool
	LogInterval        time.Duration
	Options            []string
}

func (h *Health) Certificate() (certificateFile, keyFile string) {
	return h.CertificateFile, h.KeyFile
}

// NewHealth creates a Health instance from this instance's configuration.  If the Address
// field is not supplied, this method returns nil.
func (h *Health) NewHealth(logger log.Logger, options ...health.Option) *health.Health {
	if len(h.Address) == 0 {
		return nil
	}

	for _, value := range h.Options {
		options = append(options, health.Stat(value))
	}

	return health.New(
		h.LogInterval,
		logger,
		options...,
	)
}

// New creates an HTTP server instance for serving health statistics.  If the health parameter
// is nil, then h.NewHealth is used to create a Health instance.  Otherwise, the health parameter
// is returned as is.
//
// If the Address option is not supplied, the health module is considered to be disabled.  In that
// case, this method simply returns the health parameter as the monitor and a nil server instance.
func (h *Health) New(logger log.Logger, chain alice.Chain, health *health.Health) (*health.Health, *http.Server) {
	if len(h.Address) == 0 {
		// health is disabled
		return nil, nil
	}

	if health == nil {
		if health = h.NewHealth(logger); health == nil {
			// should never hit this case, since NewHealth performs the same
			// Address field check as this method.  but, just to be safe ...
			return nil, nil
		}
	}

	mux := http.NewServeMux()
	mux.Handle("/health", chain.Then(health))

	server := &http.Server{
		Addr:     h.Address,
		Handler:  mux,
		ErrorLog: NewErrorLog(h.Name, logger),
	}

	if h.LogConnectionState {
		server.ConnState = NewConnectionStateLogger(h.Name, logger)
	}

	return health, server
}

// WebPA represents a server component within the WebPA cluster.  It is used for both
// primary servers (e.g. petasos) and supporting, embedded servers such as pprof.
type WebPA struct {
	// ApplicationName is the short identifier for the enclosing application, e.g. "talaria".
	// This value is defaulted to what's passed in via Initialize, but can be changed via injection.
	ApplicationName string

	// Primary is the main server for this application, e.g. petasos.
	Primary Basic

	// Alternate is an alternate server which serves the primary application logic.
	// Used to have the same API served on more than one port and possibly more than
	// one protocol, e.g. HTTP and HTTPS.
	Alternate Basic

	// Health describes the health server for this application.  Note that if the Address
	// is empty, no health server is started.
	Health Health

	// Pprof describes the pprof server for this application.  Note that if the Address
	// is empty, no pprof server is started.
	Pprof Basic

	// Metric describes the metrics provider server for this application
	Metric Metric

	// Build is the build string for the current codebase
	Build string

	// Server is the fully-qualified domain name of this server, typically injected as a fact
	Server string

	// Region is the region in which this server is running, typically injected as a fact
	Region string

	// Flavor is the flavor of this server, typically injected as a fact
	Flavor string

	// Log is the logging configuration for this application.
	Log *logging.Options
}

// build returns the injected build string if available, DefaultBuild otherwise
func (w *WebPA) build() string {
	if w != nil && len(w.Build) > 0 {
		return w.Build
	}

	return DefaultBuild
}

// server returns the injected fully-qualified domain name if available, DefaultServer otherwise
func (w *WebPA) server() string {
	if w != nil && len(w.Server) > 0 {
		return w.Server
	}

	return DefaultServer
}

// region returns the region in which this server is running, or DefaultRegion otherwise
func (w *WebPA) region() string {
	if w != nil && len(w.Region) > 0 {
		return w.Region
	}

	return DefaultRegion
}

// flavor returns the region in which this server is running, or DefaultRegion otherwise
func (w *WebPA) flavor() string {
	if w != nil && len(w.Flavor) > 0 {
		return w.Flavor
	}

	return DefaultFlavor
}

// Prepare gets a WebPA server ready for execution.  This method does not return errors, but the returned
// Runnable may return an error.  The supplied logger will usually come from the New function, but the
// WebPA.Log object can be used to create a different logger if desired.
//
// The caller may pass an arbitrary Health instance.  If this parameter is nil, this method will attempt to
// create one using Health.NewHealth.  In either case, if Health.Address is not supplied, no health server
// will be instantiated.
//
// The caller may also pass a gatherer type. If it is not provided, the default provided by prometheus is used.
//
// The supplied http.Handler is used for the primary server.  If the alternate server has an address,
// it will also be used for that server.  The health server uses an internally create handler, while pprof and metrics
// servers use http.DefaultServeMux.  The health Monitor created from configuration is returned so that other
// infrastructure can make use of it.
func (w *WebPA) Prepare(logger log.Logger, health *health.Health, registry xmetrics.Registry, primaryHandler http.Handler) (health.Monitor, concurrent.Runnable) {
	// allow the health instance to be non-nil, in which case it will be used in favor of
	// the WebPA-configured instance.
	var (
		staticHeaders = xhttp.StaticHeaders(http.Header{
			fmt.Sprintf("X-%s-Build", w.ApplicationName):      {w.build()},
			fmt.Sprintf("X-%s-Server", w.ApplicationName):     {w.server()},
			fmt.Sprintf("X-%s-Region", w.ApplicationName):     {w.region()},
			fmt.Sprintf("X-%s-Flavor", w.ApplicationName):     {w.flavor()},
			fmt.Sprintf("X-%s-Start-Time", w.ApplicationName): {time.Now().UTC().Format(time.RFC822)},
		})

		healthHandler, healthServer = w.Health.New(logger, alice.New(staticHeaders), health)
		infoLog                     = logging.Info(logger)
	)

	return healthHandler, concurrent.RunnableFunc(func(waitGroup *sync.WaitGroup, shutdown <-chan struct{}) error {
		if healthHandler != nil && healthServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Health.Name, "address", w.Health.Address)
			ListenAndServe(logger, &w.Health, healthServer)
			healthHandler.Run(waitGroup, shutdown)
		}

		if pprofServer := w.Pprof.New(logger, nil); pprofServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Pprof.Name, "address", w.Pprof.Address)
			ListenAndServe(logger, &w.Pprof, pprofServer)
		}

		primaryHandler = staticHeaders(w.decorateWithBasicMetrics(registry, primaryHandler))
		if primaryServer := w.Primary.New(logger, primaryHandler); primaryServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Primary.Name, "address", w.Primary.Address)
			ListenAndServe(logger, &w.Primary, primaryServer)
		} else {
			return ErrorNoPrimaryAddress
		}

		if alternateServer := w.Alternate.New(logger, primaryHandler); alternateServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Alternate.Name, "address", w.Alternate.Address)
			ListenAndServe(logger, &w.Alternate, alternateServer)
		}

		if metricsServer := w.Metric.New(logger, alice.New(staticHeaders), registry); metricsServer != nil {
			infoLog.Log(logging.MessageKey(), "starting server", "name", w.Metric.Name, "address", w.Metric.Address)
			ListenAndServe(logger, &w.Metric, metricsServer)
		}

		return nil
	})
}

//decorateWithBasicMetrics wraps a WebPA server handler with basic instrumentation metrics
func (w *WebPA) decorateWithBasicMetrics(p xmetrics.PrometheusProvider, next http.Handler) http.Handler {
	var (
		requestCounterVec    = p.NewCounterVec("api_requests_total")
		inFlightGauge        = p.NewGaugeVec("in_flight_requests").WithLabelValues()
		requestDurationVec   = p.NewHistogramVec("request_duration_seconds")
		requestSizeVec       = p.NewHistogramVec("request_size_bytes")
		responseSizeVec      = p.NewHistogramVec("response_size_bytes")
		timeToWriteHeaderVec = p.NewHistogramVec("time_writing_header_seconds")
	)

	//todo: Example documentation does something interesting with /pull vs. /push endpoints
	//https://godoc.org/github.com/prometheus/client_golang/prometheus/promhttp#InstrumentHandlerDuration
	//for now, let's keep it simple so /metrics only

	return promhttp.InstrumentHandlerInFlight(inFlightGauge,
		promhttp.InstrumentHandlerCounter(requestCounterVec,
			promhttp.InstrumentHandlerDuration(requestDurationVec,
				promhttp.InstrumentHandlerResponseSize(responseSizeVec,
					promhttp.InstrumentHandlerRequestSize(requestSizeVec,
						promhttp.InstrumentHandlerTimeToWriteHeader(timeToWriteHeaderVec, next))),
			),
		),
	)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/litl/galaxy/log"
)

var (
	httpRouter *HostRouter
)

// This works along with the ServiceRegistry, and the individual Services to
// route http requests based on the Host header. The Resgistry hold the mapping
// of VHost names to individual services, and each service has it's own
// ReeverseProxy to fulfill the request.
// HostRouter contains the ReverseProxy http Listener, and has an http.Handler
// to service the requets.
type HostRouter struct {
	sync.Mutex
	// the http frontend
	server *http.Server

	// automatically redirect to https
	SSLOnly bool

	// HTTP/HTTPS
	Scheme string

	// track our listener so we can kill the server
	listener net.Listener
}

func NewHostRouter(httpServer *http.Server) *HostRouter {
	r := &HostRouter{
		Scheme: "http",
	}
	httpServer.Handler = r
	r.server = httpServer
	return r
}

func (r *HostRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	reqId := genId()
	req.Header.Set("X-Request-Id", reqId)

	if r.SSLOnly {
		if req.TLS != nil || req.Header.Get("X-Forwarded-Proto") != "https" {
			//TODO: verify RequestURI
			redirLoc := "https://" + req.Host + req.RequestURI
			http.Redirect(w, req, redirLoc, http.StatusMovedPermanently)
			return
		}
	}

	var err error
	host := req.Host

	if strings.Contains(host, ":") {
		host, _, err = net.SplitHostPort(req.Host)
		if err != nil {
			log.Warnf("%s", err)
		}
	}

	svc := Registry.GetVHostService(host)

	if svc != nil && svc.httpProxy != nil {
		// The vhost has a service registered, give it to the proxy
		svc.ServeHTTP(w, req)
		return
	}

	r.noHostHandler(w, req)
}

func (r *HostRouter) noHostHandler(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(http.StatusNotFound)
	fmt.Fprintln(w, "Not found")
}

// TODO: collect more stats?

// Start the HTTP Router frontend.
// Takes a channel to notify when the listener is started
// to safely synchronize tests.
func (r *HostRouter) Start(ready chan bool) {
	//FIXME: poor locking strategy
	r.Lock()
	var err error
	r.listener, err = newTimeoutListener("tcp", r.server.Addr, 300*time.Second)
	if err != nil {
		log.Errorf("%s", err)
		r.Unlock()
		return
	}

	listener := r.listener
	if r.Scheme == "https" {
		listener = tls.NewListener(listener, r.server.TLSConfig)
	}

	r.Unlock()

	log.Printf("%s server listening at %s", r.Scheme, r.server.Addr)
	if ready != nil {
		close(ready)
	}

	// This will log a closed connection error every time we Stop
	// but that's mostly a testing issue.
	log.Errorf("%s", r.server.Serve(listener))
}

func (r *HostRouter) Stop() {
	r.listener.Close()
}

func startHTTPServer(wg *sync.WaitGroup) {
	defer wg.Done()

	//TODO: configure these timeouts somewhere
	httpServer := &http.Server{
		Addr:           httpAddr,
		ReadTimeout:    10 * time.Minute,
		WriteTimeout:   10 * time.Minute,
		MaxHeaderBytes: 1 << 20,
	}

	httpRouter = NewHostRouter(httpServer)
	httpRouter.SSLOnly = sslOnly

	httpRouter.Start(nil)
}

// find certs in and is the named directory, and match them up by their base
// name using '.pem' and '.key' as extensions.
func loadCerts(certDir string) (*tls.Config, error) {
	abs, err := filepath.Abs(certDir)
	if err != nil {
		return nil, err
	}

	dir, err := ioutil.ReadDir(abs)
	if err != nil {
		return nil, err
	}

	// [cert, key] pairs
	pairs := make(map[string][2]string)

	for _, f := range dir {
		name := f.Name()
		if strings.HasSuffix(name, ".pem") {
			p := pairs[name[:len(name)-4]]
			p[0] = filepath.Join(abs, name)
			pairs[name[:len(name)-4]] = p
		}
		if strings.HasSuffix(name, ".key") {
			p := pairs[name[:len(name)-4]]
			p[1] = filepath.Join(abs, name)
			pairs[name[:len(name)-4]] = p
		}
	}

	tlsCfg := &tls.Config{
		NextProtos: []string{"http/1.1"},
	}

	for key, pair := range pairs {
		if pair[0] == "" {
			log.Errorf("missing cert for key: %s", pair[1])
			continue
		}
		if pair[1] == "" {
			log.Errorf("missing key for cert: %s", pair[0])
			continue
		}

		cert, err := tls.LoadX509KeyPair(pair[0], pair[1])
		if err != nil {
			log.Error(err)
			continue
		}
		tlsCfg.Certificates = append(tlsCfg.Certificates, cert)
		log.Debugf("loaded X509KeyPair for %s", key)
	}

	if len(tlsCfg.Certificates) == 0 {
		return nil, fmt.Errorf("no tls certificates loaded")
	}

	tlsCfg.BuildNameToCertificate()

	return tlsCfg, nil
}

func startHTTPSServer(wg *sync.WaitGroup) {
	defer wg.Done()

	tlsCfg, err := loadCerts(certDir)
	if err != nil {
		log.Error(err)
		return
	}

	//TODO: configure these timeouts somewhere
	httpsServer := &http.Server{
		Addr:           httpsAddr,
		ReadTimeout:    10 * time.Minute,
		WriteTimeout:   10 * time.Minute,
		MaxHeaderBytes: 1 << 20,
		TLSConfig:      tlsCfg,
	}

	httpRouter = NewHostRouter(httpsServer)
	httpRouter.Scheme = "https"
	httpRouter.SSLOnly = sslOnly

	httpRouter.Start(nil)
}

type ErrorPage struct {
	// The Mutex protects access to the body slice, and headers
	// Everything else should be static once the ErrorPage is created.
	sync.Mutex

	Location    string
	StatusCodes []int

	// body contains the cached error page
	body []byte
	// important headers
	header http.Header
}

func (e *ErrorPage) Body() []byte {
	e.Lock()
	defer e.Unlock()
	return e.body
}

func (e *ErrorPage) SetBody(b []byte) {
	e.Lock()
	defer e.Unlock()
	e.body = b
}

func (e *ErrorPage) Header() http.Header {
	e.Lock()
	defer e.Unlock()
	return e.header
}

func (e *ErrorPage) SetHeader(h http.Header) {
	e.Lock()
	defer e.Unlock()
	e.header = h
}

// List of headers we want to cache for ErrorPages
var ErrorHeaders = []string{
	"Content-Type",
	"Content-Encoding",
	"Cache-Control",
	"Last-Modified",
	"Retry-After",
	"Set-Cookie",
}

// ErrorResponse provides a ReverProxy callback to process a response and
// insert custom error pages for a virtual host.
type ErrorResponse struct {
	sync.Mutex

	// map them by status for responses
	pages map[int]*ErrorPage

	// keep this handy to refresh the pages
	client *http.Client
}

func NewErrorResponse(pages map[string][]int) *ErrorResponse {
	errors := &ErrorResponse{
		pages: make(map[int]*ErrorPage),
	}

	// aggressively timeout connections
	errors.client = &http.Client{
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 2 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 2 * time.Second,
		},
		Timeout: 5 * time.Second,
	}

	if pages != nil {
		errors.Update(pages)
	}
	return errors
}

// Get the ErrorPage, returning nil if the page was incomplete.
// We permanently cache error pages and headers once we've seen them.
func (e *ErrorResponse) Get(code int) *ErrorPage {
	e.Lock()
	page, ok := e.pages[code]
	e.Unlock()

	if !ok {
		// this is a code we don't handle
		return nil
	}

	body := page.Body()
	if body != nil {
		return page
	}

	// we haven't successfully fetched this error
	e.fetch(page)
	return page
}

func (e *ErrorResponse) fetch(page *ErrorPage) {
	log.Debugf("Fetching error page from %s", page.Location)
	resp, err := e.client.Get(page.Location)
	if err != nil {
		log.Warnf("Could not fetch %s: %s", page.Location, err.Error())
		return
	}
	defer resp.Body.Close()

	// If the StatusCode matches any of our registered codes, it's OK
	for _, code := range page.StatusCodes {
		if resp.StatusCode == code {
			resp.StatusCode = http.StatusOK
			break
		}
	}

	if resp.StatusCode != http.StatusOK {
		log.Warnf("Server returned %d when fetching %s", resp.StatusCode, page.Location)
		return
	}

	header := make(map[string][]string)
	for _, key := range ErrorHeaders {
		if hdr, ok := resp.Header[key]; ok {
			header[key] = hdr
		}
	}
	// set the headers along with the body below

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Warnf("Error reading response from %s: %s", page.Location, err.Error())
		return
	}

	if len(body) > 0 {
		page.SetHeader(header)
		page.SetBody(body)
		return
	}
	log.Warnf("Empty response from %s", page.Location)
}

// This replaces all existing ErrorPages
func (e *ErrorResponse) Update(pages map[string][]int) {
	e.Lock()
	defer e.Unlock()

	e.pages = make(map[int]*ErrorPage)

	for loc, codes := range pages {
		page := &ErrorPage{
			StatusCodes: codes,
			Location:    loc,
		}

		for _, code := range codes {
			e.pages[code] = page
		}
		go e.fetch(page)
	}
}

func (e *ErrorResponse) CheckResponse(pr *ProxyRequest) bool {

	errPage := e.Get(pr.Response.StatusCode)
	if errPage != nil {
		// load the cached headers
		header := pr.ResponseWriter.Header()
		for key, val := range errPage.Header() {
			header[key] = val
		}

		pr.ResponseWriter.WriteHeader(pr.Response.StatusCode)
		pr.ResponseWriter.Write(errPage.Body())
		return false
	}

	return true
}

func logProxyRequest(pr *ProxyRequest) bool {
	// TODO: we may to be able to switch this off
	if pr == nil || pr.Request == nil {
		return true
	}

	var id, method, clientIP, url, backend, agent string
	var status int

	duration := pr.FinishTime.Sub(pr.StartTime)

	id = pr.Request.Header.Get("X-Request-Id")
	method = pr.Request.Method
	url = pr.Request.Host + pr.Request.RequestURI
	agent = pr.Request.UserAgent()
	status = pr.Response.StatusCode

	clientIP = pr.Request.Header.Get("X-Forwarded-For")
	if clientIP == "" {
		clientIP = pr.Request.RemoteAddr
	}

	if pr.Response != nil && pr.Response.Request != nil && pr.Response.Request.URL != nil {
		backend = pr.Response.Request.URL.Host
	}

	err := fmt.Sprintf("%v", pr.ProxyError)

	fmtStr := "id=%s method=%s clientIp=%s url=%s backend=%s status=%d duration=%s agent=%s, err=%s"

	log.Printf(fmtStr, id, method, clientIP, url, backend, status, duration, agent, err)
	return true
}
<commit_msg>shuttle: Log http(s) in all uppercase<commit_after>package main

import (
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/litl/galaxy/log"
)

var (
	httpRouter *HostRouter
)

// This works along with the ServiceRegistry, and the individual Services to
// route http requests based on the Host header. The Resgistry hold the mapping
// of VHost names to individual services, and each service has it's own
// ReeverseProxy to fulfill the request.
// HostRouter contains the ReverseProxy http Listener, and has an http.Handler
// to service the requets.
type HostRouter struct {
	sync.Mutex
	// the http frontend
	server *http.Server

	// automatically redirect to https
	SSLOnly bool

	// HTTP/HTTPS
	Scheme string

	// track our listener so we can kill the server
	listener net.Listener
}

func NewHostRouter(httpServer *http.Server) *HostRouter {
	r := &HostRouter{
		Scheme: "http",
	}
	httpServer.Handler = r
	r.server = httpServer
	return r
}

func (r *HostRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	reqId := genId()
	req.Header.Set("X-Request-Id", reqId)

	if r.SSLOnly {
		if req.TLS != nil || req.Header.Get("X-Forwarded-Proto") != "https" {
			//TODO: verify RequestURI
			redirLoc := "https://" + req.Host + req.RequestURI
			http.Redirect(w, req, redirLoc, http.StatusMovedPermanently)
			return
		}
	}

	var err error
	host := req.Host

	if strings.Contains(host, ":") {
		host, _, err = net.SplitHostPort(req.Host)
		if err != nil {
			log.Warnf("%s", err)
		}
	}

	svc := Registry.GetVHostService(host)

	if svc != nil && svc.httpProxy != nil {
		// The vhost has a service registered, give it to the proxy
		svc.ServeHTTP(w, req)
		return
	}

	r.noHostHandler(w, req)
}

func (r *HostRouter) noHostHandler(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(http.StatusNotFound)
	fmt.Fprintln(w, "Not found")
}

// TODO: collect more stats?

// Start the HTTP Router frontend.
// Takes a channel to notify when the listener is started
// to safely synchronize tests.
func (r *HostRouter) Start(ready chan bool) {
	//FIXME: poor locking strategy
	r.Lock()
	var err error
	r.listener, err = newTimeoutListener("tcp", r.server.Addr, 300*time.Second)
	if err != nil {
		log.Errorf("%s", err)
		r.Unlock()
		return
	}

	listener := r.listener
	if r.Scheme == "https" {
		listener = tls.NewListener(listener, r.server.TLSConfig)
	}

	r.Unlock()

	log.Printf("%s server listening at %s", strings.ToUpper(r.Scheme), r.server.Addr)
	if ready != nil {
		close(ready)
	}

	// This will log a closed connection error every time we Stop
	// but that's mostly a testing issue.
	log.Errorf("%s", r.server.Serve(listener))
}

func (r *HostRouter) Stop() {
	r.listener.Close()
}

func startHTTPServer(wg *sync.WaitGroup) {
	defer wg.Done()

	//TODO: configure these timeouts somewhere
	httpServer := &http.Server{
		Addr:           httpAddr,
		ReadTimeout:    10 * time.Minute,
		WriteTimeout:   10 * time.Minute,
		MaxHeaderBytes: 1 << 20,
	}

	httpRouter = NewHostRouter(httpServer)
	httpRouter.SSLOnly = sslOnly

	httpRouter.Start(nil)
}

// find certs in and is the named directory, and match them up by their base
// name using '.pem' and '.key' as extensions.
func loadCerts(certDir string) (*tls.Config, error) {
	abs, err := filepath.Abs(certDir)
	if err != nil {
		return nil, err
	}

	dir, err := ioutil.ReadDir(abs)
	if err != nil {
		return nil, err
	}

	// [cert, key] pairs
	pairs := make(map[string][2]string)

	for _, f := range dir {
		name := f.Name()
		if strings.HasSuffix(name, ".pem") {
			p := pairs[name[:len(name)-4]]
			p[0] = filepath.Join(abs, name)
			pairs[name[:len(name)-4]] = p
		}
		if strings.HasSuffix(name, ".key") {
			p := pairs[name[:len(name)-4]]
			p[1] = filepath.Join(abs, name)
			pairs[name[:len(name)-4]] = p
		}
	}

	tlsCfg := &tls.Config{
		NextProtos: []string{"http/1.1"},
	}

	for key, pair := range pairs {
		if pair[0] == "" {
			log.Errorf("missing cert for key: %s", pair[1])
			continue
		}
		if pair[1] == "" {
			log.Errorf("missing key for cert: %s", pair[0])
			continue
		}

		cert, err := tls.LoadX509KeyPair(pair[0], pair[1])
		if err != nil {
			log.Error(err)
			continue
		}
		tlsCfg.Certificates = append(tlsCfg.Certificates, cert)
		log.Debugf("loaded X509KeyPair for %s", key)
	}

	if len(tlsCfg.Certificates) == 0 {
		return nil, fmt.Errorf("no tls certificates loaded")
	}

	tlsCfg.BuildNameToCertificate()

	return tlsCfg, nil
}

func startHTTPSServer(wg *sync.WaitGroup) {
	defer wg.Done()

	tlsCfg, err := loadCerts(certDir)
	if err != nil {
		log.Error(err)
		return
	}

	//TODO: configure these timeouts somewhere
	httpsServer := &http.Server{
		Addr:           httpsAddr,
		ReadTimeout:    10 * time.Minute,
		WriteTimeout:   10 * time.Minute,
		MaxHeaderBytes: 1 << 20,
		TLSConfig:      tlsCfg,
	}

	httpRouter = NewHostRouter(httpsServer)
	httpRouter.Scheme = "https"
	httpRouter.SSLOnly = sslOnly

	httpRouter.Start(nil)
}

type ErrorPage struct {
	// The Mutex protects access to the body slice, and headers
	// Everything else should be static once the ErrorPage is created.
	sync.Mutex

	Location    string
	StatusCodes []int

	// body contains the cached error page
	body []byte
	// important headers
	header http.Header
}

func (e *ErrorPage) Body() []byte {
	e.Lock()
	defer e.Unlock()
	return e.body
}

func (e *ErrorPage) SetBody(b []byte) {
	e.Lock()
	defer e.Unlock()
	e.body = b
}

func (e *ErrorPage) Header() http.Header {
	e.Lock()
	defer e.Unlock()
	return e.header
}

func (e *ErrorPage) SetHeader(h http.Header) {
	e.Lock()
	defer e.Unlock()
	e.header = h
}

// List of headers we want to cache for ErrorPages
var ErrorHeaders = []string{
	"Content-Type",
	"Content-Encoding",
	"Cache-Control",
	"Last-Modified",
	"Retry-After",
	"Set-Cookie",
}

// ErrorResponse provides a ReverProxy callback to process a response and
// insert custom error pages for a virtual host.
type ErrorResponse struct {
	sync.Mutex

	// map them by status for responses
	pages map[int]*ErrorPage

	// keep this handy to refresh the pages
	client *http.Client
}

func NewErrorResponse(pages map[string][]int) *ErrorResponse {
	errors := &ErrorResponse{
		pages: make(map[int]*ErrorPage),
	}

	// aggressively timeout connections
	errors.client = &http.Client{
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 2 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 2 * time.Second,
		},
		Timeout: 5 * time.Second,
	}

	if pages != nil {
		errors.Update(pages)
	}
	return errors
}

// Get the ErrorPage, returning nil if the page was incomplete.
// We permanently cache error pages and headers once we've seen them.
func (e *ErrorResponse) Get(code int) *ErrorPage {
	e.Lock()
	page, ok := e.pages[code]
	e.Unlock()

	if !ok {
		// this is a code we don't handle
		return nil
	}

	body := page.Body()
	if body != nil {
		return page
	}

	// we haven't successfully fetched this error
	e.fetch(page)
	return page
}

func (e *ErrorResponse) fetch(page *ErrorPage) {
	log.Debugf("Fetching error page from %s", page.Location)
	resp, err := e.client.Get(page.Location)
	if err != nil {
		log.Warnf("Could not fetch %s: %s", page.Location, err.Error())
		return
	}
	defer resp.Body.Close()

	// If the StatusCode matches any of our registered codes, it's OK
	for _, code := range page.StatusCodes {
		if resp.StatusCode == code {
			resp.StatusCode = http.StatusOK
			break
		}
	}

	if resp.StatusCode != http.StatusOK {
		log.Warnf("Server returned %d when fetching %s", resp.StatusCode, page.Location)
		return
	}

	header := make(map[string][]string)
	for _, key := range ErrorHeaders {
		if hdr, ok := resp.Header[key]; ok {
			header[key] = hdr
		}
	}
	// set the headers along with the body below

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Warnf("Error reading response from %s: %s", page.Location, err.Error())
		return
	}

	if len(body) > 0 {
		page.SetHeader(header)
		page.SetBody(body)
		return
	}
	log.Warnf("Empty response from %s", page.Location)
}

// This replaces all existing ErrorPages
func (e *ErrorResponse) Update(pages map[string][]int) {
	e.Lock()
	defer e.Unlock()

	e.pages = make(map[int]*ErrorPage)

	for loc, codes := range pages {
		page := &ErrorPage{
			StatusCodes: codes,
			Location:    loc,
		}

		for _, code := range codes {
			e.pages[code] = page
		}
		go e.fetch(page)
	}
}

func (e *ErrorResponse) CheckResponse(pr *ProxyRequest) bool {

	errPage := e.Get(pr.Response.StatusCode)
	if errPage != nil {
		// load the cached headers
		header := pr.ResponseWriter.Header()
		for key, val := range errPage.Header() {
			header[key] = val
		}

		pr.ResponseWriter.WriteHeader(pr.Response.StatusCode)
		pr.ResponseWriter.Write(errPage.Body())
		return false
	}

	return true
}

func logProxyRequest(pr *ProxyRequest) bool {
	// TODO: we may to be able to switch this off
	if pr == nil || pr.Request == nil {
		return true
	}

	var id, method, clientIP, url, backend, agent string
	var status int

	duration := pr.FinishTime.Sub(pr.StartTime)

	id = pr.Request.Header.Get("X-Request-Id")
	method = pr.Request.Method
	url = pr.Request.Host + pr.Request.RequestURI
	agent = pr.Request.UserAgent()
	status = pr.Response.StatusCode

	clientIP = pr.Request.Header.Get("X-Forwarded-For")
	if clientIP == "" {
		clientIP = pr.Request.RemoteAddr
	}

	if pr.Response != nil && pr.Response.Request != nil && pr.Response.Request.URL != nil {
		backend = pr.Response.Request.URL.Host
	}

	err := fmt.Sprintf("%v", pr.ProxyError)

	fmtStr := "id=%s method=%s clientIp=%s url=%s backend=%s status=%d duration=%s agent=%s, err=%s"

	log.Printf(fmtStr, id, method, clientIP, url, backend, status, duration, agent, err)
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/yunge/sphinx"
        "fmt"
"time"
"strconv"
	"github.com/yvasiyarov/newrelic_platform_go"
)

const (
    MIN_PAUSE_TIME = 30
)

type SphinxStatusData map[string]string
type MetricsDataSource struct {
	SphinxHost        string
	Port              int
	ConnectionTimeout int

        PreviousData SphinxStatusData
        LastData     SphinxStatusData
        LastUpdateTime time.Time
}

func NewMetricsDataSource(sphinxHost string, port int, connectionTimeout int) *MetricsDataSource {
	if port == 0 {
		port = 9312
	}
	ds := &MetricsDataSource{
		SphinxHost:        sphinxHost,
		Port:              port,
		ConnectionTimeout: connectionTimeout,
	}
	return ds
}

func (ds *MetricsDataSource) GetData(key string) (float64, error) {
    if err := ds.CheckAndUpdateData(); err != nil {
        return 0, err
    }

    prev, last, err := ds.GetOriginalData(key)

    if err != nil {
        return 0, err
    }
    return last - prev, nil
}

func (ds *MetricsDataSource) GetOriginalData(key string) (float64, float64, error) {
    previousValue, ok := ds.PreviousData[key]
    if !ok {
        return 0, 0, fmt.Errorf("Can not get data from source \n")
    }
    currentValue, ok := ds.LastData[key]
    if !ok {
        return 0, 0, fmt.Errorf("Can not get data from source \n")
    }

    //some metric calculation can be turned off by sphinx settings
    if previousValue == "OFF" || currentValue == "OFF" {
        return 0, 0, nil
    }

    previousValueConverted, err := strconv.ParseFloat(previousValue, 64)
    if err != nil {
        return 0, 0, fmt.Errorf("Can not convert previous value of %s to int \n", key)
    }
    currentValueConverted, err := strconv.ParseFloat(currentValue, 64)
    if err != nil {
        return 0, 0, fmt.Errorf("Can not convert current value of %s to int \n", key)
    }

    return previousValueConverted, currentValueConverted, nil
}

func (ds *MetricsDataSource) CheckAndUpdateData() error {
        startTime := time.Now()
        if startTime.Sub(ds.LastUpdateTime) > time.Second * MIN_PAUSE_TIME {
            newData, err := ds.QueryData()
            if err != nil {
                return err
            }

            if ds.PreviousData == nil {
                ds.PreviousData = newData
            } else {
                ds.PreviousData = ds.LastData
            }
            ds.LastData = newData
        }

        // check uptime
        //If uptime is less then in previous run - then server were restarted
        if prev, last, err := ds.GetOriginalData("uptime"); err != nil {
            return err
        } else {
            if last < prev {
               ds.PreviousData = ds.LastData
            }
        }
        return nil
}

func (ds *MetricsDataSource) QueryData() (SphinxStatusData, error) {
	client := sphinx.NewClient().SetServer(ds.SphinxHost, ds.Port)

	if ds.ConnectionTimeout != 0 {
		client.SetConnectTimeout(ds.ConnectionTimeout)
	}
	if err := client.Error(); err != nil {
		return nil, err
	}

	defer client.Close()
	status, err := client.Status()
	if err != nil {
		return nil, err
	}

        data := make(SphinxStatusData, len(status))
	for _, row := range status {
		data[row[0]] = row[1]
	}

	return data, nil
}

type Metrica struct {
	Name       string
	Units      string
	DataKey    string
	DataSource *MetricsDataSource
}

func (metrica *Metrica) GetName() string {
	return metrica.Name
}
func (metrica *Metrica) GetUnits() string {
	return metrica.Units
}
func (metrica *Metrica) GetValue() (float64, error) {
    return metrica.DataSource.GetData(metrica.DataKey)
}

func AddMetrcas(component newrelic_platform_go.IComponent, dataSource *MetricsDataSource) {
    keys := []string{
       "connections", "maxed_out", "command_search", "command_excerpt", "command_update", "command_keywords",
       "command_persist", "command_status", "command_flushattrs", "agent_connect", "agent_retry", "queries",
       "dist_queries", "query_wall", "query_cpu", "dist_wall", "dist_local", "dist_wait", "query_reads",
       "query_readkb", "query_readtime", "avg_query_wall", "avg_query_cpu", "avg_dist_wall", "avg_dist_local",
       "avg_dist_wait", "avg_query_reads", "avg_query_readkb", "avg_query_readtime",
    }

    names := []string{
       "Connections", "Maxed out", "Command search", "Command excerpt", "Command update", "Command keywords",
       "Command persist", "Command status", "Command flushattrs", "Agent connect", "Agent retry", "Queries",
       "Dist queries", "Query wall", "Query cpu", "Dist wall", "Dist local", "Dist wait", "Query reads",
       "Query readkb", "Query readtime", "Avg query wall", "Avg query cpu", "Avg dist wall", "Avg dist local",
       "Avg dist wait", "Avg query reads", "Avg query readkb", "Avg query readtime",
    }
    for i, key := range keys {
        m := &Metrica{
            DataKey: key,
            Name: names[i],
            Units: "val/second",
            DataSource: dataSource,
        }
        component.AddMetrica(m)
    }
}

func main() {
        plugin := newrelic_platform_go.NewNewrelicPlugin("0.0.1", "7bceac019c7dcafae1ef95be3e3a3ff8866de246", 60)
        component := newrelic_platform_go.NewPluginComponent("Sphinx component", "com.github.yvasiyarov.Sphinx")
        plugin.AddComponent(component)

	ds := NewMetricsDataSource("web-d5.butik.ru", 0, 0)
        AddMetrcas(component, ds)

        plugin.Harvest()
        plugin.Harvest()
}
<commit_msg>It finally worked in console mode<commit_after>package main

import (
	"github.com/yunge/sphinx"
        "fmt"
"time"
"strconv"
	"github.com/yvasiyarov/newrelic_platform_go"
)

const (
    MIN_PAUSE_TIME = 30
)

type SphinxStatusData map[string]string
type MetricsDataSource struct {
	SphinxHost        string
	Port              int
	ConnectionTimeout int

        PreviousData SphinxStatusData
        LastData     SphinxStatusData
        LastUpdateTime time.Time
}

func NewMetricsDataSource(sphinxHost string, port int, connectionTimeout int) *MetricsDataSource {
	if port == 0 {
		port = 9312
	}
	ds := &MetricsDataSource{
		SphinxHost:        sphinxHost,
		Port:              port,
		ConnectionTimeout: connectionTimeout,
	}
	return ds
}

func (ds *MetricsDataSource) GetData(key string) (float64, error) {
    if err := ds.CheckAndUpdateData(); err != nil {
        return 0, err
    }

    prev, last, err := ds.GetOriginalData(key)

    if err != nil {
        return 0, err
    }
    return last - prev, nil
}

func (ds *MetricsDataSource) GetOriginalData(key string) (float64, float64, error) {
    previousValue, ok := ds.PreviousData[key]
    if !ok {
        return 0, 0, fmt.Errorf("Can not get data from source \n")
    }
    currentValue, ok := ds.LastData[key]
    if !ok {
        return 0, 0, fmt.Errorf("Can not get data from source \n")
    }

    //some metric calculation can be turned off by sphinx settings
    if previousValue == "OFF" || currentValue == "OFF" {
        return 0, 0, nil
    }

    previousValueConverted, err := strconv.ParseFloat(previousValue, 64)
    if err != nil {
        return 0, 0, fmt.Errorf("Can not convert previous value of %s to int \n", key)
    }
    currentValueConverted, err := strconv.ParseFloat(currentValue, 64)
    if err != nil {
        return 0, 0, fmt.Errorf("Can not convert current value of %s to int \n", key)
    }

    return previousValueConverted, currentValueConverted, nil
}

func (ds *MetricsDataSource) CheckAndUpdateData() error {
        startTime := time.Now()
        if startTime.Sub(ds.LastUpdateTime) > time.Second * MIN_PAUSE_TIME {
            newData, err := ds.QueryData()
            if err != nil {
                return err
            }

            if ds.PreviousData == nil {
                ds.PreviousData = newData
            } else {
                ds.PreviousData = ds.LastData
            }
            ds.LastData = newData
            ds.LastUpdateTime = startTime
        }

        // check uptime
        //If uptime is less then in previous run - then server were restarted
        if prev, last, err := ds.GetOriginalData("uptime"); err != nil {
            return err
        } else {
            if last < prev {
               ds.PreviousData = ds.LastData
            }
        }
        return nil
}

func (ds *MetricsDataSource) QueryData() (SphinxStatusData, error) {
	client := sphinx.NewClient().SetServer(ds.SphinxHost, ds.Port)

	if ds.ConnectionTimeout != 0 {
		client.SetConnectTimeout(ds.ConnectionTimeout)
	}
	if err := client.Error(); err != nil {
		return nil, err
	}

	defer client.Close()
	status, err := client.Status()
	if err != nil {
		return nil, err
	}

        data := make(SphinxStatusData, len(status))
	for _, row := range status {
		data[row[0]] = row[1]
	}

	return data, nil
}

type Metrica struct {
	Name       string
	Units      string
	DataKey    string
	DataSource *MetricsDataSource
}

func (metrica *Metrica) GetName() string {
	return metrica.Name
}
func (metrica *Metrica) GetUnits() string {
	return metrica.Units
}
func (metrica *Metrica) GetValue() (float64, error) {
    return metrica.DataSource.GetData(metrica.DataKey)
}

func AddMetrcas(component newrelic_platform_go.IComponent, dataSource *MetricsDataSource) {
    metricas := []*Metrica{
        &Metrica{
            DataKey: "queries",
            Name: "Queries",
            Units: "Queries/second",
        },
        &Metrica{
            DataKey: "connections",
            Name: "Connections",
            Units: "connections/second",
        },
        &Metrica{
            DataKey: "maxed_out",
            Name: "Maxed out",
            Units: "connections/second",
        },
        &Metrica{
            DataKey: "command_search",
            Name: "Command search",
            Units: "command/second",
        },
        &Metrica{
            DataKey: "command_excerpt",
            Name: "Command excerpt",
            Units: "command/second",
        },
        &Metrica{
            DataKey: "command_update",
            Name: "Command update",
            Units: "command/second",
        },
        &Metrica{
            DataKey: "command_keywords",
            Name: "Command keywords",
            Units: "command/second",
        },
        &Metrica{
            DataKey: "command_persist",
            Name: "Command persist",
            Units: "command/second",
        },
        &Metrica{
            DataKey: "command_flushattrs",
            Name: "Command flushattrs",
            Units: "command/second",
        },
    }
    for _, m := range metricas {
        m.DataSource = dataSource
        component.AddMetrica(m)
    }
}

func main() {
        plugin := newrelic_platform_go.NewNewrelicPlugin("0.0.1", "7bceac019c7dcafae1ef95be3e3a3ff8866de246", 60)
        component := newrelic_platform_go.NewPluginComponent("Sphinx component", "com.github.yvasiyarov.Sphinx")
        plugin.AddComponent(component)

	ds := NewMetricsDataSource("web-d5.butik.ru", 0, 0)
        AddMetrcas(component, ds)

        plugin.Verbose = true

        plugin.Run()
        //plugin.Harvest()
}
<|endoftext|>


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

// This package implements a basic LISP interpretor for embedding in a go program for scripting.
// This file contains the list access primitive functions.

package golisp

import (
	"fmt"
)

func RegisterListFunctionsPrimitives() {
	MakePrimitiveFunction("map", ">=2", MapImpl)
	MakePrimitiveFunction("for-each", ">=2", ForEachImpl)
	MakePrimitiveFunction("reduce", "3", ReduceImpl)
	MakePrimitiveFunction("filter", "2", FilterImpl)
	MakePrimitiveFunction("remove", "2", RemoveImpl)
	MakePrimitiveFunction("memq", "2", MemqImpl)
	MakePrimitiveFunction("memp", "2", FindTailImpl)
	MakePrimitiveFunction("find-tail", "2", FindTailImpl)
	MakePrimitiveFunction("find", "2", FindImpl)
}

func intMin(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}

func MapImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError(fmt.Sprintf("map needs a function as its first argument, but got %s.", String(f)), env)
		return
	}

	var collections []*Data = make([]*Data, 0, Length(args)-1)
	var loopCount = 30000
	var col *Data
	for a := Cdr(args); NotNilP(a); a = Cdr(a) {
		col = Car(a)
		if !ListP(col) {
			err = ProcessError(fmt.Sprintf("map needs lists as its other arguments, but got %s.", String(col)), env)
			return
		}
		collections = append(collections, col)
		loopCount = intMin(loopCount, Length(col))
	}

	var d []*Data = make([]*Data, 0, loopCount)
	var v *Data
	var a *Data
	for index := 1; index <= loopCount; index++ {
		mapArgs := make([]*Data, 0, len(collections))
		for _, mapArgCollection := range collections {
			a = Nth(mapArgCollection, index)
			mapArgs = append(mapArgs, a)
		}
		v, err = ApplyWithoutEval(f, ArrayToList(mapArgs), env)
		if err != nil {
			return
		}
		d = append(d, v)
	}

	return ArrayToList(d), nil
}

func ForEachImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError(fmt.Sprintf("foreach needs a function as its first argument, but got %s.", String(f)), env)
		return
	}

	var collections []*Data = make([]*Data, 0, Length(args)-1)
	var loopCount = 30000
	var col *Data
	for a := Cdr(args); NotNilP(a); a = Cdr(a) {
		col = Car(a)
		if !ListP(col) {
			err = ProcessError(fmt.Sprintf("foreach needs lists as its other arguments, but got %s.", String(col)), env)
			return
		}
		collections = append(collections, col)
		loopCount = intMin(loopCount, Length(col))
	}

	var a *Data
	for index := 1; index <= loopCount; index++ {
		mapArgs := make([]*Data, 0, len(collections))
		for _, mapArgCollection := range collections {
			a = Nth(mapArgCollection, index)
			mapArgs = append(mapArgs, a)
		}
		_, err = ApplyWithoutEval(f, ArrayToList(mapArgs), env)
		if err != nil {
			return
		}
	}

	return nil, nil
}

func ReduceImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError("reduce needs a function as its first argument", env)
		return
	}

	initial := Second(args)
	col := Third(args)

	if !ListP(col) {
		err = ProcessError("map needs a list as its third argument", env)
		return
	}

	if Length(col) == 0 {
		return initial, nil
	}

	if Length(col) == 1 {
		return Car(col), nil
	}

	result = Car(col)
	for c := Cdr(col); NotNilP(c); c = Cdr(c) {
		result, err = ApplyWithoutEval(f, InternalMakeList(result, Car(c)), env)
		if err != nil {
			return
		}
	}

	return
}

func FilterImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError(fmt.Sprintf("filter needs a function as its first argument, but got %s.", String(f)), env)
		return
	}

	col := Second(args)
	if !ListP(col) {
		err = ProcessError(fmt.Sprintf("filter needs a list as its second argument, but got %s.", String(col)), env)
		return
	}

	var d []*Data = make([]*Data, 0, Length(col))
	var v *Data
	for c := col; NotNilP(c); c = Cdr(c) {
		v, err = ApplyWithoutEval(f, Cons(Car(c), nil), env)
		if err != nil {
			return
		}
		if !BooleanP(v) {
			err = ProcessError("filter needs a predicate function as its first argument.", env)
			return
		}

		if BooleanValue(v) {
			d = append(d, Car(c))
		}
	}

	return ArrayToList(d), nil
}

func RemoveImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError(fmt.Sprintf("remove needs a function as its first argument, but got %s.", String(f)), env)
		return
	}

	col := Second(args)
	if !ListP(col) {
		err = ProcessError(fmt.Sprintf("remove needs a list as its second argument, but got %s.", String(col)), env)
		return
	}

	var d []*Data = make([]*Data, 0, Length(col))
	var v *Data
	for c := col; NotNilP(c); c = Cdr(c) {
		v, err = ApplyWithoutEval(f, Cons(Car(c), nil), env)
		if err != nil {
			return
		}
		if !BooleanP(v) {
			err = ProcessError("remove needs a predicate function as its first argument.", env)
			return
		}

		if !BooleanValue(v) {
			d = append(d, Car(c))
		}
	}

	return ArrayToList(d), nil
}

func MemqImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	key := First(args)

	l := Second(args)

	for c := l; NotNilP(c); c = Cdr(c) {
		if IsEqual(key, Car(c)) {
			return c, nil
		}
	}

	return LispFalse, nil
}

func FindTailImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError("find-tasil/memp needs a function as its first argument", env)
		return
	}

	l := Second(args)

	var found *Data
	for c := l; NotNilP(c); c = Cdr(c) {
		found, err = ApplyWithoutEval(f, InternalMakeList(Car(c)), env)
		if BooleanValue(found) {
			return c, nil
		}
	}

	return LispFalse, nil
}

func FindImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError("find needs a function as its first argument", env)
		return
	}

	l := Second(args)

	var found *Data
	for c := l; NotNilP(c); c = Cdr(c) {
		found, err = ApplyWithoutEval(f, InternalMakeList(Car(c)), env)
		if BooleanValue(found) {
			return Car(c), nil
		}
	}

	return LispFalse, nil
}
<commit_msg>Fix typo<commit_after>// Copyright 2014 SteelSeries ApS.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This package implements a basic LISP interpretor for embedding in a go program for scripting.
// This file contains the list access primitive functions.

package golisp

import (
	"fmt"
)

func RegisterListFunctionsPrimitives() {
	MakePrimitiveFunction("map", ">=2", MapImpl)
	MakePrimitiveFunction("for-each", ">=2", ForEachImpl)
	MakePrimitiveFunction("reduce", "3", ReduceImpl)
	MakePrimitiveFunction("filter", "2", FilterImpl)
	MakePrimitiveFunction("remove", "2", RemoveImpl)
	MakePrimitiveFunction("memq", "2", MemqImpl)
	MakePrimitiveFunction("memp", "2", FindTailImpl)
	MakePrimitiveFunction("find-tail", "2", FindTailImpl)
	MakePrimitiveFunction("find", "2", FindImpl)
}

func intMin(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}

func MapImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError(fmt.Sprintf("map needs a function as its first argument, but got %s.", String(f)), env)
		return
	}

	var collections []*Data = make([]*Data, 0, Length(args)-1)
	var loopCount = 30000
	var col *Data
	for a := Cdr(args); NotNilP(a); a = Cdr(a) {
		col = Car(a)
		if !ListP(col) {
			err = ProcessError(fmt.Sprintf("map needs lists as its other arguments, but got %s.", String(col)), env)
			return
		}
		collections = append(collections, col)
		loopCount = intMin(loopCount, Length(col))
	}

	var d []*Data = make([]*Data, 0, loopCount)
	var v *Data
	var a *Data
	for index := 1; index <= loopCount; index++ {
		mapArgs := make([]*Data, 0, len(collections))
		for _, mapArgCollection := range collections {
			a = Nth(mapArgCollection, index)
			mapArgs = append(mapArgs, a)
		}
		v, err = ApplyWithoutEval(f, ArrayToList(mapArgs), env)
		if err != nil {
			return
		}
		d = append(d, v)
	}

	return ArrayToList(d), nil
}

func ForEachImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError(fmt.Sprintf("foreach needs a function as its first argument, but got %s.", String(f)), env)
		return
	}

	var collections []*Data = make([]*Data, 0, Length(args)-1)
	var loopCount = 30000
	var col *Data
	for a := Cdr(args); NotNilP(a); a = Cdr(a) {
		col = Car(a)
		if !ListP(col) {
			err = ProcessError(fmt.Sprintf("foreach needs lists as its other arguments, but got %s.", String(col)), env)
			return
		}
		collections = append(collections, col)
		loopCount = intMin(loopCount, Length(col))
	}

	var a *Data
	for index := 1; index <= loopCount; index++ {
		mapArgs := make([]*Data, 0, len(collections))
		for _, mapArgCollection := range collections {
			a = Nth(mapArgCollection, index)
			mapArgs = append(mapArgs, a)
		}
		_, err = ApplyWithoutEval(f, ArrayToList(mapArgs), env)
		if err != nil {
			return
		}
	}

	return nil, nil
}

func ReduceImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError("reduce needs a function as its first argument", env)
		return
	}

	initial := Second(args)
	col := Third(args)

	if !ListP(col) {
		err = ProcessError("map needs a list as its third argument", env)
		return
	}

	if Length(col) == 0 {
		return initial, nil
	}

	if Length(col) == 1 {
		return Car(col), nil
	}

	result = Car(col)
	for c := Cdr(col); NotNilP(c); c = Cdr(c) {
		result, err = ApplyWithoutEval(f, InternalMakeList(result, Car(c)), env)
		if err != nil {
			return
		}
	}

	return
}

func FilterImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError(fmt.Sprintf("filter needs a function as its first argument, but got %s.", String(f)), env)
		return
	}

	col := Second(args)
	if !ListP(col) {
		err = ProcessError(fmt.Sprintf("filter needs a list as its second argument, but got %s.", String(col)), env)
		return
	}

	var d []*Data = make([]*Data, 0, Length(col))
	var v *Data
	for c := col; NotNilP(c); c = Cdr(c) {
		v, err = ApplyWithoutEval(f, Cons(Car(c), nil), env)
		if err != nil {
			return
		}
		if !BooleanP(v) {
			err = ProcessError("filter needs a predicate function as its first argument.", env)
			return
		}

		if BooleanValue(v) {
			d = append(d, Car(c))
		}
	}

	return ArrayToList(d), nil
}

func RemoveImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError(fmt.Sprintf("remove needs a function as its first argument, but got %s.", String(f)), env)
		return
	}

	col := Second(args)
	if !ListP(col) {
		err = ProcessError(fmt.Sprintf("remove needs a list as its second argument, but got %s.", String(col)), env)
		return
	}

	var d []*Data = make([]*Data, 0, Length(col))
	var v *Data
	for c := col; NotNilP(c); c = Cdr(c) {
		v, err = ApplyWithoutEval(f, Cons(Car(c), nil), env)
		if err != nil {
			return
		}
		if !BooleanP(v) {
			err = ProcessError("remove needs a predicate function as its first argument.", env)
			return
		}

		if !BooleanValue(v) {
			d = append(d, Car(c))
		}
	}

	return ArrayToList(d), nil
}

func MemqImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	key := First(args)

	l := Second(args)

	for c := l; NotNilP(c); c = Cdr(c) {
		if IsEqual(key, Car(c)) {
			return c, nil
		}
	}

	return LispFalse, nil
}

func FindTailImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError("find-tail/memp needs a function as its first argument", env)
		return
	}

	l := Second(args)

	var found *Data
	for c := l; NotNilP(c); c = Cdr(c) {
		found, err = ApplyWithoutEval(f, InternalMakeList(Car(c)), env)
		if BooleanValue(found) {
			return c, nil
		}
	}

	return LispFalse, nil
}

func FindImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	f := First(args)
	if !FunctionP(f) {
		err = ProcessError("find needs a function as its first argument", env)
		return
	}

	l := Second(args)

	var found *Data
	for c := l; NotNilP(c); c = Cdr(c) {
		found, err = ApplyWithoutEval(f, InternalMakeList(Car(c)), env)
		if BooleanValue(found) {
			return Car(c), nil
		}
	}

	return LispFalse, nil
}
<|endoftext|>


<commit_before>package primitive

import (
	"math"
	"math/rand"

	"github.com/fogleman/gg"
)

type Rectangle struct {
	W, H   int
	X1, Y1 int
	X2, Y2 int
}

func NewRandomRectangle(w, h int) *Rectangle {
	x1 := rand.Intn(w)
	y1 := rand.Intn(h)
	x2 := rand.Intn(w)
	y2 := rand.Intn(h)
	return &Rectangle{w, h, x1, y1, x2, y2}
}

func (r *Rectangle) Draw(dc *gg.Context) {
	x1, y1 := r.X1, r.Y1
	x2, y2 := r.X2, r.Y2
	if x1 > x2 {
		x1, x2 = x2, x1
	}
	if y1 > y2 {
		y1, y2 = y2, y1
	}
	dc.DrawRectangle(float64(x1), float64(y1), float64(x2-x1+1), float64(y2-y1+1))
	dc.Fill()
}

func (r *Rectangle) Copy() Shape {
	a := *r
	return &a
}

func (r *Rectangle) Mutate() {
	switch rand.Intn(2) {
	case 0:
		r.X1 = clampInt(r.X1+rand.Intn(21)-10, 0, r.W-1)
		r.Y1 = clampInt(r.Y1+rand.Intn(21)-10, 0, r.H-1)
	case 1:
		r.X2 = clampInt(r.X2+rand.Intn(21)-10, 0, r.W-1)
		r.Y2 = clampInt(r.Y2+rand.Intn(21)-10, 0, r.H-1)
	}
}

func (r *Rectangle) Rasterize() []Scanline {
	x1, y1 := r.X1, r.Y1
	x2, y2 := r.X2, r.Y2
	if x1 > x2 {
		x1, x2 = x2, x1
	}
	if y1 > y2 {
		y1, y2 = y2, y1
	}
	lines := make([]Scanline, y2-y1+1)
	i := 0
	for y := y1; y <= y2; y++ {
		lines[i] = Scanline{y, x1, x2}
		i++
	}
	return lines
}

type RotatedRectangle struct {
	W, H   int
	X, Y   int
	Sx, Sy int
	Angle  int
}

func NewRandomRotatedRectangle(w, h int) *RotatedRectangle {
	x := rand.Intn(w)
	y := rand.Intn(h)
	sx := rand.Intn(w / 2)
	sy := rand.Intn(h / 2)
	a := rand.Intn(360)
	r := &RotatedRectangle{w, h, x, y, sx, sy, a}
	r.Mutate()
	return r
}

func (r *RotatedRectangle) Draw(dc *gg.Context) {
	sx, sy := float64(r.Sx), float64(r.Sy)
	dc.Push()
	dc.Translate(float64(r.X), float64(r.Y))
	dc.Rotate(radians(float64(r.Angle)))
	dc.DrawRectangle(-sx/2, -sy/2, sx, sy)
	dc.Pop()
	dc.Fill()
}

func (r *RotatedRectangle) Copy() Shape {
	a := *r
	return &a
}

func (r *RotatedRectangle) Mutate() {
	for {
		switch rand.Intn(2) {
		case 0:
			r.X = clampInt(r.X+rand.Intn(21)-10, 0, r.W-1)
			r.Y = clampInt(r.Y+rand.Intn(21)-10, 0, r.H-1)
		case 1:
			r.Sx = clampInt(r.Sx+rand.Intn(21)-10, 0, r.W-1)
			r.Sy = clampInt(r.Sy+rand.Intn(21)-10, 0, r.H-1)
		case 2:
			r.Angle = r.Angle + rand.Intn(21) - 10
		}
		if r.Valid() {
			break
		}
	}
}

func (r *RotatedRectangle) Valid() bool {
	a, b := r.Sx, r.Sy
	if a < b {
		a, b = b, a
	}
	aspect := float64(a) / float64(b)
	return aspect <= 5
}

func (r *RotatedRectangle) Rasterize() []Scanline {
	sx, sy := float64(r.Sx), float64(r.Sy)
	angle := radians(float64(r.Angle))
	rx1, ry1 := rotate(-sx/2, -sy/2, angle)
	rx2, ry2 := rotate(sx/2, -sy/2, angle)
	rx3, ry3 := rotate(sx/2, sy/2, angle)
	rx4, ry4 := rotate(-sx/2, sy/2, angle)
	x1, y1 := int(rx1)+r.X, int(ry1)+r.Y
	x2, y2 := int(rx2)+r.X, int(ry2)+r.Y
	x3, y3 := int(rx3)+r.X, int(ry3)+r.Y
	x4, y4 := int(rx4)+r.X, int(ry4)+r.Y
	miny := minInt(y1, minInt(y2, minInt(y3, y4)))
	maxy := maxInt(y1, maxInt(y2, maxInt(y3, y4)))
	n := maxy - miny + 1
	min := make([]int, n)
	max := make([]int, n)
	for i := range min {
		min[i] = r.W
	}
	xs := []int{x1, x2, x3, x4, x1}
	ys := []int{y1, y2, y3, y4, y1}
	// TODO: this could be better probably
	for i := 0; i < 4; i++ {
		x, y := float64(xs[i]), float64(ys[i])
		dx, dy := float64(xs[i+1]-xs[i]), float64(ys[i+1]-ys[i])
		count := int(math.Sqrt(dx*dx+dy*dy)) * 2
		for j := 0; j < count; j++ {
			t := float64(j) / float64(count-1)
			xi := int(x + dx*t)
			yi := int(y+dy*t) - miny
			min[yi] = minInt(min[yi], xi)
			max[yi] = maxInt(max[yi], xi)
		}
	}
	lines := make([]Scanline, 0, n)
	for i := 0; i < n; i++ {
		y := miny + i
		if y < 0 || y >= r.H {
			continue
		}
		a := maxInt(min[i], 0)
		b := minInt(max[i], r.W-1)
		lines = append(lines, Scanline{y, a, b})
	}
	return lines
}
<commit_msg>rotated rect fix<commit_after>package primitive

import (
	"math"
	"math/rand"

	"github.com/fogleman/gg"
)

type Rectangle struct {
	W, H   int
	X1, Y1 int
	X2, Y2 int
}

func NewRandomRectangle(w, h int) *Rectangle {
	x1 := rand.Intn(w)
	y1 := rand.Intn(h)
	x2 := rand.Intn(w)
	y2 := rand.Intn(h)
	return &Rectangle{w, h, x1, y1, x2, y2}
}

func (r *Rectangle) Draw(dc *gg.Context) {
	x1, y1 := r.X1, r.Y1
	x2, y2 := r.X2, r.Y2
	if x1 > x2 {
		x1, x2 = x2, x1
	}
	if y1 > y2 {
		y1, y2 = y2, y1
	}
	dc.DrawRectangle(float64(x1), float64(y1), float64(x2-x1+1), float64(y2-y1+1))
	dc.Fill()
}

func (r *Rectangle) Copy() Shape {
	a := *r
	return &a
}

func (r *Rectangle) Mutate() {
	switch rand.Intn(2) {
	case 0:
		r.X1 = clampInt(r.X1+rand.Intn(21)-10, 0, r.W-1)
		r.Y1 = clampInt(r.Y1+rand.Intn(21)-10, 0, r.H-1)
	case 1:
		r.X2 = clampInt(r.X2+rand.Intn(21)-10, 0, r.W-1)
		r.Y2 = clampInt(r.Y2+rand.Intn(21)-10, 0, r.H-1)
	}
}

func (r *Rectangle) Rasterize() []Scanline {
	x1, y1 := r.X1, r.Y1
	x2, y2 := r.X2, r.Y2
	if x1 > x2 {
		x1, x2 = x2, x1
	}
	if y1 > y2 {
		y1, y2 = y2, y1
	}
	lines := make([]Scanline, y2-y1+1)
	i := 0
	for y := y1; y <= y2; y++ {
		lines[i] = Scanline{y, x1, x2}
		i++
	}
	return lines
}

type RotatedRectangle struct {
	W, H   int
	X, Y   int
	Sx, Sy int
	Angle  int
}

func NewRandomRotatedRectangle(w, h int) *RotatedRectangle {
	x := rand.Intn(w)
	y := rand.Intn(h)
	sx := rand.Intn(w / 2)
	sy := rand.Intn(h / 2)
	a := rand.Intn(360)
	r := &RotatedRectangle{w, h, x, y, sx, sy, a}
	r.Mutate()
	return r
}

func (r *RotatedRectangle) Draw(dc *gg.Context) {
	sx, sy := float64(r.Sx), float64(r.Sy)
	dc.Push()
	dc.Translate(float64(r.X), float64(r.Y))
	dc.Rotate(radians(float64(r.Angle)))
	dc.DrawRectangle(-sx/2, -sy/2, sx, sy)
	dc.Pop()
	dc.Fill()
}

func (r *RotatedRectangle) Copy() Shape {
	a := *r
	return &a
}

func (r *RotatedRectangle) Mutate() {
	switch rand.Intn(3) {
	case 0:
		r.X = clampInt(r.X+rand.Intn(21)-10, 0, r.W-1)
		r.Y = clampInt(r.Y+rand.Intn(21)-10, 0, r.H-1)
	case 1:
		r.Sx = clampInt(r.Sx+rand.Intn(21)-10, 0, r.W-1)
		r.Sy = clampInt(r.Sy+rand.Intn(21)-10, 0, r.H-1)
	case 2:
		r.Angle = r.Angle + rand.Intn(41) - 20
	}
	for !r.Valid() {
		r.Sx = clampInt(r.Sx+rand.Intn(21)-10, 0, r.W-1)
		r.Sy = clampInt(r.Sy+rand.Intn(21)-10, 0, r.H-1)
	}
}

func (r *RotatedRectangle) Valid() bool {
	a, b := r.Sx, r.Sy
	if a < b {
		a, b = b, a
	}
	aspect := float64(a) / float64(b)
	return aspect <= 5
}

func (r *RotatedRectangle) Rasterize() []Scanline {
	sx, sy := float64(r.Sx), float64(r.Sy)
	angle := radians(float64(r.Angle))
	rx1, ry1 := rotate(-sx/2, -sy/2, angle)
	rx2, ry2 := rotate(sx/2, -sy/2, angle)
	rx3, ry3 := rotate(sx/2, sy/2, angle)
	rx4, ry4 := rotate(-sx/2, sy/2, angle)
	x1, y1 := int(rx1)+r.X, int(ry1)+r.Y
	x2, y2 := int(rx2)+r.X, int(ry2)+r.Y
	x3, y3 := int(rx3)+r.X, int(ry3)+r.Y
	x4, y4 := int(rx4)+r.X, int(ry4)+r.Y
	miny := minInt(y1, minInt(y2, minInt(y3, y4)))
	maxy := maxInt(y1, maxInt(y2, maxInt(y3, y4)))
	n := maxy - miny + 1
	min := make([]int, n)
	max := make([]int, n)
	for i := range min {
		min[i] = r.W
	}
	xs := []int{x1, x2, x3, x4, x1}
	ys := []int{y1, y2, y3, y4, y1}
	// TODO: this could be better probably
	for i := 0; i < 4; i++ {
		x, y := float64(xs[i]), float64(ys[i])
		dx, dy := float64(xs[i+1]-xs[i]), float64(ys[i+1]-ys[i])
		count := int(math.Sqrt(dx*dx+dy*dy)) * 2
		for j := 0; j < count; j++ {
			t := float64(j) / float64(count-1)
			xi := int(x + dx*t)
			yi := int(y+dy*t) - miny
			min[yi] = minInt(min[yi], xi)
			max[yi] = maxInt(max[yi], xi)
		}
	}
	lines := make([]Scanline, 0, n)
	for i := 0; i < n; i++ {
		y := miny + i
		if y < 0 || y >= r.H {
			continue
		}
		a := maxInt(min[i], 0)
		b := minInt(max[i], r.W-1)
		lines = append(lines, Scanline{y, a, b})
	}
	return lines
}
<|endoftext|>


<commit_before>package redisc

import (
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/PuerkitoBio/juggler/internal/redistest"
	"github.com/garyburd/redigo/redis"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestClusterRefreshNormalServer(t *testing.T) {
	cmd, port := redistest.StartServer(t, nil, "")
	defer cmd.Process.Kill()

	c := &Cluster{
		StartupNodes: []string{":" + port},
	}
	err := c.Refresh()
	if assert.Error(t, err, "Refresh") {
		assert.Contains(t, err.Error(), "redisc: all nodes failed", "expected error message")
	}
}

func TestClusterRefresh(t *testing.T) {
	fn, ports := redistest.StartCluster(t, nil)
	defer fn()

	for i, p := range ports {
		ports[i] = ":" + p
	}
	c := &Cluster{
		StartupNodes: ports,
	}

	err := c.Refresh()
	if assert.NoError(t, err, "Refresh") {
		var prev string
		pix := -1
		for ix, master := range c.mapping {
			if master != prev || ix == len(c.mapping)-1 {
				prev = master
				t.Logf("%5d: %s\n", ix, master)
				pix++
			}
			if assert.NotEmpty(t, master) {
				split := strings.Index(master, ":")
				assert.Contains(t, ports, master[split:], "expected master")
			}
		}
	}
}

func TestClusterClose(t *testing.T) {
	c := &Cluster{
		StartupNodes: []string{":6379"},
		DialOptions:  []redis.DialOption{redis.DialConnectTimeout(2 * time.Second)},
		CreatePool:   createPool,
	}
	assert.NoError(t, c.Close(), "Close")
	if err := c.Close(); assert.Error(t, err, "Close after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
	if conn := c.Get(); assert.Error(t, conn.Err(), "Get after Close") {
		assert.Contains(t, conn.Err().Error(), "redisc: closed", "expected message")
	}
	if _, err := c.Dial(); assert.Error(t, err, "Dial after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
	if err := c.Refresh(); assert.Error(t, err, "Refresh after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
}

func createPool(addr string, opts ...redis.DialOption) (*redis.Pool, error) {
	return &redis.Pool{
		MaxIdle:     5,
		MaxActive:   10,
		IdleTimeout: time.Minute,
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp", addr, opts...)
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}, nil
}

type redisCmd struct {
	name   string
	args   redis.Args
	resp   interface{} // if resp is of type lenResult, asserts that there is a result at least this long
	errMsg string
}

type lenResult int

func TestCommands(t *testing.T) {
	fn, ports := redistest.StartCluster(t, nil)
	defer fn()

	cmdsPerGroup := map[string][]redisCmd{
		"cluster": {
			{"CLUSTER", redis.Args{"INFO"}, lenResult(10), ""},
			{"READONLY", nil, "OK", ""},
			{"READWRITE", nil, "OK", ""},
			{"CLUSTER", redis.Args{"COUNTKEYSINSLOT", 12345}, int64(0), ""},
			{"CLUSTER", redis.Args{"KEYSLOT", "a"}, int64(15495), ""},
			{"CLUSTER", redis.Args{"NODES"}, lenResult(100), ""},
		},
		"connection": {
			{"AUTH", redis.Args{"pwd"}, nil, "ERR Client sent AUTH, but no password is set"},
			{"ECHO", redis.Args{"a"}, []byte("a"), ""},
			{"PING", nil, "PONG", ""},
			{"SELECT", redis.Args{1}, nil, "ERR SELECT is not allowed in cluster mode"},
			{"QUIT", nil, "OK", ""},
		},
		"hashes": {
			{"HSET", redis.Args{"ha", "f1", "1"}, int64(1), ""},
			{"HLEN", redis.Args{"ha"}, int64(1), ""},
			{"HEXISTS", redis.Args{"ha", "f1"}, int64(1), ""},
			{"HDEL", redis.Args{"ha", "f1", "f2"}, int64(1), ""},
			{"HINCRBY", redis.Args{"hb", "f1", "1"}, int64(1), ""},
			{"HINCRBYFLOAT", redis.Args{"hb", "f2", "0.5"}, []byte("0.5"), ""},
			{"HKEYS", redis.Args{"hb"}, []interface{}{[]byte("f1"), []byte("f2")}, ""},
			{"HMGET", redis.Args{"hb", "f1", "f2"}, []interface{}{[]byte("1"), []byte("0.5")}, ""},
			{"HMSET", redis.Args{"hc", "f1", "a", "f2", "b"}, "OK", ""},
			{"HSET", redis.Args{"ha", "f1", "2"}, int64(1), ""},
			{"HGET", redis.Args{"ha", "f1"}, []byte("2"), ""},
			{"HGETALL", redis.Args{"ha"}, []interface{}{[]byte("f1"), []byte("2")}, ""},
			{"HSETNX", redis.Args{"ha", "f2", "3"}, int64(1), ""},
			//{"HSTRLEN", redis.Args{"hb", "f2"}, int64(3), ""}, // redis 3.2 only
			{"HVALS", redis.Args{"hb"}, []interface{}{[]byte("1"), []byte("0.5")}, ""},
			{"HSCAN", redis.Args{"hb", 0}, lenResult(2), ""},
		},
		"hyperloglog": {
			{"PFADD", redis.Args{"hll", "a", "b", "c"}, int64(1), ""},
			{"PFCOUNT", redis.Args{"hll"}, int64(3), ""},
			{"PFADD", redis.Args{"hll2", "d"}, int64(1), ""},
			{"PFMERGE", redis.Args{"hll", "hll2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
		},
		"keys": {
			{"SET", redis.Args{"k1", "z"}, "OK", ""},
			{"EXISTS", redis.Args{"k1"}, int64(1), ""},
			{"DUMP", redis.Args{"k1"}, lenResult(10), ""},
			{"EXPIRE", redis.Args{"k1", 10}, int64(1), ""},
			{"EXPIREAT", redis.Args{"k1", time.Now().Add(time.Hour).Unix()}, int64(1), ""},
			{"KEYS", redis.Args{"z*"}, []interface{}{}, ""}, // KEYS is supported, but uses a random node and returns keys from that node (undeterministic)
			{"MOVE", redis.Args{"k1", 2}, nil, "ERR MOVE is not allowed in cluster mode"},
			{"PERSIST", redis.Args{"k1"}, int64(1), ""},
			{"PEXPIRE", redis.Args{"k1", 10000}, int64(1), ""},
			{"PEXPIREAT", redis.Args{"k1", time.Now().Add(time.Hour).UnixNano() / int64(time.Millisecond)}, int64(1), ""},
			{"PTTL", redis.Args{"k1"}, lenResult(3500000), ""},
			// RANDOMKEY is not deterministic
			{"RENAME", redis.Args{"k1", "k2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"RENAMENX", redis.Args{"k1", "k2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"SCAN", redis.Args{0}, lenResult(2), ""}, // works, but only for the keys on that random node
			{"TTL", redis.Args{"k1"}, lenResult(3000), ""},
			{"TYPE", redis.Args{"k1"}, "string", ""},
			{"DEL", redis.Args{"k1"}, int64(1), ""},
		},
		"lists": {
			{"LPUSH", redis.Args{"l1", "a", "b", "c"}, int64(3), ""},
			{"LINDEX", redis.Args{"l1", 1}, []byte("b"), ""},
			{"LINSERT", redis.Args{"l1", "BEFORE", "b", "d"}, int64(4), ""},
			{"LLEN", redis.Args{"l1"}, int64(4), ""},
			{"LPOP", redis.Args{"l1"}, []byte("c"), ""},
			{"LPUSHX", redis.Args{"l1", "e"}, int64(4), ""},
			{"LRANGE", redis.Args{"l1", 0, 1}, []interface{}{[]byte("e"), []byte("d")}, ""},
			{"LREM", redis.Args{"l1", 0, "d"}, int64(1), ""},
			{"LSET", redis.Args{"l1", 0, "f"}, "OK", ""},
			{"LTRIM", redis.Args{"l1", 0, 3}, "OK", ""},
			{"RPOP", redis.Args{"l1"}, []byte("a"), ""},
			{"RPOPLPUSH", redis.Args{"l1", "l2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"RPUSH", redis.Args{"l1", "g"}, int64(3), ""},
			{"RPUSH", redis.Args{"l1", "h"}, int64(4), ""},
			{"BLPOP", redis.Args{"l1", 1}, lenResult(2), ""},
			{"BRPOP", redis.Args{"l1", 1}, lenResult(2), ""},
			{"BRPOPLPUSH", redis.Args{"l1", "l2", 1}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
		},
	}

	for i, p := range ports {
		ports[i] = ":" + p
	}
	c := &Cluster{
		StartupNodes: ports,
		DialOptions:  []redis.DialOption{redis.DialConnectTimeout(2 * time.Second)},
		CreatePool:   createPool,
	}
	require.NoError(t, c.Refresh(), "Refresh")

	var wg sync.WaitGroup
	wg.Add(len(cmdsPerGroup))
	for _, cmds := range cmdsPerGroup {
		go func(cmds []redisCmd) {
			defer wg.Done()
			runCommands(t, c, cmds)
		}(cmds)
	}
	wg.Wait()
}

func runCommands(t *testing.T, c *Cluster, cmds []redisCmd) {
	for _, cmd := range cmds {
		conn := c.Get()
		res, err := conn.Do(cmd.name, cmd.args...)
		if cmd.errMsg != "" {
			if assert.Error(t, err, cmd.name) {
				assert.Contains(t, err.Error(), cmd.errMsg, cmd.name)
			}
		} else {
			assert.NoError(t, err, cmd.name)
			if lr, ok := cmd.resp.(lenResult); ok {
				switch res := res.(type) {
				case []byte:
					assert.True(t, len(res) >= int(lr), "result has at least %d bytes, has %d", lr, len(res))
				case []interface{}:
					assert.Equal(t, int(lr), len(res), "result array has %d items, has %d", lr, len(res))
				case int64:
					assert.True(t, res >= int64(lr), "result is at least %d, is %d", lr, res)
				default:
					t.Errorf("unexpected result type %T", res)
				}
			} else {
				if !assert.Equal(t, cmd.resp, res, cmd.name) {
					t.Logf("%T vs %T", cmd.resp, res)
				}
			}
		}
		require.NoError(t, conn.Close(), "Close")
	}
}
<commit_msg>redisc: test pubsub commands<commit_after>package redisc

import (
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/PuerkitoBio/juggler/internal/redistest"
	"github.com/garyburd/redigo/redis"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestClusterRefreshNormalServer(t *testing.T) {
	cmd, port := redistest.StartServer(t, nil, "")
	defer cmd.Process.Kill()

	c := &Cluster{
		StartupNodes: []string{":" + port},
	}
	err := c.Refresh()
	if assert.Error(t, err, "Refresh") {
		assert.Contains(t, err.Error(), "redisc: all nodes failed", "expected error message")
	}
}

func TestClusterRefresh(t *testing.T) {
	fn, ports := redistest.StartCluster(t, nil)
	defer fn()

	for i, p := range ports {
		ports[i] = ":" + p
	}
	c := &Cluster{
		StartupNodes: ports,
	}

	err := c.Refresh()
	if assert.NoError(t, err, "Refresh") {
		var prev string
		pix := -1
		for ix, master := range c.mapping {
			if master != prev || ix == len(c.mapping)-1 {
				prev = master
				t.Logf("%5d: %s\n", ix, master)
				pix++
			}
			if assert.NotEmpty(t, master) {
				split := strings.Index(master, ":")
				assert.Contains(t, ports, master[split:], "expected master")
			}
		}
	}
}

func TestClusterClose(t *testing.T) {
	c := &Cluster{
		StartupNodes: []string{":6379"},
		DialOptions:  []redis.DialOption{redis.DialConnectTimeout(2 * time.Second)},
		CreatePool:   createPool,
	}
	assert.NoError(t, c.Close(), "Close")
	if err := c.Close(); assert.Error(t, err, "Close after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
	if conn := c.Get(); assert.Error(t, conn.Err(), "Get after Close") {
		assert.Contains(t, conn.Err().Error(), "redisc: closed", "expected message")
	}
	if _, err := c.Dial(); assert.Error(t, err, "Dial after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
	if err := c.Refresh(); assert.Error(t, err, "Refresh after Close") {
		assert.Contains(t, err.Error(), "redisc: closed", "expected message")
	}
}

func createPool(addr string, opts ...redis.DialOption) (*redis.Pool, error) {
	return &redis.Pool{
		MaxIdle:     5,
		MaxActive:   10,
		IdleTimeout: time.Minute,
		Dial: func() (redis.Conn, error) {
			return redis.Dial("tcp", addr, opts...)
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}, nil
}

type redisCmd struct {
	name   string
	args   redis.Args
	resp   interface{} // if resp is of type lenResult, asserts that there is a result at least this long
	errMsg string
}

type lenResult int

func TestCommands(t *testing.T) {
	fn, ports := redistest.StartCluster(t, nil)
	defer fn()

	cmdsPerGroup := map[string][]redisCmd{
		"cluster": {
			{"CLUSTER", redis.Args{"INFO"}, lenResult(10), ""},
			{"READONLY", nil, "OK", ""},
			{"READWRITE", nil, "OK", ""},
			{"CLUSTER", redis.Args{"COUNTKEYSINSLOT", 12345}, int64(0), ""},
			{"CLUSTER", redis.Args{"KEYSLOT", "a"}, int64(15495), ""},
			{"CLUSTER", redis.Args{"NODES"}, lenResult(100), ""},
		},
		"connection": {
			{"AUTH", redis.Args{"pwd"}, nil, "ERR Client sent AUTH, but no password is set"},
			{"ECHO", redis.Args{"a"}, []byte("a"), ""},
			{"PING", nil, "PONG", ""},
			{"SELECT", redis.Args{1}, nil, "ERR SELECT is not allowed in cluster mode"},
			{"QUIT", nil, "OK", ""},
		},
		"hashes": {
			{"HSET", redis.Args{"ha", "f1", "1"}, int64(1), ""},
			{"HLEN", redis.Args{"ha"}, int64(1), ""},
			{"HEXISTS", redis.Args{"ha", "f1"}, int64(1), ""},
			{"HDEL", redis.Args{"ha", "f1", "f2"}, int64(1), ""},
			{"HINCRBY", redis.Args{"hb", "f1", "1"}, int64(1), ""},
			{"HINCRBYFLOAT", redis.Args{"hb", "f2", "0.5"}, []byte("0.5"), ""},
			{"HKEYS", redis.Args{"hb"}, []interface{}{[]byte("f1"), []byte("f2")}, ""},
			{"HMGET", redis.Args{"hb", "f1", "f2"}, []interface{}{[]byte("1"), []byte("0.5")}, ""},
			{"HMSET", redis.Args{"hc", "f1", "a", "f2", "b"}, "OK", ""},
			{"HSET", redis.Args{"ha", "f1", "2"}, int64(1), ""},
			{"HGET", redis.Args{"ha", "f1"}, []byte("2"), ""},
			{"HGETALL", redis.Args{"ha"}, []interface{}{[]byte("f1"), []byte("2")}, ""},
			{"HSETNX", redis.Args{"ha", "f2", "3"}, int64(1), ""},
			//{"HSTRLEN", redis.Args{"hb", "f2"}, int64(3), ""}, // redis 3.2 only
			{"HVALS", redis.Args{"hb"}, []interface{}{[]byte("1"), []byte("0.5")}, ""},
			{"HSCAN", redis.Args{"hb", 0}, lenResult(2), ""},
		},
		"hyperloglog": {
			{"PFADD", redis.Args{"hll", "a", "b", "c"}, int64(1), ""},
			{"PFCOUNT", redis.Args{"hll"}, int64(3), ""},
			{"PFADD", redis.Args{"hll2", "d"}, int64(1), ""},
			{"PFMERGE", redis.Args{"hll", "hll2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
		},
		"keys": {
			{"SET", redis.Args{"k1", "z"}, "OK", ""},
			{"EXISTS", redis.Args{"k1"}, int64(1), ""},
			{"DUMP", redis.Args{"k1"}, lenResult(10), ""},
			{"EXPIRE", redis.Args{"k1", 10}, int64(1), ""},
			{"EXPIREAT", redis.Args{"k1", time.Now().Add(time.Hour).Unix()}, int64(1), ""},
			{"KEYS", redis.Args{"z*"}, []interface{}{}, ""}, // KEYS is supported, but uses a random node and returns keys from that node (undeterministic)
			{"MOVE", redis.Args{"k1", 2}, nil, "ERR MOVE is not allowed in cluster mode"},
			{"PERSIST", redis.Args{"k1"}, int64(1), ""},
			{"PEXPIRE", redis.Args{"k1", 10000}, int64(1), ""},
			{"PEXPIREAT", redis.Args{"k1", time.Now().Add(time.Hour).UnixNano() / int64(time.Millisecond)}, int64(1), ""},
			{"PTTL", redis.Args{"k1"}, lenResult(3500000), ""},
			// RANDOMKEY is not deterministic
			{"RENAME", redis.Args{"k1", "k2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"RENAMENX", redis.Args{"k1", "k2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"SCAN", redis.Args{0}, lenResult(2), ""}, // works, but only for the keys on that random node
			{"TTL", redis.Args{"k1"}, lenResult(3000), ""},
			{"TYPE", redis.Args{"k1"}, "string", ""},
			{"DEL", redis.Args{"k1"}, int64(1), ""},
		},
		"lists": {
			{"LPUSH", redis.Args{"l1", "a", "b", "c"}, int64(3), ""},
			{"LINDEX", redis.Args{"l1", 1}, []byte("b"), ""},
			{"LINSERT", redis.Args{"l1", "BEFORE", "b", "d"}, int64(4), ""},
			{"LLEN", redis.Args{"l1"}, int64(4), ""},
			{"LPOP", redis.Args{"l1"}, []byte("c"), ""},
			{"LPUSHX", redis.Args{"l1", "e"}, int64(4), ""},
			{"LRANGE", redis.Args{"l1", 0, 1}, []interface{}{[]byte("e"), []byte("d")}, ""},
			{"LREM", redis.Args{"l1", 0, "d"}, int64(1), ""},
			{"LSET", redis.Args{"l1", 0, "f"}, "OK", ""},
			{"LTRIM", redis.Args{"l1", 0, 3}, "OK", ""},
			{"RPOP", redis.Args{"l1"}, []byte("a"), ""},
			{"RPOPLPUSH", redis.Args{"l1", "l2"}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
			{"RPUSH", redis.Args{"l1", "g"}, int64(3), ""},
			{"RPUSH", redis.Args{"l1", "h"}, int64(4), ""},
			{"BLPOP", redis.Args{"l1", 1}, lenResult(2), ""},
			{"BRPOP", redis.Args{"l1", 1}, lenResult(2), ""},
			{"BRPOPLPUSH", redis.Args{"l1", "l2", 1}, nil, "CROSSSLOT Keys in request don't hash to the same slot"},
		},
		"pubsub": {
			{"PUBSUB", redis.Args{"NUMPAT"}, lenResult(0), ""},
			{"PUBLISH", redis.Args{"ev1", "a"}, lenResult(0), ""},
		},
	}

	for i, p := range ports {
		ports[i] = ":" + p
	}
	c := &Cluster{
		StartupNodes: ports,
		DialOptions:  []redis.DialOption{redis.DialConnectTimeout(2 * time.Second)},
		CreatePool:   createPool,
	}
	require.NoError(t, c.Refresh(), "Refresh")

	var wg sync.WaitGroup
	// start a goroutine that subscribes and listens to events
	ok := make(chan int)
	done := make(chan int)
	go func() {
		conn, err := c.Dial()
		require.NoError(t, err, "Dial for PubSub")
		psc := redis.PubSubConn{Conn: conn}
		assert.NoError(t, psc.PSubscribe("ev*"), "PSubscribe")
		assert.NoError(t, psc.Subscribe("e1"), "Subscribe")
		ok <- 1

		var received bool
	loop:
		for {
			select {
			case <-done:
				break loop
			default:
			}

			v := psc.Receive()
			switch v := v.(type) {
			case redis.PMessage:
				if !assert.Equal(t, []byte("a"), v.Data, "Received value") {
					t.Logf("%T", v)
				}
				received = true
				break loop
			}
		}

		<-done

		assert.NoError(t, psc.Unsubscribe("e1"), "Unsubscribe")
		assert.NoError(t, psc.PUnsubscribe("ev*"), "PUnsubscribe")
		assert.NoError(t, psc.Close(), "Close for PubSub")
		assert.True(t, received, "Did receive event")
		ok <- 1
	}()

	<-ok
	wg.Add(len(cmdsPerGroup))
	for _, cmds := range cmdsPerGroup {
		go func(cmds []redisCmd) {
			defer wg.Done()
			runCommands(t, c, cmds)
		}(cmds)
	}
	wg.Wait()
	close(done)
	<-ok
}

func runCommands(t *testing.T, c *Cluster, cmds []redisCmd) {
	for _, cmd := range cmds {
		conn := c.Get()
		res, err := conn.Do(cmd.name, cmd.args...)
		if cmd.errMsg != "" {
			if assert.Error(t, err, cmd.name) {
				assert.Contains(t, err.Error(), cmd.errMsg, cmd.name)
			}
		} else {
			assert.NoError(t, err, cmd.name)
			if lr, ok := cmd.resp.(lenResult); ok {
				switch res := res.(type) {
				case []byte:
					assert.True(t, len(res) >= int(lr), "result has at least %d bytes, has %d", lr, len(res))
				case []interface{}:
					assert.Equal(t, int(lr), len(res), "result array has %d items, has %d", lr, len(res))
				case int64:
					assert.True(t, res >= int64(lr), "result is at least %d, is %d", lr, res)
				default:
					t.Errorf("unexpected result type %T", res)
				}
			} else {
				if !assert.Equal(t, cmd.resp, res, cmd.name) {
					t.Logf("%T vs %T", cmd.resp, res)
				}
			}
		}
		require.NoError(t, conn.Close(), "Close")
	}
}
<|endoftext|>


<commit_before>package sftp

import (
	"github.com/stretchr/testify/assert"
	"io"
	"testing"
)

func clientRequestServerPair(t *testing.T) (*Client, *RequestServer) {
	cr, sw := io.Pipe()
	sr, cw := io.Pipe()
	server, err := NewRequestServer(struct {
		io.Reader
		io.WriteCloser
	}{sr, sw})
	if err != nil { t.Fatal(err) }
	go server.Serve()
	client, err := NewClientPipe(cr, cw)
	if err != nil { t.Fatalf("%+v\n", err) }
	return client, server
}

func TestPsRequestCache(t *testing.T) {
	_, ps := clientRequestServerPair(t)
	foo := &Request{Filepath: "foo"}
	bar := &Request{Filepath: "bar"}
	ps.nextRequest(foo)
	ps.nextRequest(bar)
	assert.Len(t, ps.openRequests, 2)
	_foo, ok := ps.getRequest("foo")
	assert.Equal(t, foo, _foo)
	assert.True(t, ok)
	_, ok = ps.getRequest("zed")
	assert.False(t, ok)
	ps.closeRequest("foo")
	ps.closeRequest("bar")
	assert.Len(t, ps.openRequests, 0)
}
<commit_msg>variable name fix<commit_after>package sftp

import (
	"github.com/stretchr/testify/assert"
	"io"
	"testing"
)

func clientRequestServerPair(t *testing.T) (*Client, *RequestServer) {
	cr, sw := io.Pipe()
	sr, cw := io.Pipe()
	server, err := NewRequestServer(struct {
		io.Reader
		io.WriteCloser
	}{sr, sw})
	if err != nil { t.Fatal(err) }
	go server.Serve()
	client, err := NewClientPipe(cr, cw)
	if err != nil { t.Fatalf("%+v\n", err) }
	return client, server
}

func TestPsRequestCache(t *testing.T) {
	_, rs := clientRequestServerPair(t)
	foo := &Request{Filepath: "foo"}
	bar := &Request{Filepath: "bar"}
	rs.nextRequest(foo)
	rs.nextRequest(bar)
	assert.Len(t, rs.openRequests, 2)
	_foo, ok := rs.getRequest("foo")
	assert.Equal(t, foo, _foo)
	assert.True(t, ok)
	_, ok = rs.getRequest("zed")
	assert.False(t, ok)
	rs.closeRequest("foo")
	rs.closeRequest("bar")
	assert.Len(t, rs.openRequests, 0)
}
<|endoftext|>