<commit_before>package emitter

import (
	"fmt"
	"net"
	"net/http"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc/db"
	"github.com/concourse/atc/metric"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

type PrometheusEmitter struct {
	buildsStarted     prometheus.Counter
	buildsFinished    prometheus.Counter
	buildsSucceeded   prometheus.Counter
	buildsErrored     prometheus.Counter
	buildsFailed      prometheus.Counter
	buildsAborted     prometheus.Counter
	buildsFinishedVec *prometheus.CounterVec
	buildDurationsVec *prometheus.HistogramVec

	workerContainers *prometheus.GaugeVec
	workerVolumes    *prometheus.GaugeVec

	httpRequestsDuration *prometheus.HistogramVec

	schedulingFullDuration    *prometheus.GaugeVec
	schedulingLoadingDuration *prometheus.GaugeVec
	schedulingJobDuration     *prometheus.GaugeVec

	dbQueriesTotal prometheus.Counter
	dbConnections  *prometheus.GaugeVec
}

type PrometheusConfig struct {
	BindIP   string `long:"prometheus-bind-ip" description:"IP to listen on to expose Prometheus metrics."`
	BindPort string `long:"prometheus-bind-port" description:"Port to listen on to expose Prometheus metrics."`
}

func init() {
	metric.RegisterEmitter(&PrometheusConfig{})
}

func (config *PrometheusConfig) Description() string { return "Prometheus" }
func (config *PrometheusConfig) IsConfigured() bool {
	return config.BindPort != "" && config.BindIP != ""
}
func (config *PrometheusConfig) bind() string {
	return fmt.Sprintf("%s:%s", config.BindIP, config.BindPort)
}

func (config *PrometheusConfig) NewEmitter() (metric.Emitter, error) {
	// build metrics
	buildsStarted := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "started_total",
		Help:      "Total number of Concourse builds started.",
	})
	prometheus.MustRegister(buildsStarted)

	buildsFinished := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "finished_total",
		Help:      "Total number of Concourse builds finished.",
	})
	prometheus.MustRegister(buildsFinished)

	buildsSucceeded := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "succeeded_total",
		Help:      "Total number of Concourse builds succeeded.",
	})
	prometheus.MustRegister(buildsSucceeded)

	buildsErrored := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "errored_total",
		Help:      "Total number of Concourse builds errored.",
	})
	prometheus.MustRegister(buildsErrored)

	buildsFailed := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "failed_total",
		Help:      "Total number of Concourse builds failed.",
	})
	prometheus.MustRegister(buildsFailed)

	buildsAborted := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "aborted_total",
		Help:      "Total number of Concourse builds aborted.",
	})
	prometheus.MustRegister(buildsAborted)

	buildsFinishedVec := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "concourse",
			Subsystem: "builds",
			Name:      "finished",
			Help:      "Count of builds finished across various dimensions.",
		},
		[]string{"team", "pipeline", "status"},
	)
	prometheus.MustRegister(buildsFinishedVec)
	buildDurationsVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Namespace: "concourse",
			Subsystem: "builds",
			Name:      "duration_seconds",
			Help:      "Build time in seconds",
			Buckets:   []float64{1, 60, 180, 300, 600, 900, 1200, 1800, 2700, 3600, 7200, 18000, 36000},
		},
		[]string{"team", "pipeline"},
	)
	prometheus.MustRegister(buildDurationsVec)

	// worker metrics
	workerContainers := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "workers",
			Name:      "containers",
			Help:      "Number of containers per worker",
		},
		[]string{"worker"},
	)
	prometheus.MustRegister(workerContainers)

	workerVolumes := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "workers",
			Name:      "volumes",
			Help:      "Number of volumes per worker",
		},
		[]string{"worker"},
	)
	prometheus.MustRegister(workerVolumes)

	// http metrics
	httpRequestsDuration := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Namespace: "concourse",
			Subsystem: "http_responses",
			Name:      "duration_seconds",
			Help:      "Response time in seconds",
		},
		[]string{"method", "route"},
	)
	prometheus.MustRegister(httpRequestsDuration)

	// scheduling metrics
	schedulingFullDuration := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "scheduling",
			Name:      "full_duration_seconds",
			Help:      "Last time taken to schedule an entire pipeline.",
		},
		[]string{"pipeline"},
	)
	prometheus.MustRegister(schedulingFullDuration)

	schedulingLoadingDuration := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "scheduling",
			Name:      "loading_duration_seconds",
			Help:      "Last time taken to load version information from the database for a pipeline.",
		},
		[]string{"pipeline"},
	)
	prometheus.MustRegister(schedulingLoadingDuration)

	schedulingJobDuration := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "scheduling",
			Name:      "job_duration_seconds",
			Help:      "Last time taken to calculate the set of valid input versions for a pipeline.",
		},
		[]string{"pipeline"},
	)
	prometheus.MustRegister(schedulingJobDuration)

	dbQueriesTotal := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "db",
		Name:      "queries_total",
		Help:      "Total number of database Concourse database queries",
	})
	prometheus.MustRegister(dbQueriesTotal)

	dbConnections := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "db",
			Name:      "connections",
			Help:      "Current number of concourse database connections",
		},
		[]string{"dbname"},
	)
	prometheus.MustRegister(dbConnections)

	listener, err := net.Listen("tcp", config.bind())
	if err != nil {
		return nil, err
	}

	go http.Serve(listener, promhttp.Handler())

	return &PrometheusEmitter{
		buildsStarted:     buildsStarted,
		buildsFinished:    buildsFinished,
		buildsFinishedVec: buildsFinishedVec,
		buildDurationsVec: buildDurationsVec,
		buildsSucceeded:   buildsSucceeded,
		buildsErrored:     buildsErrored,
		buildsFailed:      buildsFailed,
		buildsAborted:     buildsAborted,

		workerContainers: workerContainers,
		workerVolumes:    workerVolumes,

		httpRequestsDuration: httpRequestsDuration,

		schedulingFullDuration:    schedulingFullDuration,
		schedulingLoadingDuration: schedulingLoadingDuration,
		schedulingJobDuration:     schedulingJobDuration,

		dbQueriesTotal: dbQueriesTotal,
		dbConnections:  dbConnections,
	}, nil
}

// Emit processes incoming metrics.
// In order to provide idiomatic Prometheus metrics, we'll have to convert the various
// Event types (differentiated by the less-than-ideal string Name field) into different
// Prometheus metrics.
func (emitter *PrometheusEmitter) Emit(logger lager.Logger, event metric.Event) {
	switch event.Name {
	case "build started":
		emitter.buildsStarted.Inc()
	case "build finished":
		emitter.buildFinishedMetrics(logger, event)
	case "worker containers":
		emitter.workerContainersMetric(logger, event)
	case "worker volumes":
		emitter.workerVolumesMetric(logger, event)
	case "http response time":
		emitter.httpResponseTimeMetrics(logger, event)
	case "scheduling: full duration (ms)":
		emitter.schedulingMetrics(logger, event)
	case "scheduling: loading versions duration (ms)":
		emitter.schedulingMetrics(logger, event)
	case "scheduling: job duration (ms)":
		emitter.schedulingMetrics(logger, event)
	case "database queries":
		emitter.databaseMetrics(logger, event)
	case "database connections":
		emitter.databaseMetrics(logger, event)
	default:
		// unless we have a specific metric, we do nothing
	}
}

func (emitter *PrometheusEmitter) buildFinishedMetrics(logger lager.Logger, event metric.Event) {
	// concourse_builds_finished_total
	emitter.buildsFinished.Inc()

	// concourse_builds_finished
	team, exists := event.Attributes["team_name"]
	if !exists {
		logger.Error("failed-to-find-team-name-in-event", fmt.Errorf("expected team_name to exist in event.Attributes"))
	}

	pipeline, exists := event.Attributes["pipeline"]
	if !exists {
		logger.Error("failed-to-find-pipeline-in-event", fmt.Errorf("expected pipeline to exist in event.Attributes"))
	}

	buildStatus, exists := event.Attributes["build_status"]
	if !exists {
		logger.Error("failed-to-find-build_status-in-event", fmt.Errorf("expected build_status to exist in event.Attributes"))
	}
	emitter.buildsFinishedVec.WithLabelValues(team, pipeline, buildStatus).Inc()

	// concourse_builds_(aborted|succeeded|failed|errored)_total
	switch buildStatus {
	case string(db.BuildStatusAborted):
		// concourse_builds_aborted_total
		emitter.buildsAborted.Inc()
	case string(db.BuildStatusSucceeded):
		// concourse_builds_succeeded_total
		emitter.buildsSucceeded.Inc()
	case string(db.BuildStatusFailed):
		// concourse_builds_failed_total
		emitter.buildsFailed.Inc()
	case string(db.BuildStatusErrored):
		// concourse_builds_errored_total
		emitter.buildsErrored.Inc()
	}

	// concourse_builds_duration_seconds
	duration, ok := event.Value.(float64)
	if !ok {
		logger.Error("build-finished-event-value-type-mismatch", fmt.Errorf("expected event.Value to be a float64"))
	}
	// seconds are the standard prometheus base unit for time
	duration = duration / 1000
	emitter.buildDurationsVec.WithLabelValues(team, pipeline).Observe(duration)
}

func (emitter *PrometheusEmitter) workerContainersMetric(logger lager.Logger, event metric.Event) {
	worker, exists := event.Attributes["worker"]
	if !exists {
		logger.Error("failed-to-find-worker-in-event", fmt.Errorf("expected worker to exist in event.Attributes"))
	}

	containers, ok := event.Value.(int)
	if !ok {
		logger.Error("worker-volumes-event-value-type-mismatch", fmt.Errorf("expected event.Value to be an int"))
	}

	emitter.workerContainers.WithLabelValues(worker).Set(float64(containers))
}

func (emitter *PrometheusEmitter) workerVolumesMetric(logger lager.Logger, event metric.Event) {
	worker, exists := event.Attributes["worker"]
	if !exists {
		logger.Error("failed-to-find-worker-in-event", fmt.Errorf("expected worker to exist in event.Attributes"))
	}

	volumes, ok := event.Value.(int)
	if !ok {
		logger.Error("worker-volumes-event-value-type-mismatch", fmt.Errorf("expected event.Value to be an int"))
	}

	emitter.workerVolumes.WithLabelValues(worker).Set(float64(volumes))
}

func (emitter *PrometheusEmitter) httpResponseTimeMetrics(logger lager.Logger, event metric.Event) {
	route, exists := event.Attributes["route"]
	if !exists {
		logger.Error("failed-to-find-route-in-event", fmt.Errorf("expected method to exist in event.Attributes"))
	}

	method, exists := event.Attributes["method"]
	if !exists {
		logger.Error("failed-to-find-method-in-event", fmt.Errorf("expected method to exist in event.Attributes"))
	}

	responseTime, ok := event.Value.(float64)
	if !ok {
		logger.Error("http-response-time-event-value-type-mismatch", fmt.Errorf("expected event.Value to be a float64"))
	}

	emitter.httpRequestsDuration.WithLabelValues(method, route).Observe(responseTime / 1000)
}

func (emitter *PrometheusEmitter) schedulingMetrics(logger lager.Logger, event metric.Event) {
	pipeline, exists := event.Attributes["pipeline"]
	if !exists {
		logger.Error("failed-to-find-pipeline-in-event", fmt.Errorf("expected pipeline to exist in event.Attributes"))
	}

	duration, ok := event.Value.(float64)
	if !ok {
		logger.Error("scheduling-full-duration-value-type-mismatch", fmt.Errorf("expected event.Value to be a float64"))
	}

	switch event.Name {
	case "scheduling: full duration (ms)":
		// concourse_scheduling_full_duration_seconds
		emitter.schedulingFullDuration.WithLabelValues(pipeline).Set(duration / 1000)
	case "scheduling: loading versions duration (ms)":
		// concourse_scheduling_loading_duration_seconds
		emitter.schedulingLoadingDuration.WithLabelValues(pipeline).Set(duration / 1000)
	case "scheduling: job duration (ms)":
		// concourse_scheduling_job_duration_seconds
		emitter.schedulingJobDuration.WithLabelValues(pipeline).Set(duration / 1000)
	default:
	}
}

func (emitter *PrometheusEmitter) databaseMetrics(logger lager.Logger, event metric.Event) {
	value, ok := event.Value.(int)
	if !ok {
		logger.Error("db-value-type-mismatch", fmt.Errorf("expected event.Value to be a int"))
	}
	switch event.Name {
	case "database queries":
		emitter.dbQueriesTotal.Add(float64(value))
	case "database connections":
		connectionName, exists := event.Attributes["ConnectionName"]
		if !exists {
			logger.Error("failed-to-connection-name-in-event", fmt.Errorf("expected ConnectionName to exist in event.Attributes"))
		}
		emitter.dbConnections.WithLabelValues(connectionName).Set(float64(value))
	default:
	}

}
<commit_msg>prometheus: Add job label to job scheduling metric<commit_after>package emitter

import (
	"fmt"
	"net"
	"net/http"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/atc/db"
	"github.com/concourse/atc/metric"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

type PrometheusEmitter struct {
	buildsStarted     prometheus.Counter
	buildsFinished    prometheus.Counter
	buildsSucceeded   prometheus.Counter
	buildsErrored     prometheus.Counter
	buildsFailed      prometheus.Counter
	buildsAborted     prometheus.Counter
	buildsFinishedVec *prometheus.CounterVec
	buildDurationsVec *prometheus.HistogramVec

	workerContainers *prometheus.GaugeVec
	workerVolumes    *prometheus.GaugeVec

	httpRequestsDuration *prometheus.HistogramVec

	schedulingFullDuration    *prometheus.GaugeVec
	schedulingLoadingDuration *prometheus.GaugeVec
	schedulingJobDuration     *prometheus.GaugeVec

	dbQueriesTotal prometheus.Counter
	dbConnections  *prometheus.GaugeVec
}

type PrometheusConfig struct {
	BindIP   string `long:"prometheus-bind-ip" description:"IP to listen on to expose Prometheus metrics."`
	BindPort string `long:"prometheus-bind-port" description:"Port to listen on to expose Prometheus metrics."`
}

func init() {
	metric.RegisterEmitter(&PrometheusConfig{})
}

func (config *PrometheusConfig) Description() string { return "Prometheus" }
func (config *PrometheusConfig) IsConfigured() bool {
	return config.BindPort != "" && config.BindIP != ""
}
func (config *PrometheusConfig) bind() string {
	return fmt.Sprintf("%s:%s", config.BindIP, config.BindPort)
}

func (config *PrometheusConfig) NewEmitter() (metric.Emitter, error) {
	// build metrics
	buildsStarted := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "started_total",
		Help:      "Total number of Concourse builds started.",
	})
	prometheus.MustRegister(buildsStarted)

	buildsFinished := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "finished_total",
		Help:      "Total number of Concourse builds finished.",
	})
	prometheus.MustRegister(buildsFinished)

	buildsSucceeded := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "succeeded_total",
		Help:      "Total number of Concourse builds succeeded.",
	})
	prometheus.MustRegister(buildsSucceeded)

	buildsErrored := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "errored_total",
		Help:      "Total number of Concourse builds errored.",
	})
	prometheus.MustRegister(buildsErrored)

	buildsFailed := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "failed_total",
		Help:      "Total number of Concourse builds failed.",
	})
	prometheus.MustRegister(buildsFailed)

	buildsAborted := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "builds",
		Name:      "aborted_total",
		Help:      "Total number of Concourse builds aborted.",
	})
	prometheus.MustRegister(buildsAborted)

	buildsFinishedVec := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "concourse",
			Subsystem: "builds",
			Name:      "finished",
			Help:      "Count of builds finished across various dimensions.",
		},
		[]string{"team", "pipeline", "status"},
	)
	prometheus.MustRegister(buildsFinishedVec)
	buildDurationsVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Namespace: "concourse",
			Subsystem: "builds",
			Name:      "duration_seconds",
			Help:      "Build time in seconds",
			Buckets:   []float64{1, 60, 180, 300, 600, 900, 1200, 1800, 2700, 3600, 7200, 18000, 36000},
		},
		[]string{"team", "pipeline"},
	)
	prometheus.MustRegister(buildDurationsVec)

	// worker metrics
	workerContainers := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "workers",
			Name:      "containers",
			Help:      "Number of containers per worker",
		},
		[]string{"worker"},
	)
	prometheus.MustRegister(workerContainers)

	workerVolumes := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "workers",
			Name:      "volumes",
			Help:      "Number of volumes per worker",
		},
		[]string{"worker"},
	)
	prometheus.MustRegister(workerVolumes)

	// http metrics
	httpRequestsDuration := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Namespace: "concourse",
			Subsystem: "http_responses",
			Name:      "duration_seconds",
			Help:      "Response time in seconds",
		},
		[]string{"method", "route"},
	)
	prometheus.MustRegister(httpRequestsDuration)

	// scheduling metrics
	schedulingFullDuration := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "scheduling",
			Name:      "full_duration_seconds",
			Help:      "Last time taken to schedule an entire pipeline.",
		},
		[]string{"pipeline"},
	)
	prometheus.MustRegister(schedulingFullDuration)

	schedulingLoadingDuration := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "scheduling",
			Name:      "loading_duration_seconds",
			Help:      "Last time taken to load version information from the database for a pipeline.",
		},
		[]string{"pipeline"},
	)
	prometheus.MustRegister(schedulingLoadingDuration)

	schedulingJobDuration := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "scheduling",
			Name:      "job_duration_seconds",
			Help:      "Last time taken to calculate the set of valid input versions for a pipeline.",
		},
		[]string{"pipeline", "job"},
	)
	prometheus.MustRegister(schedulingJobDuration)

	dbQueriesTotal := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "concourse",
		Subsystem: "db",
		Name:      "queries_total",
		Help:      "Total number of database Concourse database queries",
	})
	prometheus.MustRegister(dbQueriesTotal)

	dbConnections := prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Namespace: "concourse",
			Subsystem: "db",
			Name:      "connections",
			Help:      "Current number of concourse database connections",
		},
		[]string{"dbname"},
	)
	prometheus.MustRegister(dbConnections)

	listener, err := net.Listen("tcp", config.bind())
	if err != nil {
		return nil, err
	}

	go http.Serve(listener, promhttp.Handler())

	return &PrometheusEmitter{
		buildsStarted:     buildsStarted,
		buildsFinished:    buildsFinished,
		buildsFinishedVec: buildsFinishedVec,
		buildDurationsVec: buildDurationsVec,
		buildsSucceeded:   buildsSucceeded,
		buildsErrored:     buildsErrored,
		buildsFailed:      buildsFailed,
		buildsAborted:     buildsAborted,

		workerContainers: workerContainers,
		workerVolumes:    workerVolumes,

		httpRequestsDuration: httpRequestsDuration,

		schedulingFullDuration:    schedulingFullDuration,
		schedulingLoadingDuration: schedulingLoadingDuration,
		schedulingJobDuration:     schedulingJobDuration,

		dbQueriesTotal: dbQueriesTotal,
		dbConnections:  dbConnections,
	}, nil
}

// Emit processes incoming metrics.
// In order to provide idiomatic Prometheus metrics, we'll have to convert the various
// Event types (differentiated by the less-than-ideal string Name field) into different
// Prometheus metrics.
func (emitter *PrometheusEmitter) Emit(logger lager.Logger, event metric.Event) {
	switch event.Name {
	case "build started":
		emitter.buildsStarted.Inc()
	case "build finished":
		emitter.buildFinishedMetrics(logger, event)
	case "worker containers":
		emitter.workerContainersMetric(logger, event)
	case "worker volumes":
		emitter.workerVolumesMetric(logger, event)
	case "http response time":
		emitter.httpResponseTimeMetrics(logger, event)
	case "scheduling: full duration (ms)":
		emitter.schedulingMetrics(logger, event)
	case "scheduling: loading versions duration (ms)":
		emitter.schedulingMetrics(logger, event)
	case "scheduling: job duration (ms)":
		emitter.schedulingMetrics(logger, event)
	case "database queries":
		emitter.databaseMetrics(logger, event)
	case "database connections":
		emitter.databaseMetrics(logger, event)
	default:
		// unless we have a specific metric, we do nothing
	}
}

func (emitter *PrometheusEmitter) buildFinishedMetrics(logger lager.Logger, event metric.Event) {
	// concourse_builds_finished_total
	emitter.buildsFinished.Inc()

	// concourse_builds_finished
	team, exists := event.Attributes["team_name"]
	if !exists {
		logger.Error("failed-to-find-team-name-in-event", fmt.Errorf("expected team_name to exist in event.Attributes"))
	}

	pipeline, exists := event.Attributes["pipeline"]
	if !exists {
		logger.Error("failed-to-find-pipeline-in-event", fmt.Errorf("expected pipeline to exist in event.Attributes"))
	}

	buildStatus, exists := event.Attributes["build_status"]
	if !exists {
		logger.Error("failed-to-find-build_status-in-event", fmt.Errorf("expected build_status to exist in event.Attributes"))
	}
	emitter.buildsFinishedVec.WithLabelValues(team, pipeline, buildStatus).Inc()

	// concourse_builds_(aborted|succeeded|failed|errored)_total
	switch buildStatus {
	case string(db.BuildStatusAborted):
		// concourse_builds_aborted_total
		emitter.buildsAborted.Inc()
	case string(db.BuildStatusSucceeded):
		// concourse_builds_succeeded_total
		emitter.buildsSucceeded.Inc()
	case string(db.BuildStatusFailed):
		// concourse_builds_failed_total
		emitter.buildsFailed.Inc()
	case string(db.BuildStatusErrored):
		// concourse_builds_errored_total
		emitter.buildsErrored.Inc()
	}

	// concourse_builds_duration_seconds
	duration, ok := event.Value.(float64)
	if !ok {
		logger.Error("build-finished-event-value-type-mismatch", fmt.Errorf("expected event.Value to be a float64"))
	}
	// seconds are the standard prometheus base unit for time
	duration = duration / 1000
	emitter.buildDurationsVec.WithLabelValues(team, pipeline).Observe(duration)
}

func (emitter *PrometheusEmitter) workerContainersMetric(logger lager.Logger, event metric.Event) {
	worker, exists := event.Attributes["worker"]
	if !exists {
		logger.Error("failed-to-find-worker-in-event", fmt.Errorf("expected worker to exist in event.Attributes"))
	}

	containers, ok := event.Value.(int)
	if !ok {
		logger.Error("worker-volumes-event-value-type-mismatch", fmt.Errorf("expected event.Value to be an int"))
	}

	emitter.workerContainers.WithLabelValues(worker).Set(float64(containers))
}

func (emitter *PrometheusEmitter) workerVolumesMetric(logger lager.Logger, event metric.Event) {
	worker, exists := event.Attributes["worker"]
	if !exists {
		logger.Error("failed-to-find-worker-in-event", fmt.Errorf("expected worker to exist in event.Attributes"))
	}

	volumes, ok := event.Value.(int)
	if !ok {
		logger.Error("worker-volumes-event-value-type-mismatch", fmt.Errorf("expected event.Value to be an int"))
	}

	emitter.workerVolumes.WithLabelValues(worker).Set(float64(volumes))
}

func (emitter *PrometheusEmitter) httpResponseTimeMetrics(logger lager.Logger, event metric.Event) {
	route, exists := event.Attributes["route"]
	if !exists {
		logger.Error("failed-to-find-route-in-event", fmt.Errorf("expected method to exist in event.Attributes"))
	}

	method, exists := event.Attributes["method"]
	if !exists {
		logger.Error("failed-to-find-method-in-event", fmt.Errorf("expected method to exist in event.Attributes"))
	}

	responseTime, ok := event.Value.(float64)
	if !ok {
		logger.Error("http-response-time-event-value-type-mismatch", fmt.Errorf("expected event.Value to be a float64"))
	}

	emitter.httpRequestsDuration.WithLabelValues(method, route).Observe(responseTime / 1000)
}

func (emitter *PrometheusEmitter) schedulingMetrics(logger lager.Logger, event metric.Event) {
	pipeline, exists := event.Attributes["pipeline"]
	if !exists {
		logger.Error("failed-to-find-pipeline-in-event", fmt.Errorf("expected pipeline to exist in event.Attributes"))
	}

	duration, ok := event.Value.(float64)
	if !ok {
		logger.Error("scheduling-full-duration-value-type-mismatch", fmt.Errorf("expected event.Value to be a float64"))
	}

	switch event.Name {
	case "scheduling: full duration (ms)":
		// concourse_scheduling_full_duration_seconds
		emitter.schedulingFullDuration.WithLabelValues(pipeline).Set(duration / 1000)
	case "scheduling: loading versions duration (ms)":
		// concourse_scheduling_loading_duration_seconds
		emitter.schedulingLoadingDuration.WithLabelValues(pipeline).Set(duration / 1000)
	case "scheduling: job duration (ms)":
		job, exists := event.Attributes["job"]
		if !exists {
			logger.Error("failed-to-find-job-in-event", fmt.Errorf("expected job to exist in event.Attributes"))
		}
		// concourse_scheduling_job_duration_seconds
		emitter.schedulingJobDuration.WithLabelValues(pipeline, job).Set(duration / 1000)
	default:
	}
}

func (emitter *PrometheusEmitter) databaseMetrics(logger lager.Logger, event metric.Event) {
	value, ok := event.Value.(int)
	if !ok {
		logger.Error("db-value-type-mismatch", fmt.Errorf("expected event.Value to be a int"))
	}
	switch event.Name {
	case "database queries":
		emitter.dbQueriesTotal.Add(float64(value))
	case "database connections":
		connectionName, exists := event.Attributes["ConnectionName"]
		if !exists {
			logger.Error("failed-to-connection-name-in-event", fmt.Errorf("expected ConnectionName to exist in event.Attributes"))
		}
		emitter.dbConnections.WithLabelValues(connectionName).Set(float64(value))
	default:
	}

}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/grafana/grafana/pkg/log"
	"github.com/raintank/raintank-metric/metric_tank/consolidation"
	"runtime"
	"sort"
)

// doRecover is the handler that turns panics into returns from the top level of getTarget.
func doRecover(errp *error) {
	e := recover()
	if e != nil {
		if _, ok := e.(runtime.Error); ok {
			panic(e)
		}
		if err, ok := e.(error); ok {
			*errp = err
		} else if errStr, ok := e.(string); ok {
			*errp = errors.New(errStr)
		} else {
			*errp = fmt.Errorf("%v", e)
		}
	}
	return
}

func divide(pointsA, pointsB []Point) []Point {
	if len(pointsA) != len(pointsB) {
		panic(fmt.Errorf("divide of a series with len %d by a series with len %d", len(pointsA), len(pointsB)))
	}
	out := make([]Point, len(pointsA))
	for i, a := range pointsA {
		b := pointsB[i]
		out[i] = Point{a.Val / b.Val, a.Ts}
	}
	return out
}

func consolidate(in []Point, num int, consolidator consolidation.Consolidator) []Point {
	aggFunc := consolidation.GetAggFunc(consolidator)
	buf := make([]float64, num)
	bufpos := -1
	points := make([]Point, (len(in)/num)+1)
	for inpos, p := range in {
		bufpos = inpos % num
		buf[bufpos] = p.Val
		if bufpos == num-1 {
			points = append(points, Point{aggFunc(buf), p.Ts})
		}
	}
	if bufpos != -1 && bufpos < num-1 {
		// we have an incomplete buf of some points that didn't get aggregated yet
		points = append(points, Point{aggFunc(buf[:bufpos+1]), in[len(in)-1].Ts})
	}
	return points
}

// returns how many points should be aggregated together so that you end up with as many points as possible,
// but never more than maxPoints
func aggEvery(numPoints, maxPoints uint32) int {
	return int((numPoints + maxPoints - 1) / maxPoints)
}

type planOption struct {
	archive  string
	interval uint32
	intestim bool
	points   uint32
	comment  string
}

type plan []planOption

func (a plan) Len() int           { return len(a) }
func (a plan) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a plan) Less(i, j int) bool { return a[i].points > a[j].points }

func getTarget(req Req, aggSettings []aggSetting, metaCache *MetaCache) (points []Point, err error) {
	defer doRecover(&err)
	archive := -1 // -1 means original data, 0 last agg level, 1 2nd last, etc.

	p := make([]planOption, len(aggSettings)+1)
	guess := false

	// note: the metacache is clearly not a perfect all-knowning entity, it just knows the last interval of metrics seen since program start
	// and we assume we can use that interval through history.
	// TODO: no support for interval changes, metrics not seen yet, missing datablocks, ...
	meta := metaCache.Get(req.key)
	interval := uint32(meta.interval)

	// we don't have the data yet, let's assume the interval is 10 seconds
	if interval == 0 {
		guess = true
		interval = 10
	}
	numPoints := (req.to - req.from) / interval

	p[0] = planOption{"raw", interval, guess, numPoints, ""}

	aggs := aggSettingsSpanDesc(aggSettings)
	sort.Sort(aggs)
	finished := false
	for i, aggSetting := range aggs {
		numPointsHere := (req.to - req.from) / aggSetting.span
		p[i+1] = planOption{fmt.Sprintf("agg %d", i), aggSetting.span, false, numPointsHere, ""}
		if numPointsHere >= req.minPoints && !finished {
			archive = i
			interval = aggSetting.span
			numPoints = numPointsHere
			finished = true
		}
	}

	p[archive+1].comment = "<-- chosen"

	// note, it should always be safe to dynamically switch on/off consolidation based on how well our data stacks up against the request
	// i.e. whether your data got consolidated or not, it should be pretty equivalent.
	// for that reason, stdev should not be done as a consolidation. but sos is still useful for when we explicitly (and always, not optionally) want the stdev.

	readConsolidated := (archive != -1)                 // do we need to read from a downsampled series?
	runtimeConsolidation := (numPoints > req.maxPoints) // do we need to compress any points at runtime?

	log.Debug("getTarget()         %s", req)
	log.Debug("type   interval   points")
	sortedPlan := plan(p)
	sort.Sort(sortedPlan)
	for _, opt := range p {
		iStr := fmt.Sprintf("%d", opt.interval)
		if opt.intestim {
			iStr = fmt.Sprintf("%d (guess)", opt.interval)
		}
		log.Debug("%-6s %-10s %-6d %s", opt.archive, iStr, opt.points, opt.comment)
	}
	log.Debug("runtimeConsolidation: %t", runtimeConsolidation)

	if !readConsolidated && !runtimeConsolidation {
		return getSeries(req.key, consolidation.None, 0, req.from, req.to), nil
	} else if !readConsolidated && runtimeConsolidation {
		return consolidate(
			getSeries(req.key, consolidation.None, 0, req.from, req.to),
			aggEvery(numPoints, req.maxPoints),
			req.consolidator), nil
	} else if readConsolidated && !runtimeConsolidation {
		if req.consolidator == consolidation.Avg {
			return divide(
				getSeries(req.key, consolidation.Sum, interval, req.from, req.to),
				getSeries(req.key, consolidation.Cnt, interval, req.from, req.to),
			), nil
		} else {
			return getSeries(req.key, req.consolidator, interval, req.from, req.to), nil
		}
	} else {
		// readConsolidated && runtimeConsolidation
		aggNum := aggEvery(numPoints, req.maxPoints)
		if req.consolidator == consolidation.Avg {
			return divide(
				consolidate(
					getSeries(req.key, consolidation.Sum, interval, req.from, req.to),
					aggNum,
					consolidation.Sum),
				consolidate(
					getSeries(req.key, consolidation.Cnt, interval, req.from, req.to),
					aggNum,
					consolidation.Cnt),
			), nil
		} else {
			return consolidate(
				getSeries(req.key, req.consolidator, interval, req.from, req.to),
				aggNum, req.consolidator), nil
		}
	}
}

func logLoad(typ, key string, from, to uint32) {
	log.Debug("load from %-6s %-20s %d - %d (%s - %s) span:%ds", typ, key, from, to, TS(from), TS(to), to-from-1)
}

func aggMetricKey(key, archive string, aggSpan uint32) string {
	return fmt.Sprintf("%s_%s_%d", key, archive, aggSpan)
}

// getSeries just gets the needed raw iters from mem and/or cassandra, based on from/to
// it can query for data within aggregated archives, by using fn min/max/sos/sum/cnt and providing the matching agg span.
func getSeries(key string, consolidator consolidation.Consolidator, aggSpan, fromUnix, toUnix uint32) []Point {
	iters := make([]Iter, 0)
	memIters := make([]Iter, 0)
	oldest := toUnix
	if metric, ok := metrics.Get(key); ok {
		if consolidator != consolidation.None {
			logLoad("memory", aggMetricKey(key, consolidator.Archive(), aggSpan), fromUnix, toUnix)
			oldest, memIters = metric.GetAggregated(consolidator, aggSpan, fromUnix, toUnix)
		} else {
			logLoad("memory", key, fromUnix, toUnix)
			oldest, memIters = metric.Get(fromUnix, toUnix)
		}
	}
	if oldest > fromUnix {
		reqSpanBoth.Value(int64(toUnix - fromUnix))
		if consolidator != consolidation.None {
			key = aggMetricKey(key, consolidator.Archive(), aggSpan)
		}
		logLoad("cassan", key, fromUnix, oldest)
		storeIters, err := searchCassandra(key, fromUnix, oldest)
		if err != nil {
			panic(err)
		}
		iters = append(iters, storeIters...)
	} else {
		reqSpanMem.Value(int64(toUnix - fromUnix))
	}
	iters = append(iters, memIters...)

	points := make([]Point, 0)
	for _, iter := range iters {
		total := 0
		good := 0
		for iter.Next() {
			total += 1
			ts, val := iter.Values()
			if ts >= fromUnix && ts < toUnix {
				good += 1
				points = append(points, Point{val, ts})
			}
		}
		log.Debug("getSeries: iter %s  values good/total %d/%d", iter.cmt, good, total)
	}
	return points
}
<commit_msg>fix consolidation. fix #69<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/grafana/grafana/pkg/log"
	"github.com/raintank/raintank-metric/metric_tank/consolidation"
	"runtime"
	"sort"
)

// doRecover is the handler that turns panics into returns from the top level of getTarget.
func doRecover(errp *error) {
	e := recover()
	if e != nil {
		if _, ok := e.(runtime.Error); ok {
			panic(e)
		}
		if err, ok := e.(error); ok {
			*errp = err
		} else if errStr, ok := e.(string); ok {
			*errp = errors.New(errStr)
		} else {
			*errp = fmt.Errorf("%v", e)
		}
	}
	return
}

func divide(pointsA, pointsB []Point) []Point {
	if len(pointsA) != len(pointsB) {
		panic(fmt.Errorf("divide of a series with len %d by a series with len %d", len(pointsA), len(pointsB)))
	}
	out := make([]Point, len(pointsA))
	for i, a := range pointsA {
		b := pointsB[i]
		out[i] = Point{a.Val / b.Val, a.Ts}
	}
	return out
}

func consolidate(in []Point, num int, consolidator consolidation.Consolidator) []Point {
	aggFunc := consolidation.GetAggFunc(consolidator)
	buf := make([]float64, num)
	bufpos := -1
	outLen := len(in) / num
	if len(in)%num != 0 {
		outLen += 1
	}
	points := make([]Point, 0, outLen)
	for inpos, p := range in {
		bufpos = inpos % num
		buf[bufpos] = p.Val
		if bufpos == num-1 {
			points = append(points, Point{aggFunc(buf), p.Ts})
		}
	}
	if bufpos != -1 && bufpos < num-1 {
		// we have an incomplete buf of some points that didn't get aggregated yet
		points = append(points, Point{aggFunc(buf[:bufpos+1]), in[len(in)-1].Ts})
	}
	return points
}

// returns how many points should be aggregated together so that you end up with as many points as possible,
// but never more than maxPoints
func aggEvery(numPoints, maxPoints uint32) int {
	return int((numPoints + maxPoints - 1) / maxPoints)
}

type planOption struct {
	archive  string
	interval uint32
	intestim bool
	points   uint32
	comment  string
}

type plan []planOption

func (a plan) Len() int           { return len(a) }
func (a plan) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a plan) Less(i, j int) bool { return a[i].points > a[j].points }

func getTarget(req Req, aggSettings []aggSetting, metaCache *MetaCache) (points []Point, err error) {
	defer doRecover(&err)
	archive := -1 // -1 means original data, 0 last agg level, 1 2nd last, etc.

	p := make([]planOption, len(aggSettings)+1)
	guess := false

	// note: the metacache is clearly not a perfect all-knowning entity, it just knows the last interval of metrics seen since program start
	// and we assume we can use that interval through history.
	// TODO: no support for interval changes, metrics not seen yet, missing datablocks, ...
	meta := metaCache.Get(req.key)
	interval := uint32(meta.interval)

	// we don't have the data yet, let's assume the interval is 10 seconds
	if interval == 0 {
		guess = true
		interval = 10
	}
	numPoints := (req.to - req.from) / interval

	p[0] = planOption{"raw", interval, guess, numPoints, ""}

	aggs := aggSettingsSpanDesc(aggSettings)
	sort.Sort(aggs)
	finished := false
	for i, aggSetting := range aggs {
		numPointsHere := (req.to - req.from) / aggSetting.span
		p[i+1] = planOption{fmt.Sprintf("agg %d", i), aggSetting.span, false, numPointsHere, ""}
		if numPointsHere >= req.minPoints && !finished {
			archive = i
			interval = aggSetting.span
			numPoints = numPointsHere
			finished = true
		}
	}

	p[archive+1].comment = "<-- chosen"

	// note, it should always be safe to dynamically switch on/off consolidation based on how well our data stacks up against the request
	// i.e. whether your data got consolidated or not, it should be pretty equivalent.
	// for that reason, stdev should not be done as a consolidation. but sos is still useful for when we explicitly (and always, not optionally) want the stdev.

	readConsolidated := (archive != -1)                 // do we need to read from a downsampled series?
	runtimeConsolidation := (numPoints > req.maxPoints) // do we need to compress any points at runtime?

	log.Debug("getTarget()         %s", req)
	log.Debug("type   interval   points")
	sortedPlan := plan(p)
	sort.Sort(sortedPlan)
	for _, opt := range p {
		iStr := fmt.Sprintf("%d", opt.interval)
		if opt.intestim {
			iStr = fmt.Sprintf("%d (guess)", opt.interval)
		}
		log.Debug("%-6s %-10s %-6d %s", opt.archive, iStr, opt.points, opt.comment)
	}
	log.Debug("runtimeConsolidation: %t", runtimeConsolidation)

	if !readConsolidated && !runtimeConsolidation {
		return getSeries(req.key, consolidation.None, 0, req.from, req.to), nil
	} else if !readConsolidated && runtimeConsolidation {
		return consolidate(
			getSeries(req.key, consolidation.None, 0, req.from, req.to),
			aggEvery(numPoints, req.maxPoints),
			req.consolidator), nil
	} else if readConsolidated && !runtimeConsolidation {
		if req.consolidator == consolidation.Avg {
			return divide(
				getSeries(req.key, consolidation.Sum, interval, req.from, req.to),
				getSeries(req.key, consolidation.Cnt, interval, req.from, req.to),
			), nil
		} else {
			return getSeries(req.key, req.consolidator, interval, req.from, req.to), nil
		}
	} else {
		// readConsolidated && runtimeConsolidation
		aggNum := aggEvery(numPoints, req.maxPoints)
		if req.consolidator == consolidation.Avg {
			return divide(
				consolidate(
					getSeries(req.key, consolidation.Sum, interval, req.from, req.to),
					aggNum,
					consolidation.Sum),
				consolidate(
					getSeries(req.key, consolidation.Cnt, interval, req.from, req.to),
					aggNum,
					consolidation.Cnt),
			), nil
		} else {
			return consolidate(
				getSeries(req.key, req.consolidator, interval, req.from, req.to),
				aggNum, req.consolidator), nil
		}
	}
}

func logLoad(typ, key string, from, to uint32) {
	log.Debug("load from %-6s %-20s %d - %d (%s - %s) span:%ds", typ, key, from, to, TS(from), TS(to), to-from-1)
}

func aggMetricKey(key, archive string, aggSpan uint32) string {
	return fmt.Sprintf("%s_%s_%d", key, archive, aggSpan)
}

// getSeries just gets the needed raw iters from mem and/or cassandra, based on from/to
// it can query for data within aggregated archives, by using fn min/max/sos/sum/cnt and providing the matching agg span.
func getSeries(key string, consolidator consolidation.Consolidator, aggSpan, fromUnix, toUnix uint32) []Point {
	iters := make([]Iter, 0)
	memIters := make([]Iter, 0)
	oldest := toUnix
	if metric, ok := metrics.Get(key); ok {
		if consolidator != consolidation.None {
			logLoad("memory", aggMetricKey(key, consolidator.Archive(), aggSpan), fromUnix, toUnix)
			oldest, memIters = metric.GetAggregated(consolidator, aggSpan, fromUnix, toUnix)
		} else {
			logLoad("memory", key, fromUnix, toUnix)
			oldest, memIters = metric.Get(fromUnix, toUnix)
		}
	}
	if oldest > fromUnix {
		reqSpanBoth.Value(int64(toUnix - fromUnix))
		if consolidator != consolidation.None {
			key = aggMetricKey(key, consolidator.Archive(), aggSpan)
		}
		logLoad("cassan", key, fromUnix, oldest)
		storeIters, err := searchCassandra(key, fromUnix, oldest)
		if err != nil {
			panic(err)
		}
		iters = append(iters, storeIters...)
	} else {
		reqSpanMem.Value(int64(toUnix - fromUnix))
	}
	iters = append(iters, memIters...)

	points := make([]Point, 0)
	for _, iter := range iters {
		total := 0
		good := 0
		for iter.Next() {
			total += 1
			ts, val := iter.Values()
			if ts >= fromUnix && ts < toUnix {
				good += 1
				points = append(points, Point{val, ts})
			}
		}
		log.Debug("getSeries: iter %s  values good/total %d/%d", iter.cmt, good, total)
	}
	return points
}
<|endoftext|>


<commit_before>package aggregators

import (
	"testing"
)

func TestGoalsScoreSpecName(t *testing.T) {
	name := "a"
	as := MustNew(name, "goalsscore")
	got := as.Name()
	if got != name {
		t.Errorf("Name - got: %s, want: %s", got, name)
	}
}

func TestGoalsScoreSpecKind(t *testing.T) {
	kind := "goalsscore"
	as := MustNew("a", kind)
	got := as.Kind()
	if got != kind {
		t.Errorf("Kind - got: %s, want: %s", got, kind)
	}
}

func TestGoalsScoreSpecArg(t *testing.T) {
	arg := ""
	as := MustNew("a", "goalsscore")
	got := as.Arg()
	if got != arg {
		t.Errorf("Arg - got: %s, want: %s", got, arg)
	}
}
<commit_msg>Expand goalsscore aggregator test<commit_after>package aggregators

import (
	"github.com/lawrencewoodman/dexpr"
	"github.com/lawrencewoodman/dlit"
	"github.com/vlifesystems/rhkit/goal"
	"testing"
)

func TestGoalsScoreSpecName(t *testing.T) {
	name := "a"
	as := MustNew(name, "goalsscore")
	got := as.Name()
	if got != name {
		t.Errorf("Name - got: %s, want: %s", got, name)
	}
}

func TestGoalsScoreSpecKind(t *testing.T) {
	kind := "goalsscore"
	as := MustNew("a", kind)
	got := as.Kind()
	if got != kind {
		t.Errorf("Kind - got: %s, want: %s", got, kind)
	}
}

func TestGoalsScoreSpecArg(t *testing.T) {
	arg := ""
	as := MustNew("a", "goalsscore")
	got := as.Arg()
	if got != arg {
		t.Errorf("Arg - got: %s, want: %s", got, arg)
	}
}

func TestGoalsScoreNextRecord(t *testing.T) {
	as := MustNew("a", "goalsscore")
	ai := as.New()
	record := map[string]*dlit.Literal{}
	got := ai.NextRecord(record, true)
	if got != nil {
		t.Errorf("NextRecord: got: %s, want: %s", got, nil)
	}
}

func TestGoalsScoreResult(t *testing.T) {
	aggregatorSpecs := []AggregatorSpec{
		MustNew("income", "calc", "3 + 4"),
		MustNew("costs", "calc", "5 + 6"),
		MustNew("profit", "calc", "costs - income"),
		MustNew("highest", "calc", "7+10"),
		MustNew("lowest", "calc", "5"),
		MustNew("mid", "calc", "3"),
		MustNew("goalsScore", "goalsscore"),
	}
	cases := []struct {
		goals []*goal.Goal
		want  *dlit.Literal
	}{
		{goals: []*goal.Goal{
			goal.MustNew("income > 10"),
			goal.MustNew("costs < 2"),
			goal.MustNew("profit >= 6"),
			goal.MustNew("highest >= 18"),
			goal.MustNew("lowest <= 4"),
		},
			want: dlit.MustNew(0),
		},
		{goals: []*goal.Goal{
			goal.MustNew("income > 10"),
			goal.MustNew("costs < 2"),
			goal.MustNew("profit >= 2"),
			goal.MustNew("highest >= 15"),
			goal.MustNew("lowest <= 4"),
		},
			want: dlit.MustNew(0.002),
		},
		{goals: []*goal.Goal{
			goal.MustNew("income > 6"),
			goal.MustNew("costs < 2"),
			goal.MustNew("profit >= 2"),
		},
			want: dlit.MustNew(1.001),
		},
		{goals: []*goal.Goal{
			goal.MustNew("income > 6"),
			goal.MustNew("costs < 20"),
			goal.MustNew("profit >= 10"),
			goal.MustNew("highest >= 15"),
			goal.MustNew("lowest <= 7"),
		},
			want: dlit.MustNew(2.002),
		},
		{goals: []*goal.Goal{
			goal.MustNew("income > 6"),
			goal.MustNew("costs < 20"),
			goal.MustNew("profit >= nothing"),
			goal.MustNew("highest >= 15"),
			goal.MustNew("lowest <= 7"),
		},
			want: dlit.MustNew(dexpr.InvalidExprError{
				Expr: "profit >= nothing",
				Err:  dexpr.VarNotExistError("nothing"),
			}),
		},
	}
	numRecords := int64(12)
	instances := make([]AggregatorInstance, len(aggregatorSpecs))
	for i, aggregatorSpec := range aggregatorSpecs {
		instances[i] = aggregatorSpec.New()
	}
	goalsScoreInstance := instances[len(instances)-1]
	for i, c := range cases {
		got := goalsScoreInstance.Result(instances, c.goals, numRecords)
		if got.String() != c.want.String() {
			t.Errorf("(%d) Result: got: %s, want: %s", i, got, c.want)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"fmt"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	. "github.com/claudetech/loggo/default"
	"golang.org/x/net/context"
)

// Mount the fuse volume
func Mount(client *Drive, mountpoint string, mountOptions []string) error {
	Log.Infof("Mounting path %v", mountpoint)

	if _, err := os.Stat(mountpoint); os.IsNotExist(err) {
		Log.Debugf("Mountpoint doesn't exist, creating...")
		if err := os.MkdirAll(mountpoint, 0644); nil != err {
			Log.Debugf("%v", err)
			return fmt.Errorf("Could not create mount directory %v", mountpoint)
		}
	}

	fuse.Debug = func(msg interface{}) {
		Log.Tracef("FUSE %v", msg)
	}

	options := []fuse.MountOption{
		fuse.NoAppleDouble(),
		fuse.NoAppleXattr(),
		fuse.ReadOnly(),
	}
	for _, option := range mountOptions {
		var o fuse.MountOption
		switch option {
		case "allow_other":
			o = fuse.AllowOther()
			break
		case "allow_root":
			o = fuse.AllowRoot()
			break
		case "allow_dev":
			o = fuse.AllowDev()
			break
		case "allow_non_empty_mount":
			o = fuse.AllowNonEmptyMount()
			break
		case "allow_suid":
			o = fuse.AllowSUID()
			break
		}
		options = append(options, o)
	}

	c, err := fuse.Mount(mountpoint, options...)
	if err != nil {
		return err
	}
	defer c.Close()

	object, err := client.GetObject("root")
	if nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Could not get root node")
	}

	filesys := &FS{
		client: client,
		rootID: object.ObjectID,
	}
	if err := fs.Serve(c, filesys); err != nil {
		return err
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; err != nil {
		return err
	}

	Log.Infof("Unmounting path %v", mountpoint)
	err = fuse.Unmount(mountpoint)
	if nil != err {
		return err
	}

	return nil
}

// FS the fuse filesystem
type FS struct {
	client *Drive
	rootID string
}

// Root returns the root path
func (f *FS) Root() (fs.Node, error) {
	object, err := f.client.GetObject(f.rootID)
	if nil != err {
		Log.Warningf("%v", err)
		return nil, fmt.Errorf("Could not get root object")
	}
	return &Object{
		client: f.client,
		object: &object,
	}, nil
}

// Object represents one drive object
type Object struct {
	client *Drive
	object *APIObject
	buffer *Buffer
}

// Attr returns the attributes for a directory
func (o *Object) Attr(ctx context.Context, attr *fuse.Attr) error {
	if o.object.IsDir {
		attr.Mode = os.ModeDir | 0755
		attr.Size = 0
	} else {
		attr.Mode = 0644
		attr.Size = o.object.Size
	}

	attr.Mtime = o.object.LastModified
	attr.Crtime = o.object.LastModified
	attr.Ctime = o.object.LastModified

	return nil
}

// ReadDirAll shows all files in the current directory
func (o *Object) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	objects, err := o.client.GetObjectsByParent(o.object.ObjectID)
	if nil != err {
		Log.Warningf("%v", err)
		return nil, err
	}

	dirs := []fuse.Dirent{}
	for _, object := range objects {
		dirs = append(dirs, fuse.Dirent{
			Name: object.Name,
			Type: fuse.DT_File,
		})
	}
	return dirs, nil
}

// Lookup tests if a file is existent in the current directory
func (o *Object) Lookup(ctx context.Context, name string) (fs.Node, error) {
	object, err := o.client.GetObjectByParentAndName(o.object.ObjectID, name)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, err
	}

	return &Object{
		client: o.client,
		object: &object,
	}, nil
}

// Open opens a file for reading
func (o *Object) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {
	if req.Dir {
		return o, nil
	}

	buffer, err := o.client.Open(o.object)
	if nil != err {
		Log.Warningf("%v", err)
		return o, fuse.ENOENT
	}
	o.buffer = buffer

	return o, nil
}

// Release a stream
func (o *Object) Release(ctx context.Context, req *fuse.ReleaseRequest) error {
	if nil != o.buffer {
		if err := o.buffer.Close(); nil != err {
			Log.Debugf("%v", err)
			Log.Warningf("Could not close buffer stream")
		}
	}
	return nil
}

// Read reads some bytes or the whole file
func (o *Object) Read(ctx context.Context, req *fuse.ReadRequest, resp *fuse.ReadResponse) error {
	buf, err := o.buffer.ReadBytes(req.Offset, int64(req.Size), false)
	if nil != err {
		Log.Warningf("%v", err)
		return err
	}

	resp.Data = buf[:]
	return nil
}
<commit_msg>do not panic if fuse option is not supported<commit_after>package main

import (
	"os"

	"fmt"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	. "github.com/claudetech/loggo/default"
	"golang.org/x/net/context"
)

// Mount the fuse volume
func Mount(client *Drive, mountpoint string, mountOptions []string) error {
	Log.Infof("Mounting path %v", mountpoint)

	if _, err := os.Stat(mountpoint); os.IsNotExist(err) {
		Log.Debugf("Mountpoint doesn't exist, creating...")
		if err := os.MkdirAll(mountpoint, 0644); nil != err {
			Log.Debugf("%v", err)
			return fmt.Errorf("Could not create mount directory %v", mountpoint)
		}
	}

	fuse.Debug = func(msg interface{}) {
		Log.Tracef("FUSE %v", msg)
	}

	options := []fuse.MountOption{
		fuse.NoAppleDouble(),
		fuse.NoAppleXattr(),
		fuse.ReadOnly(),
	}
	for _, option := range mountOptions {
		switch option {
		case "allow_other":
			options = append(options, fuse.AllowOther())
			break
		case "allow_root":
			options = append(options, fuse.AllowRoot())
			break
		case "allow_dev":
			options = append(options, fuse.AllowDev())
			break
		case "allow_non_empty_mount":
			options = append(options, fuse.AllowNonEmptyMount())
			break
		case "allow_suid":
			options = append(options, fuse.AllowSUID())
			break
		default:
			Log.Warningf("Fuse option %v is not supported, yet", option)
			break
		}
	}

	c, err := fuse.Mount(mountpoint, options...)
	if err != nil {
		return err
	}
	defer c.Close()

	object, err := client.GetObject("root")
	if nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Could not get root node")
	}

	filesys := &FS{
		client: client,
		rootID: object.ObjectID,
	}
	if err := fs.Serve(c, filesys); err != nil {
		return err
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; err != nil {
		return err
	}

	Log.Infof("Unmounting path %v", mountpoint)
	err = fuse.Unmount(mountpoint)
	if nil != err {
		return err
	}

	return nil
}

// FS the fuse filesystem
type FS struct {
	client *Drive
	rootID string
}

// Root returns the root path
func (f *FS) Root() (fs.Node, error) {
	object, err := f.client.GetObject(f.rootID)
	if nil != err {
		Log.Warningf("%v", err)
		return nil, fmt.Errorf("Could not get root object")
	}
	return &Object{
		client: f.client,
		object: &object,
	}, nil
}

// Object represents one drive object
type Object struct {
	client *Drive
	object *APIObject
	buffer *Buffer
}

// Attr returns the attributes for a directory
func (o *Object) Attr(ctx context.Context, attr *fuse.Attr) error {
	if o.object.IsDir {
		attr.Mode = os.ModeDir | 0755
		attr.Size = 0
	} else {
		attr.Mode = 0644
		attr.Size = o.object.Size
	}

	attr.Mtime = o.object.LastModified
	attr.Crtime = o.object.LastModified
	attr.Ctime = o.object.LastModified

	return nil
}

// ReadDirAll shows all files in the current directory
func (o *Object) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	objects, err := o.client.GetObjectsByParent(o.object.ObjectID)
	if nil != err {
		Log.Warningf("%v", err)
		return nil, err
	}

	dirs := []fuse.Dirent{}
	for _, object := range objects {
		dirs = append(dirs, fuse.Dirent{
			Name: object.Name,
			Type: fuse.DT_File,
		})
	}
	return dirs, nil
}

// Lookup tests if a file is existent in the current directory
func (o *Object) Lookup(ctx context.Context, name string) (fs.Node, error) {
	object, err := o.client.GetObjectByParentAndName(o.object.ObjectID, name)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, err
	}

	return &Object{
		client: o.client,
		object: &object,
	}, nil
}

// Open opens a file for reading
func (o *Object) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {
	if req.Dir {
		return o, nil
	}

	buffer, err := o.client.Open(o.object)
	if nil != err {
		Log.Warningf("%v", err)
		return o, fuse.ENOENT
	}
	o.buffer = buffer

	return o, nil
}

// Release a stream
func (o *Object) Release(ctx context.Context, req *fuse.ReleaseRequest) error {
	if nil != o.buffer {
		if err := o.buffer.Close(); nil != err {
			Log.Debugf("%v", err)
			Log.Warningf("Could not close buffer stream")
		}
	}
	return nil
}

// Read reads some bytes or the whole file
func (o *Object) Read(ctx context.Context, req *fuse.ReadRequest, resp *fuse.ReadResponse) error {
	buf, err := o.buffer.ReadBytes(req.Offset, int64(req.Size), false)
	if nil != err {
		Log.Warningf("%v", err)
		return err
	}

	resp.Data = buf[:]
	return nil
}
<|endoftext|>


<commit_before>package sudoku

// Sudoku has one rule - no value can be repeated horizontally, vertically,
// or in the same section. This gives rise to a few simple rules - and those
// rules are applied across each cluster - without even knowning the orientation
// of the cluster.
//
// In order to make solving possible however, that one rule is enforced with
// the following rules:
//
// 1) If all cells are solved, that cluster is solved.
// 2) If any cell is solved, it has no possibles.
// 3) If any cell is solved, that value is not possible in other cells.
// 4) If any cell only has one possible value, that is that cell's value.
// 5) If any x cells only have x possible values, those values are not possible
//  outside of those cells - those values are constrained to those cells.
// 6) If any value only has one possible cell, that is that cell's value.
// 7) If any x values only have x possible cells, those cells only have those
//  possible values - those cells are constrained to those values.
//
// Additional Helper functions are included and explained later.

// indexedCLuster is a datatype for an index for the values of the cluster.
// Each possible value is a key in the map. THe value of each key is an array of
// possible locations for that value - the index and order are not defined,
// instead the values of the array are the indexes of possible cells in for that
// value.
type intArray []int
type indexedCluster map[int]intArray

func indexCluster(in []cell) (out indexedCluster) {
	for id, each := range in {
		for _, onePossible_ := range each.possible {
			out[onePossible] = append(out[onePossible], id)
		}
	}
	return out
}

// This covers rule 1 from above:
// 1) If all cells are solved, that cluster is solved.
func clusterSolved(cluster []cell, u chan cell) (solved bool) {
	solved = true
	for _, each := range cluster {
		if each.actual == 0 {
			solved = false
		} else if len(each.possible) > 0 {
			solved = false
		}
	}
	return solved
}

// This covers rule 2 from above:
// 2) If any cell is solved, it has no possibles.
func solvedNoPossible(cluster []cell, u chan cell) (changed bool) {
	for _, each := range cluster {
		if each.actual == 0 {
			continue
		}
		if len(each.possible) > 0 {
			changed = true
			u <- cell{
				location: each.location,
				possible: each.possible,
			}
		}
	}
	return changed
}

// Removes known values from other possibles in the same cluster
// Covers rule 3 from above:
// 3) If any cell is solved, that value is not possible in other cells.
func eliminateKnowns(workingCluster []cell, u chan cell) (changed bool) {
	var knownValues []int

	// Loop thru and find all solved values.
	for _, each := range workingCluster {
		if each.actual != 0 {
			knownValues = append(knownValues, each.actual)
		}
	}

	for _, each := range workingCluster {
		if anyInArr(each, knownValues) {
			u <- cell{
				location: each.location,
				possible: subArr(each, knownValues)
			}
			changed = true
		}
	}
	return changed
}

// This covers the 4th rule from above:
// 4) If any cell only has one possible value, that is that cell's value.
func singleValueSolver(cluster []cell, u chan cell) (changed bool) {
	for _, each := range cluster {
		// if the cell is already solved, skip this.
		if each.actual != 0 {
			continue
		}

		// if there is more than one possible value for this cell, you should be good
		if len(each.possible) > 1 {
			continue
		}

		// should never happen - probably #TODO# to catch this
		if len(each.possible) < 1 {
			panic("Found an unsolved cell with no possible values")
		}

		changed = true

		// send back an update for this cell
		u <- cell{
			location: each.location,
			actual:   each.possible[0],
			possible: []int{}
		}
	}
	return changed
}

// A helper function to determine the values certain cells hit
func valuesPainted(markedCells []int, cluster []cell) (neededValues []int) {
	for _, oneCell := range markedCells {
		neededValues = addArr(neededValues, cluster[oneCell].possible)
	}
	return neededValues
}

// A helper function to determine the number of valus hit given a specific set
// of cells.
func cellsCost(markedCells []int, cluster []cell) int {
	return len(valuesPainted(markedCells, cluster))
}

func cellLimiterChild(limit int, markedCells []int, cluster []cell, u chan cell) (changed bool) {
	valueCount := cellsCost(markedCells, cluster)
	// you have overspent - it's a no-go
	if len(markedCells) > limit {
		return false
	}

	// did you fill the cells? if so, mark it
	if valueCount == len(markedCells) {
		// it's a match - so remove values
		for idCell, oneCell := range cluster {
			if oneCell.actual != 0 {
				// already solved
				continue
			}
			if markedCells[idCell] {
				// this cell is a part of the list - no exclusions needed
				continue
			}
			remove := map[int]bool{}
			for potential, _ := range oneCell.possible {
				if markedCells[idCell] {
					// this possibility exists in the map - need to remove it
					remove[potential] = true
				}
			}
			if len(remove) > 0 {
				changed = true
				u <- cell{
					location: oneCell.location,
					possible: remove,
				}
			}
		}
	}

	// you have room to add more cells (depth first?)
	if len(markedCells) < limit {
		if valueCount < len(markedCells) {
			// #TODO# probably fix this? rework into error?
			panic("less possible values than squares to fill")
		}
		if valueCount > len(markedCells) {
			// you need to try adding each other cell
			for idCell, oneCell := range cluster {
				if oneCell.actual != 0 {
					continue
				}
				if markedCells[idCell] {
					// this cell is already in the map, skip
					continue
				}

				// decend down into looking at that cell
				childMarkedCells := markedCells // #TODO# map copy
				childMarkedCells[idCell] = true
				if cellLimiterChild(limit, childMarkedCells, cluster, u) {
					// if you got true from the child, pass it on
					return true
				}
			}
		}
	}

	return changed
}

// This covers rule 5:
// 5) If any x cells have x possible values, those values are not possible
//  outside of those cells - those values are constrained to those cells.
func cellLimiter(cluster []cell, u chan cell) (changed bool) {
	upperBound := len(cluster)
	for _, eachCell := range cluster {
		if eachCell.actual != 0 {
			upperBound--
		}
	}
	for i := 2; i <= upperBound; i++ {
		if cellLimiterChild(i, map[int]bool{}, cluster, u) {
			changed = true
		}
	}
	return changed
}

// This covers rule 6 from above:
// 6) If any value only has one possible cell, that is that cell's value.
func singleCellSolver(index indexedCluster, cluster []cell, u chan cell) (changed bool) {
	for val, section := range index {
		if len(section) < 1 {
			// something went terribly wrong here - #TODO# add panic catch?
			panic("Found an unsolved cell with no possible values")
		} else if len(section) == 1 {
			u <- cell{
				location: cluster[section[0]].location,
				actual:   val,
				possible: []int{},
			}
			changed = true
		}
	}
	return changed
}

// A helper function to determine the number of cells hit by working across a map
// of values.
func valuesCost(markedVals map[int]bool, index indexedCluster, cluseter []cell) int {
	var neededCells map[int]bool
	for value, _ := range markedVals {
		for oneCell, _ := range index[value] {
			neededCells[oneCell] = true
		}
	}
	return len(neededCells)
}

func valueLimiterChild(limit int, markedValues map[int]bool, index indexedCluster,
	cluster []cell, u chan cell) (changed bool) {
	if len(markedValues) > limit {
		// we have marked more values
		return false
	}
	currentCost := valuesCost(markedValues, index, cluster)
	if currentCost > limit {
		// you're over the budget to spend
		return false
	}
	return changed
}

// THis covers rule 7 from above:
// 7) If any x values have only x possible cells, those cells only have those
//  possible values - those cells are constrained to those values.
/*
func valueLimiter(index indexedCluster, cluster []cell, u chan cell) (changed bool) {

}
*/
<commit_msg>rewrote `cellLimiter` to work with arrays, and possibly cleaned it up<commit_after>package sudoku

// Sudoku has one rule - no value can be repeated horizontally, vertically,
// or in the same section. This gives rise to a few simple rules - and those
// rules are applied across each cluster - without even knowning the orientation
// of the cluster.
//
// In order to make solving possible however, that one rule is enforced with
// the following rules:
//
// 1) If all cells are solved, that cluster is solved.
// 2) If any cell is solved, it has no possibles.
// 3) If any cell is solved, that value is not possible in other cells.
// 4) If any cell only has one possible value, that is that cell's value.
// 5) If any x cells only have x possible values, those values are not possible
//  outside of those cells - those values are constrained to those cells.
// 6) If any value only has one possible cell, that is that cell's value.
// 7) If any x values only have x possible cells, those cells only have those
//  possible values - those cells are constrained to those values.
//
// Additional Helper functions are included and explained later.

// indexedCLuster is a datatype for an index for the values of the cluster.
// Each possible value is a key in the map. THe value of each key is an array of
// possible locations for that value - the index and order are not defined,
// instead the values of the array are the indexes of possible cells in for that
// value.
type intArray []int
type indexedCluster map[int]intArray

func indexCluster(in []cell) (out indexedCluster) {
	for id, each := range in {
		for _, onePossible_ := range each.possible {
			out[onePossible] = append(out[onePossible], id)
		}
	}
	return out
}

// This covers rule 1 from above:
// 1) If all cells are solved, that cluster is solved.
func clusterSolved(cluster []cell, u chan cell) (solved bool) {
	solved = true
	for _, each := range cluster {
		if each.actual == 0 {
			solved = false
		} else if len(each.possible) > 0 {
			solved = false
		}
	}
	return solved
}

// This covers rule 2 from above:
// 2) If any cell is solved, it has no possibles.
func solvedNoPossible(cluster []cell, u chan cell) (changed bool) {
	for _, each := range cluster {
		if each.actual == 0 {
			continue
		}
		if len(each.possible) > 0 {
			changed = true
			u <- cell{
				location: each.location,
				possible: each.possible,
			}
		}
	}
	return changed
}

// Removes known values from other possibles in the same cluster
// Covers rule 3 from above:
// 3) If any cell is solved, that value is not possible in other cells.
func eliminateKnowns(workingCluster []cell, u chan cell) (changed bool) {
	var knownValues []int

	// Loop thru and find all solved values.
	for _, each := range workingCluster {
		if each.actual != 0 {
			knownValues = append(knownValues, each.actual)
		}
	}

	for _, each := range workingCluster {
		if anyInArr(each, knownValues) {
			u <- cell{
				location: each.location,
				possible: subArr(each, knownValues),
			}
			changed = true
		}
	}
	return changed
}

// This covers the 4th rule from above:
// 4) If any cell only has one possible value, that is that cell's value.
func singleValueSolver(cluster []cell, u chan cell) (changed bool) {
	for _, each := range cluster {
		// if the cell is already solved, skip this.
		if each.actual != 0 {
			continue
		}

		// if there is more than one possible value for this cell, you should be good
		if len(each.possible) > 1 {
			continue
		}

		// should never happen - probably #TODO# to catch this
		if len(each.possible) < 1 {
			panic("Found an unsolved cell with no possible values")
		}

		changed = true

		// send back an update for this cell
		u <- cell{
			location: each.location,
			actual:   each.possible[0],
			possible: []int{},
		}
	}
	return changed
}

// A helper function to determine the values certain cells hit
func valuesPainted(markedCells []int, cluster []cell) (neededValues []int) {
	for _, oneCell := range markedCells {
		neededValues = addArr(neededValues, cluster[oneCell].possible)
	}
	return neededValues
}

// A helper function to determine the number of valus hit given a specific set
// of cells.
func cellsCost(markedCells []int, cluster []cell) int {
	return len(valuesPainted(markedCells, cluster))
}

func cellLimiterChild(limit int, markedCells []int, cluster []cell, u chan cell) (changed bool) {
	valueCount := cellsCost(markedCells, cluster)
	if valueCount < len(markedCells) {
		// #TODO# probably fix this? rework into error?
		panic("less possible values than squares to fill")
	}
	// you have overspent - it's a no-go
	if len(markedCells) > limit {
		return false
	}

	// did you fill the cells? if so, mark it
	if valueCount == len(markedCells) {
		valuesCovered := valuesPainted(markedCells, cluster)
		// it's a match - so remove values
		for id, each := range cluster {
			if each.actual != 0 {
				// already solved
				continue
			}
			if inArr(markedCells, id) {
				// this cell is a part of the list - no exclusions needed
				continue
			}

			toRemove := andArr(each.possible, valuesCovered)
			if len(toRemove) > 0 {
				changed = true
				u <- cell{
					location: each.location,
					possible: remove,
				}
			}
		}
	}

	// you have room to add more cells (depth first?)
	if len(markedCells) < limit {
		// you need to try adding each other cell
		for idCell, oneCell := range cluster {
			if oneCell.actual != 0 {
				continue
			}
			if inArr(markedCells, idCell) {
				// this cell is already in the map, skip
				continue
			}

			// decend down into looking at that cell
			changed = changed || cellLimiterChild(limit, append(markedCells, idCell), cluster, u)
		}
	}

	return changed
}

// This covers rule 5:
// 5) If any x cells have x possible values, those values are not possible
//  outside of those cells - those values are constrained to those cells.
func cellLimiter(cluster []cell, u chan cell) (changed bool) {
	upperBound := len(cluster)
	for _, eachCell := range cluster {
		if eachCell.actual != 0 {
			upperBound--
		}
	}
	for i := 2; i <= upperBound; i++ {
		changed = changed || cellLimiterChild(i, []int{}, cluster, u)
	}
	return changed
}

// This covers rule 6 from above:
// 6) If any value only has one possible cell, that is that cell's value.
func singleCellSolver(index indexedCluster, cluster []cell, u chan cell) (changed bool) {
	for val, section := range index {
		if len(section) < 1 {
			// something went terribly wrong here - #TODO# add panic catch?
			panic("Found an unsolved cell with no possible values")
		} else if len(section) == 1 {
			u <- cell{
				location: cluster[section[0]].location,
				actual:   val,
				possible: []int{},
			}
			changed = true
		}
	}
	return changed
}

// A helper function to determine the number of cells hit by working across a map
// of values.
func valuesCost(markedVals map[int]bool, index indexedCluster, cluseter []cell) int {
	var neededCells map[int]bool
	for value, _ := range markedVals {
		for oneCell, _ := range index[value] {
			neededCells[oneCell] = true
		}
	}
	return len(neededCells)
}

func valueLimiterChild(limit int, markedValues map[int]bool, index indexedCluster,
	cluster []cell, u chan cell) (changed bool) {
	if len(markedValues) > limit {
		// we have marked more values
		return false
	}
	currentCost := valuesCost(markedValues, index, cluster)
	if currentCost > limit {
		// you're over the budget to spend
		return false
	}
	return changed
}

// THis covers rule 7 from above:
// 7) If any x values have only x possible cells, those cells only have those
//  possible values - those cells are constrained to those values.
/*
func valueLimiter(index indexedCluster, cluster []cell, u chan cell) (changed bool) {

}
*/
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/bcicen/ctop/config"
	"github.com/bcicen/ctop/container"
	"github.com/bcicen/ctop/widgets"
	"github.com/bcicen/ctop/widgets/menu"
	ui "github.com/gizak/termui"
)

// MenuFn executes a menu window, returning the next menu or nil
type MenuFn func() MenuFn

var helpDialog = []menu.Item{
	{"<enter> - open container menu", ""},
	{"", ""},
	{"[a] - toggle display of all containers", ""},
	{"[f] - filter displayed containers", ""},
	{"[h] - open this help dialog", ""},
	{"[H] - toggle ctop header", ""},
	{"[s] - select container sort field", ""},
	{"[r] - reverse container sort order", ""},
	{"[o] - open single view", ""},
	{"[l] - view container logs ([t] to toggle timestamp when open)", ""},
	{"[e] - exec shell", ""},
	{"[S] - save current configuration to file", ""},
	{"[q] - exit ctop", ""},
}

func HelpMenu() MenuFn {
	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.BorderLabel = "Help"
	m.AddItems(helpDialog...)
	ui.Handle("/sys/wnd/resize", func(e ui.Event) {
		ui.Clear()
		ui.Render(m)
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func FilterMenu() MenuFn {
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	i := widgets.NewInput()
	i.BorderLabel = "Filter"
	i.SetY(ui.TermHeight() - i.Height)
	i.Data = config.GetVal("filterStr")
	ui.Render(i)

	// refresh container rows on input
	stream := i.Stream()
	go func() {
		for s := range stream {
			config.Update("filterStr", s)
			RefreshDisplay()
			ui.Render(i)
		}
	}()

	i.InputHandlers()
	ui.Handle("/sys/kbd/<escape>", func(ui.Event) {
		config.Update("filterStr", "")
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		config.Update("filterStr", i.Data)
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func SortMenu() MenuFn {
	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.SortItems = true
	m.BorderLabel = "Sort Field"

	for _, field := range container.SortFields() {
		m.AddItems(menu.Item{field, ""})
	}

	// set cursor position to current sort field
	m.SetCursor(config.GetVal("sortField"))

	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)
	HandleKeys("exit", ui.StopLoop)

	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		config.Update("sortField", m.SelectedItem().Val)
		ui.StopLoop()
	})

	ui.Render(m)
	ui.Loop()
	return nil
}

func ContainerMenu() MenuFn {
	c := cursor.Selected()
	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.BorderLabel = "Menu"

	items := []menu.Item{
		menu.Item{Val: "single", Label: "single view"},
		menu.Item{Val: "logs", Label: "log view"},
	}

	if c.Meta["state"] == "running" {
		items = append(items, menu.Item{Val: "stop", Label: "stop"})
		items = append(items, menu.Item{Val: "pause", Label: "pause"})
		items = append(items, menu.Item{Val: "restart", Label: "restart"})
		items = append(items, menu.Item{Val: "exec shell", Label: "exec shell"})
	}
	if c.Meta["state"] == "exited" || c.Meta["state"] == "created" {
		items = append(items, menu.Item{Val: "start", Label: "start"})
		items = append(items, menu.Item{Val: "remove", Label: "remove"})
	}
	if c.Meta["state"] == "paused" {
		items = append(items, menu.Item{Val: "unpause", Label: "unpause"})
	}
	items = append(items, menu.Item{Val: "cancel", Label: "cancel"})

	m.AddItems(items...)
	ui.Render(m)

	var nextMenu MenuFn
	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)
	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		switch m.SelectedItem().Val {
		case "single":
			nextMenu = SingleView
		case "logs":
			nextMenu = LogMenu
		case "exec shell":
			nextMenu = ExecShell
		case "start":
			nextMenu = Confirm(confirmTxt("start", c.GetMeta("name")), c.Start)
		case "stop":
			nextMenu = Confirm(confirmTxt("stop", c.GetMeta("name")), c.Stop)
		case "remove":
			nextMenu = Confirm(confirmTxt("remove", c.GetMeta("name")), c.Remove)
		case "pause":
			nextMenu = Confirm(confirmTxt("pause", c.GetMeta("name")), c.Pause)
		case "unpause":
			nextMenu = Confirm(confirmTxt("unpause", c.GetMeta("name")), c.Unpause)
		case "restart":
			nextMenu = Confirm(confirmTxt("restart", c.GetMeta("name")), c.Restart)
		}
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		ui.StopLoop()
	})
	ui.Loop()
	return nextMenu
}

func LogMenu() MenuFn {

	c := cursor.Selected()
	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	logs, quit := logReader(c)
	m := widgets.NewTextView(logs)
	m.BorderLabel = fmt.Sprintf("Logs [%s]", c.GetMeta("name"))
	ui.Render(m)

	ui.Handle("/sys/wnd/resize", func(e ui.Event) {
		m.Resize()
	})
	ui.Handle("/sys/kbd/t", func(ui.Event) {
		m.Toggle()
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		quit <- true
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func ExecShell() MenuFn {
	c := cursor.Selected()

	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()
	ui.StopLoop()
	defer ui.Loop()

	shell := config.Get("shell")
	if err := c.Exec([]string{shell.Val, "-c", "echo '\033[0m' && clear && " + shell.Val}); err != nil {
		log.Fatal(err)
	}

	return nil
}

// Create a confirmation dialog with a given description string and
// func to perform if confirmed
func Confirm(txt string, fn func()) MenuFn {
	menu := func() MenuFn {
		ui.DefaultEvtStream.ResetHandlers()
		defer ui.DefaultEvtStream.ResetHandlers()

		m := menu.NewMenu()
		m.Selectable = true
		m.BorderLabel = "Confirm"
		m.SubText = txt

		items := []menu.Item{
			menu.Item{Val: "cancel", Label: "[c]ancel"},
			menu.Item{Val: "yes", Label: "[y]es"},
		}

		var response bool

		m.AddItems(items...)
		ui.Render(m)

		yes := func() {
			response = true
			ui.StopLoop()
		}

		no := func() {
			response = false
			ui.StopLoop()
		}

		HandleKeys("up", m.Up)
		HandleKeys("down", m.Down)
		HandleKeys("exit", no)
		ui.Handle("/sys/kbd/c", func(ui.Event) { no() })
		ui.Handle("/sys/kbd/y", func(ui.Event) { yes() })

		ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
			switch m.SelectedItem().Val {
			case "cancel":
				no()
			case "yes":
				yes()
			}
		})

		ui.Loop()
		if response {
			fn()
		}
		return nil
	}
	return menu
}

type toggleLog struct {
	timestamp time.Time
	message   string
}

func (t *toggleLog) Toggle(on bool) string {
	if on {
		return fmt.Sprintf("%s %s", t.timestamp.Format("2006-01-02T15:04:05.999Z07:00"), t.message)
	}
	return t.message
}

func logReader(container *container.Container) (logs chan widgets.ToggleText, quit chan bool) {

	logCollector := container.Logs()
	stream := logCollector.Stream()
	logs = make(chan widgets.ToggleText)
	quit = make(chan bool)

	go func() {
		for {
			select {
			case log := <-stream:
				logs <- &toggleLog{timestamp: log.Timestamp, message: log.Message}
			case <-quit:
				logCollector.Stop()
				close(logs)
				return
			}
		}
	}()
	return
}

func confirmTxt(a, n string) string { return fmt.Sprintf("%s container %s?", a, n) }
<commit_msg>Unnecessary loop stopping<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/bcicen/ctop/config"
	"github.com/bcicen/ctop/container"
	"github.com/bcicen/ctop/widgets"
	"github.com/bcicen/ctop/widgets/menu"
	ui "github.com/gizak/termui"
)

// MenuFn executes a menu window, returning the next menu or nil
type MenuFn func() MenuFn

var helpDialog = []menu.Item{
	{"<enter> - open container menu", ""},
	{"", ""},
	{"[a] - toggle display of all containers", ""},
	{"[f] - filter displayed containers", ""},
	{"[h] - open this help dialog", ""},
	{"[H] - toggle ctop header", ""},
	{"[s] - select container sort field", ""},
	{"[r] - reverse container sort order", ""},
	{"[o] - open single view", ""},
	{"[l] - view container logs ([t] to toggle timestamp when open)", ""},
	{"[e] - exec shell", ""},
	{"[S] - save current configuration to file", ""},
	{"[q] - exit ctop", ""},
}

func HelpMenu() MenuFn {
	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.BorderLabel = "Help"
	m.AddItems(helpDialog...)
	ui.Handle("/sys/wnd/resize", func(e ui.Event) {
		ui.Clear()
		ui.Render(m)
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func FilterMenu() MenuFn {
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	i := widgets.NewInput()
	i.BorderLabel = "Filter"
	i.SetY(ui.TermHeight() - i.Height)
	i.Data = config.GetVal("filterStr")
	ui.Render(i)

	// refresh container rows on input
	stream := i.Stream()
	go func() {
		for s := range stream {
			config.Update("filterStr", s)
			RefreshDisplay()
			ui.Render(i)
		}
	}()

	i.InputHandlers()
	ui.Handle("/sys/kbd/<escape>", func(ui.Event) {
		config.Update("filterStr", "")
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		config.Update("filterStr", i.Data)
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func SortMenu() MenuFn {
	ui.Clear()
	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.SortItems = true
	m.BorderLabel = "Sort Field"

	for _, field := range container.SortFields() {
		m.AddItems(menu.Item{field, ""})
	}

	// set cursor position to current sort field
	m.SetCursor(config.GetVal("sortField"))

	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)
	HandleKeys("exit", ui.StopLoop)

	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		config.Update("sortField", m.SelectedItem().Val)
		ui.StopLoop()
	})

	ui.Render(m)
	ui.Loop()
	return nil
}

func ContainerMenu() MenuFn {
	c := cursor.Selected()
	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	m := menu.NewMenu()
	m.Selectable = true
	m.BorderLabel = "Menu"

	items := []menu.Item{
		menu.Item{Val: "single", Label: "single view"},
		menu.Item{Val: "logs", Label: "log view"},
	}

	if c.Meta["state"] == "running" {
		items = append(items, menu.Item{Val: "stop", Label: "stop"})
		items = append(items, menu.Item{Val: "pause", Label: "pause"})
		items = append(items, menu.Item{Val: "restart", Label: "restart"})
		items = append(items, menu.Item{Val: "exec shell", Label: "exec shell"})
	}
	if c.Meta["state"] == "exited" || c.Meta["state"] == "created" {
		items = append(items, menu.Item{Val: "start", Label: "start"})
		items = append(items, menu.Item{Val: "remove", Label: "remove"})
	}
	if c.Meta["state"] == "paused" {
		items = append(items, menu.Item{Val: "unpause", Label: "unpause"})
	}
	items = append(items, menu.Item{Val: "cancel", Label: "cancel"})

	m.AddItems(items...)
	ui.Render(m)

	var nextMenu MenuFn
	HandleKeys("up", m.Up)
	HandleKeys("down", m.Down)
	ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
		switch m.SelectedItem().Val {
		case "single":
			nextMenu = SingleView
		case "logs":
			nextMenu = LogMenu
		case "exec shell":
			nextMenu = ExecShell
		case "start":
			nextMenu = Confirm(confirmTxt("start", c.GetMeta("name")), c.Start)
		case "stop":
			nextMenu = Confirm(confirmTxt("stop", c.GetMeta("name")), c.Stop)
		case "remove":
			nextMenu = Confirm(confirmTxt("remove", c.GetMeta("name")), c.Remove)
		case "pause":
			nextMenu = Confirm(confirmTxt("pause", c.GetMeta("name")), c.Pause)
		case "unpause":
			nextMenu = Confirm(confirmTxt("unpause", c.GetMeta("name")), c.Unpause)
		case "restart":
			nextMenu = Confirm(confirmTxt("restart", c.GetMeta("name")), c.Restart)
		}
		ui.StopLoop()
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		ui.StopLoop()
	})
	ui.Loop()
	return nextMenu
}

func LogMenu() MenuFn {

	c := cursor.Selected()
	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	logs, quit := logReader(c)
	m := widgets.NewTextView(logs)
	m.BorderLabel = fmt.Sprintf("Logs [%s]", c.GetMeta("name"))
	ui.Render(m)

	ui.Handle("/sys/wnd/resize", func(e ui.Event) {
		m.Resize()
	})
	ui.Handle("/sys/kbd/t", func(ui.Event) {
		m.Toggle()
	})
	ui.Handle("/sys/kbd/", func(ui.Event) {
		quit <- true
		ui.StopLoop()
	})
	ui.Loop()
	return nil
}

func ExecShell() MenuFn {
	c := cursor.Selected()

	if c == nil {
		return nil
	}

	ui.DefaultEvtStream.ResetHandlers()
	defer ui.DefaultEvtStream.ResetHandlers()

	shell := config.Get("shell")
	if err := c.Exec([]string{shell.Val, "-c", "echo '\033[0m' && clear && " + shell.Val}); err != nil {
		log.Fatal(err)
	}

	return nil
}

// Create a confirmation dialog with a given description string and
// func to perform if confirmed
func Confirm(txt string, fn func()) MenuFn {
	menu := func() MenuFn {
		ui.DefaultEvtStream.ResetHandlers()
		defer ui.DefaultEvtStream.ResetHandlers()

		m := menu.NewMenu()
		m.Selectable = true
		m.BorderLabel = "Confirm"
		m.SubText = txt

		items := []menu.Item{
			menu.Item{Val: "cancel", Label: "[c]ancel"},
			menu.Item{Val: "yes", Label: "[y]es"},
		}

		var response bool

		m.AddItems(items...)
		ui.Render(m)

		yes := func() {
			response = true
			ui.StopLoop()
		}

		no := func() {
			response = false
			ui.StopLoop()
		}

		HandleKeys("up", m.Up)
		HandleKeys("down", m.Down)
		HandleKeys("exit", no)
		ui.Handle("/sys/kbd/c", func(ui.Event) { no() })
		ui.Handle("/sys/kbd/y", func(ui.Event) { yes() })

		ui.Handle("/sys/kbd/<enter>", func(ui.Event) {
			switch m.SelectedItem().Val {
			case "cancel":
				no()
			case "yes":
				yes()
			}
		})

		ui.Loop()
		if response {
			fn()
		}
		return nil
	}
	return menu
}

type toggleLog struct {
	timestamp time.Time
	message   string
}

func (t *toggleLog) Toggle(on bool) string {
	if on {
		return fmt.Sprintf("%s %s", t.timestamp.Format("2006-01-02T15:04:05.999Z07:00"), t.message)
	}
	return t.message
}

func logReader(container *container.Container) (logs chan widgets.ToggleText, quit chan bool) {

	logCollector := container.Logs()
	stream := logCollector.Stream()
	logs = make(chan widgets.ToggleText)
	quit = make(chan bool)

	go func() {
		for {
			select {
			case log := <-stream:
				logs <- &toggleLog{timestamp: log.Timestamp, message: log.Message}
			case <-quit:
				logCollector.Stop()
				close(logs)
				return
			}
		}
	}()
	return
}

func confirmTxt(a, n string) string { return fmt.Sprintf("%s container %s?", a, n) }
<|endoftext|>


<commit_before>// Package outbound implements multiple dns client and dispatcher for outbound connection.
package outbound

import (
	"math/rand"
	"net"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/miekg/dns"
	"github.com/shawn1m/overture/core/cache"
	"github.com/shawn1m/overture/core/common"
	"github.com/shawn1m/overture/core/hosts"
	"golang.org/x/net/proxy"
)

type Client struct {
	ResponseMessage *dns.Msg
	QuestionMessage *dns.Msg

	DNSUpstream           *DNSUpstream
	EDNSClientSubnetIP    string
	InboundIP             string
	ReservedIPNetworkList []*net.IPNet

	Hosts *hosts.Hosts
	Cache *cache.Cache
}

func NewClient(q *dns.Msg, u *DNSUpstream, ip string, h *hosts.Hosts, cache *cache.Cache) *Client {

	c := &Client{QuestionMessage: q, DNSUpstream: u, InboundIP: ip, Hosts: h, Cache: cache}
	c.getEDNSClientSubnetIP()
	c.ReservedIPNetworkList = getReservedIPNetworkList()
	return c
}

func (c *Client) getEDNSClientSubnetIP() {

	switch c.DNSUpstream.EDNSClientSubnet.Policy {
	case "auto":
		if !common.IsIPMatchList(net.ParseIP(c.InboundIP), c.ReservedIPNetworkList, false) {
			c.EDNSClientSubnetIP = c.InboundIP
		} else {
			c.EDNSClientSubnetIP = c.DNSUpstream.EDNSClientSubnet.ExternalIP
		}
	case "disable":
	}
}

func (c *Client) ExchangeFromRemote(isCache bool, isLog bool) {

	if c.ExchangeFromCache(isLog) {
		return
	}

	setEDNSClientSubnet(c.QuestionMessage, c.EDNSClientSubnetIP)

	var conn net.Conn
	if c.DNSUpstream.SOCKS5Address != "" {
		s, err := proxy.SOCKS5(c.DNSUpstream.Protocol, c.DNSUpstream.SOCKS5Address, nil, proxy.Direct)
		if err != nil {
			log.Warn("Get socks5 proxy dialer failed: ", err)
			return
		}
		conn, err = s.Dial(c.DNSUpstream.Protocol, c.DNSUpstream.Address)
		if err != nil {
			log.Warn("Dial DNS upstream with SOCKS5 proxy failed: ", err)
			return
		}
	} else {
		var err error
		if conn, err = net.Dial(c.DNSUpstream.Protocol, c.DNSUpstream.Address); err != nil {
			log.Warn("Dial DNS upstream failed: ", err)
			return
		}
	}

	dnsTimeout := time.Duration(c.DNSUpstream.Timeout) * time.Second / 3

	conn.SetDeadline(time.Now().Add(dnsTimeout))
	conn.SetReadDeadline(time.Now().Add(dnsTimeout))
	conn.SetWriteDeadline(time.Now().Add(dnsTimeout))

	dc := &dns.Conn{Conn: conn}
	defer dc.Close()
	err := dc.WriteMsg(c.QuestionMessage)
	if err != nil {
		log.Warn(c.DNSUpstream.Name + " Fail: Send question message failed")
		return
	}
	temp, err := dc.ReadMsg()

	if err != nil {
		if err == dns.ErrTruncated {
			log.Warn(c.DNSUpstream.Name + " Fail: Maybe this server does not support EDNS Client Subnet")
			return
		}
	}
	if temp == nil {
		log.Debug(c.DNSUpstream.Name + " Fail: Response message is nil, maybe timeout, please check your query or dns configuration")
		return
	}

	c.ResponseMessage = temp

	if isCache && c.Cache != nil {
		c.Cache.InsertMessage(cache.Key(c.QuestionMessage.Question[0], c.EDNSClientSubnetIP), c.ResponseMessage)
	}

	if isLog {
		c.logAnswer(false)
	}
}

func (c *Client) ExchangeFromCache(isLog bool) bool {

	if c.Cache == nil {
		return false
	}

	m := c.Cache.Hit(cache.Key(c.QuestionMessage.Question[0], c.EDNSClientSubnetIP), c.QuestionMessage.Id)
	if m != nil {
		log.Debug(c.DNSUpstream.Name + " Hit: " + cache.Key(c.QuestionMessage.Question[0], c.EDNSClientSubnetIP))
		c.ResponseMessage = m
		if isLog {
			c.logAnswer(false)
		}
		return true
	}

	return false
}

func (c *Client) ExchangeFromLocal() bool {
	raw_name := c.QuestionMessage.Question[0].Name

	if c.ExchangeFromHosts(raw_name) || c.ExchangeFromIP(raw_name) {
		return true
	}

	return false
}

func (c *Client) ExchangeFromHosts(raw_name string) bool {

	if c.Hosts == nil {
		return false
	}

	name := raw_name[:len(raw_name)-1]
	ipv4List, ipv6List := c.Hosts.Find(name)

	if c.QuestionMessage.Question[0].Qtype == dns.TypeA && len(ipv4List) > 0 {
		var rrl []dns.RR
		for _, ip := range ipv4List {
			a, _ := dns.NewRR(raw_name + " IN A " + ip.String())
			rrl = append(rrl, a)
		}
		c.setLocalResponseMessage(rrl)
		if c.ResponseMessage != nil {
			return true
		}
	} else if c.QuestionMessage.Question[0].Qtype == dns.TypeAAAA && len(ipv6List) > 0 {
		var rrl []dns.RR
		for _, ip := range ipv6List {
			aaaa, _ := dns.NewRR(raw_name + " IN AAAA " + ip.String())
			rrl = append(rrl, aaaa)
		}
		c.setLocalResponseMessage(rrl)
		if c.ResponseMessage != nil {
			return true
		}
	}

	return false
}

func (c *Client) ExchangeFromIP(raw_name string) bool {

	name := raw_name[:len(raw_name)-1]
	ip := net.ParseIP(name)
	if ip == nil {
		return false
	}
	if ip.To4() == nil && ip.To16() != nil && c.QuestionMessage.Question[0].Qtype == dns.TypeAAAA {
		aaaa, _ := dns.NewRR(raw_name + " IN AAAA " + ip.String())
		c.setLocalResponseMessage([]dns.RR{aaaa})
		return true
	} else if ip.To4() != nil && c.QuestionMessage.Question[0].Qtype == dns.TypeA {
		a, _ := dns.NewRR(raw_name + " IN A " + ip.String())
		c.setLocalResponseMessage([]dns.RR{a})
		return true
	}

	return false
}

func (c *Client) setLocalResponseMessage(rrl []dns.RR) {

	shuffleRRList := func(rrl []dns.RR) {
		rand.Seed(time.Now().UnixNano())
		for i := range rrl {
			j := rand.Intn(i + 1)
			rrl[i], rrl[j] = rrl[j], rrl[i]
		}
	}
	c.ResponseMessage = new(dns.Msg)
	for _, rr := range rrl {
		c.ResponseMessage.Answer = append(c.ResponseMessage.Answer, rr)
	}
	shuffleRRList(c.ResponseMessage.Answer)
	c.ResponseMessage.SetReply(c.QuestionMessage)
	c.ResponseMessage.RecursionAvailable = true
}

func (c *Client) logAnswer(isLocal bool) {

	for _, a := range c.ResponseMessage.Answer {
		var name string
		if isLocal {
			name = "Local"
		} else {
			name = c.DNSUpstream.Name
		}
		log.Debug(name + " Answer: " + a.String())
	}
}

func getReservedIPNetworkList() []*net.IPNet {

	ipnl := make([]*net.IPNet, 0)
	localCIDR := []string{"127.0.0.0/8", "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "100.64.0.0/10"}
	for _, c := range localCIDR {
		_, ip_net, err := net.ParseCIDR(c)
		if err != nil {
			break
		}
		ipnl = append(ipnl, ip_net)
	}
	return ipnl
}
<commit_msg>getEDNSClientSubnetIP with empty ReservedIP...List (#44)<commit_after>// Package outbound implements multiple dns client and dispatcher for outbound connection.
package outbound

import (
	"math/rand"
	"net"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/miekg/dns"
	"github.com/shawn1m/overture/core/cache"
	"github.com/shawn1m/overture/core/common"
	"github.com/shawn1m/overture/core/hosts"
	"golang.org/x/net/proxy"
)

type Client struct {
	ResponseMessage *dns.Msg
	QuestionMessage *dns.Msg

	DNSUpstream           *DNSUpstream
	EDNSClientSubnetIP    string
	InboundIP             string

	Hosts *hosts.Hosts
	Cache *cache.Cache
}

var ReservedIPNetworkList = getReservedIPNetworkList()

func NewClient(q *dns.Msg, u *DNSUpstream, ip string, h *hosts.Hosts, cache *cache.Cache) *Client {

	c := &Client{QuestionMessage: q, DNSUpstream: u, InboundIP: ip, Hosts: h, Cache: cache}
        c.getEDNSClientSubnetIP()
	return c
}

func (c *Client) getEDNSClientSubnetIP() {

	switch c.DNSUpstream.EDNSClientSubnet.Policy {
	case "auto":
		if !common.IsIPMatchList(net.ParseIP(c.InboundIP), ReservedIPNetworkList, false) {
			c.EDNSClientSubnetIP = c.InboundIP
		} else {
			c.EDNSClientSubnetIP = c.DNSUpstream.EDNSClientSubnet.ExternalIP
		}
	case "disable":
	}
}

func (c *Client) ExchangeFromRemote(isCache bool, isLog bool) {

	if c.ExchangeFromCache(isLog) {
		return
	}

	setEDNSClientSubnet(c.QuestionMessage, c.EDNSClientSubnetIP)

	var conn net.Conn
	if c.DNSUpstream.SOCKS5Address != "" {
		s, err := proxy.SOCKS5(c.DNSUpstream.Protocol, c.DNSUpstream.SOCKS5Address, nil, proxy.Direct)
		if err != nil {
			log.Warn("Get socks5 proxy dialer failed: ", err)
			return
		}
		conn, err = s.Dial(c.DNSUpstream.Protocol, c.DNSUpstream.Address)
		if err != nil {
			log.Warn("Dial DNS upstream with SOCKS5 proxy failed: ", err)
			return
		}
	} else {
		var err error
		if conn, err = net.Dial(c.DNSUpstream.Protocol, c.DNSUpstream.Address); err != nil {
			log.Warn("Dial DNS upstream failed: ", err)
			return
		}
	}

	dnsTimeout := time.Duration(c.DNSUpstream.Timeout) * time.Second / 3

	conn.SetDeadline(time.Now().Add(dnsTimeout))
	conn.SetReadDeadline(time.Now().Add(dnsTimeout))
	conn.SetWriteDeadline(time.Now().Add(dnsTimeout))

	dc := &dns.Conn{Conn: conn}
	defer dc.Close()
	err := dc.WriteMsg(c.QuestionMessage)
	if err != nil {
		log.Warn(c.DNSUpstream.Name + " Fail: Send question message failed")
		return
	}
	temp, err := dc.ReadMsg()

	if err != nil {
		if err == dns.ErrTruncated {
			log.Warn(c.DNSUpstream.Name + " Fail: Maybe this server does not support EDNS Client Subnet")
			return
		}
	}
	if temp == nil {
		log.Debug(c.DNSUpstream.Name + " Fail: Response message is nil, maybe timeout, please check your query or dns configuration")
		return
	}

	c.ResponseMessage = temp

	if isCache && c.Cache != nil {
		c.Cache.InsertMessage(cache.Key(c.QuestionMessage.Question[0], c.EDNSClientSubnetIP), c.ResponseMessage)
	}

	if isLog {
		c.logAnswer(false)
	}
}

func (c *Client) ExchangeFromCache(isLog bool) bool {

	if c.Cache == nil {
		return false
	}

	m := c.Cache.Hit(cache.Key(c.QuestionMessage.Question[0], c.EDNSClientSubnetIP), c.QuestionMessage.Id)
	if m != nil {
		log.Debug(c.DNSUpstream.Name + " Hit: " + cache.Key(c.QuestionMessage.Question[0], c.EDNSClientSubnetIP))
		c.ResponseMessage = m
		if isLog {
			c.logAnswer(false)
		}
		return true
	}

	return false
}

func (c *Client) ExchangeFromLocal() bool {
	raw_name := c.QuestionMessage.Question[0].Name

	if c.ExchangeFromHosts(raw_name) || c.ExchangeFromIP(raw_name) {
		return true
	}

	return false
}

func (c *Client) ExchangeFromHosts(raw_name string) bool {

	if c.Hosts == nil {
		return false
	}

	name := raw_name[:len(raw_name)-1]
	ipv4List, ipv6List := c.Hosts.Find(name)

	if c.QuestionMessage.Question[0].Qtype == dns.TypeA && len(ipv4List) > 0 {
		var rrl []dns.RR
		for _, ip := range ipv4List {
			a, _ := dns.NewRR(raw_name + " IN A " + ip.String())
			rrl = append(rrl, a)
		}
		c.setLocalResponseMessage(rrl)
		if c.ResponseMessage != nil {
			return true
		}
	} else if c.QuestionMessage.Question[0].Qtype == dns.TypeAAAA && len(ipv6List) > 0 {
		var rrl []dns.RR
		for _, ip := range ipv6List {
			aaaa, _ := dns.NewRR(raw_name + " IN AAAA " + ip.String())
			rrl = append(rrl, aaaa)
		}
		c.setLocalResponseMessage(rrl)
		if c.ResponseMessage != nil {
			return true
		}
	}

	return false
}

func (c *Client) ExchangeFromIP(raw_name string) bool {

	name := raw_name[:len(raw_name)-1]
	ip := net.ParseIP(name)
	if ip == nil {
		return false
	}
	if ip.To4() == nil && ip.To16() != nil && c.QuestionMessage.Question[0].Qtype == dns.TypeAAAA {
		aaaa, _ := dns.NewRR(raw_name + " IN AAAA " + ip.String())
		c.setLocalResponseMessage([]dns.RR{aaaa})
		return true
	} else if ip.To4() != nil && c.QuestionMessage.Question[0].Qtype == dns.TypeA {
		a, _ := dns.NewRR(raw_name + " IN A " + ip.String())
		c.setLocalResponseMessage([]dns.RR{a})
		return true
	}

	return false
}

func (c *Client) setLocalResponseMessage(rrl []dns.RR) {

	shuffleRRList := func(rrl []dns.RR) {
		rand.Seed(time.Now().UnixNano())
		for i := range rrl {
			j := rand.Intn(i + 1)
			rrl[i], rrl[j] = rrl[j], rrl[i]
		}
	}
	c.ResponseMessage = new(dns.Msg)
	for _, rr := range rrl {
		c.ResponseMessage.Answer = append(c.ResponseMessage.Answer, rr)
	}
	shuffleRRList(c.ResponseMessage.Answer)
	c.ResponseMessage.SetReply(c.QuestionMessage)
	c.ResponseMessage.RecursionAvailable = true
}

func (c *Client) logAnswer(isLocal bool) {

	for _, a := range c.ResponseMessage.Answer {
		var name string
		if isLocal {
			name = "Local"
		} else {
			name = c.DNSUpstream.Name
		}
		log.Debug(name + " Answer: " + a.String())
	}
}

func getReservedIPNetworkList() []*net.IPNet {

	ipnl := make([]*net.IPNet, 0)
	localCIDR := []string{"127.0.0.0/8", "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "100.64.0.0/10"}
	for _, c := range localCIDR {
		_, ip_net, err := net.ParseCIDR(c)
		if err != nil {
			break
		}
		ipnl = append(ipnl, ip_net)
	}
	return ipnl
}
<|endoftext|>