<commit_before>package render

import (
	"bufio"
	"bytes"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"time"

	"github.com/gogo/protobuf/proto"
	"go.uber.org/zap"

	"github.com/lomik/graphite-clickhouse/carbonzipperpb"
	"github.com/lomik/graphite-clickhouse/config"
	"github.com/lomik/graphite-clickhouse/finder"
	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
	"github.com/lomik/graphite-clickhouse/helper/log"
	"github.com/lomik/graphite-clickhouse/helper/pickle"
	"github.com/lomik/graphite-clickhouse/helper/point"
)

type Handler struct {
	config *config.Config
}

func NewHandler(config *config.Config) *Handler {
	return &Handler{
		config: config,
	}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	logger := log.FromContext(r.Context())
	target := r.URL.Query().Get("target")

	var prefix string
	var err error

	fromTimestamp, err := strconv.ParseInt(r.URL.Query().Get("from"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	untilTimestamp, err := strconv.ParseInt(r.URL.Query().Get("until"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	// Search in small index table first
	finder := finder.New(r.Context(), h.config)

	err = finder.Execute(target)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	metricList := finder.Series()

	maxStep := int32(0)

	listBuf := bytes.NewBuffer(nil)

	// make Path IN (...), calculate max step
	for index, m := range metricList {
		if len(m) == 0 {
			continue
		}
		step := h.config.Rollup.Step(unsafeString(m), int32(fromTimestamp))
		if step > maxStep {
			maxStep = step
		}

		if index > 0 {
			listBuf.Write([]byte{','})
		}

		listBuf.WriteString("'" + clickhouse.Escape(unsafeString(m)) + "'")
	}

	if listBuf.Len() == 0 {
		// Return empty response
		h.Reply(w, r, &Data{Points: make([]point.Point, 0), Finder: finder}, 0, 0, "")
		return
	}

	var pathWhere = fmt.Sprintf(
		"Path IN (%s)",
		string(listBuf.Bytes()),
	)

	until := untilTimestamp - untilTimestamp%int64(maxStep) + int64(maxStep) - 1
	dateWhere := fmt.Sprintf(
		"(Date >='%s' AND Date <= '%s' AND Time >= %d AND Time <= %d)",
		time.Unix(fromTimestamp, 0).Format("2006-01-02"),
		time.Unix(untilTimestamp, 0).Format("2006-01-02"),
		fromTimestamp,
		until,
	)

	// @TODO: change format to RowBinary
	query := fmt.Sprintf(
		`
		SELECT
			Path, Time, Value, Timestamp
		FROM %s
		WHERE (%s) AND (%s)
		FORMAT RowBinary
		`,
		h.config.ClickHouse.DataTable,
		pathWhere,
		dateWhere,
	)

	body, err := clickhouse.Query(
		r.Context(),
		h.config.ClickHouse.Url,
		query,
		h.config.ClickHouse.DataTimeout.Value(),
	)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	parseStart := time.Now()

	data, err := DataParse(body)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	d := time.Since(parseStart)
	logger.Debug("parse", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))

	sortStart := time.Now()
	sort.Sort(data)
	d = time.Since(sortStart)
	logger.Debug("sort", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))

	data.Points = point.Uniq(data.Points)
	data.Finder = finder

	// pp.Println(points)
	h.Reply(w, r, data, int32(fromTimestamp), int32(untilTimestamp), prefix)
}

func (h *Handler) Reply(w http.ResponseWriter, r *http.Request, data *Data, from, until int32, prefix string) {
	start := time.Now()
	switch r.URL.Query().Get("format") {
	case "pickle":
		h.ReplyPickle(w, r, data, from, until, prefix)
	case "protobuf":
		h.ReplyProtobuf(w, r, data, from, until, prefix)
	}
	d := time.Since(start)
	log.FromContext(r.Context()).Debug("reply", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))
}

func (h *Handler) ReplyPickle(w http.ResponseWriter, r *http.Request, data *Data, from, until int32, prefix string) {
	var rollupTime time.Duration
	var pickleTime time.Duration

	points := data.Points

	defer func() {
		log.FromContext(r.Context()).Debug("rollup",
			zap.String("runtime", rollupTime.String()),
			zap.Duration("runtime_ns", rollupTime),
		)
		log.FromContext(r.Context()).Debug("pickle",
			zap.String("runtime", pickleTime.String()),
			zap.Duration("runtime_ns", pickleTime),
		)
	}()

	if len(points) == 0 {
		w.Write(pickle.EmptyList)
		return
	}

	writer := bufio.NewWriterSize(w, 1024*1024)
	p := pickle.NewWriter(writer)
	defer writer.Flush()

	p.List()

	writeMetric := func(points []point.Point) {
		rollupStart := time.Now()
		points, step := h.config.Rollup.RollupMetric(points)
		rollupTime += time.Since(rollupStart)

		pickleStart := time.Now()
		p.Dict()

		p.String("name")
		p.Bytes(data.Finder.Abs([]byte(points[0].Metric)))
		p.SetItem()

		p.String("step")
		p.Uint32(uint32(step))
		p.SetItem()

		start := from - (from % step)
		if start < from {
			start += step
		}
		end := until - (until % step)
		last := start - step

		p.String("values")
		p.List()
		for _, point := range points {
			if point.Time < start || point.Time > end {
				continue
			}

			if point.Time > last+step {
				p.AppendNulls(int(((point.Time - last) / step) - 1))
			}

			p.AppendFloat64(point.Value)

			last = point.Time
		}

		if end > last {
			p.AppendNulls(int((end - last) / step))
		}
		p.SetItem()

		p.String("start")
		p.Uint32(uint32(start))
		p.SetItem()

		p.String("end")
		p.Uint32(uint32(end))
		p.SetItem()

		p.Append()
		pickleTime += time.Since(pickleStart)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	p.Stop()
}

func (h *Handler) ReplyProtobuf(w http.ResponseWriter, r *http.Request, data *Data, from, until int32, prefix string) {
	points := data.Points

	if len(points) == 0 {
		return
	}

	var multiResponse carbonzipperpb.MultiFetchResponse

	writeMetric := func(points []point.Point) {
		points, step := h.config.Rollup.RollupMetric(points)

		var name string
		name = string(data.Finder.Abs([]byte(points[0].Metric)))

		start := from - (from % step)
		if start < from {
			start += step
		}
		stop := until - (until % step)
		count := ((stop - start) / step) + 1

		response := carbonzipperpb.FetchResponse{
			Name:      proto.String(name),
			StartTime: &start,
			StopTime:  &stop,
			StepTime:  &step,
			Values:    make([]float64, count),
			IsAbsent:  make([]bool, count),
		}

		var index int32
		// skip points before start
		for index = 0; index < int32(len(points)) && points[index].Time < start; index++ {
		}

		for i := int32(0); i < count; i++ {
			if index < int32(len(points)) && points[index].Time == start+step*i {
				response.Values[i] = points[index].Value
				response.IsAbsent[i] = false
				index++
			} else {
				response.Values[i] = 0
				response.IsAbsent[i] = true
			}
		}

		multiResponse.Metrics = append(multiResponse.Metrics, &response)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	body, _ := proto.Marshal(&multiResponse)
	w.Write(body)
}
<commit_msg>remove todo<commit_after>package render

import (
	"bufio"
	"bytes"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"time"

	"github.com/gogo/protobuf/proto"
	"go.uber.org/zap"

	"github.com/lomik/graphite-clickhouse/carbonzipperpb"
	"github.com/lomik/graphite-clickhouse/config"
	"github.com/lomik/graphite-clickhouse/finder"
	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
	"github.com/lomik/graphite-clickhouse/helper/log"
	"github.com/lomik/graphite-clickhouse/helper/pickle"
	"github.com/lomik/graphite-clickhouse/helper/point"
)

type Handler struct {
	config *config.Config
}

func NewHandler(config *config.Config) *Handler {
	return &Handler{
		config: config,
	}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	logger := log.FromContext(r.Context())
	target := r.URL.Query().Get("target")

	var prefix string
	var err error

	fromTimestamp, err := strconv.ParseInt(r.URL.Query().Get("from"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	untilTimestamp, err := strconv.ParseInt(r.URL.Query().Get("until"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	// Search in small index table first
	finder := finder.New(r.Context(), h.config)

	err = finder.Execute(target)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	metricList := finder.Series()

	maxStep := int32(0)

	listBuf := bytes.NewBuffer(nil)

	// make Path IN (...), calculate max step
	for index, m := range metricList {
		if len(m) == 0 {
			continue
		}
		step := h.config.Rollup.Step(unsafeString(m), int32(fromTimestamp))
		if step > maxStep {
			maxStep = step
		}

		if index > 0 {
			listBuf.Write([]byte{','})
		}

		listBuf.WriteString("'" + clickhouse.Escape(unsafeString(m)) + "'")
	}

	if listBuf.Len() == 0 {
		// Return empty response
		h.Reply(w, r, &Data{Points: make([]point.Point, 0), Finder: finder}, 0, 0, "")
		return
	}

	var pathWhere = fmt.Sprintf(
		"Path IN (%s)",
		string(listBuf.Bytes()),
	)

	until := untilTimestamp - untilTimestamp%int64(maxStep) + int64(maxStep) - 1
	dateWhere := fmt.Sprintf(
		"(Date >='%s' AND Date <= '%s' AND Time >= %d AND Time <= %d)",
		time.Unix(fromTimestamp, 0).Format("2006-01-02"),
		time.Unix(untilTimestamp, 0).Format("2006-01-02"),
		fromTimestamp,
		until,
	)

	query := fmt.Sprintf(
		`
		SELECT
			Path, Time, Value, Timestamp
		FROM %s
		WHERE (%s) AND (%s)
		FORMAT RowBinary
		`,
		h.config.ClickHouse.DataTable,
		pathWhere,
		dateWhere,
	)

	body, err := clickhouse.Query(
		r.Context(),
		h.config.ClickHouse.Url,
		query,
		h.config.ClickHouse.DataTimeout.Value(),
	)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	parseStart := time.Now()

	data, err := DataParse(body)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	d := time.Since(parseStart)
	logger.Debug("parse", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))

	sortStart := time.Now()
	sort.Sort(data)
	d = time.Since(sortStart)
	logger.Debug("sort", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))

	data.Points = point.Uniq(data.Points)
	data.Finder = finder

	// pp.Println(points)
	h.Reply(w, r, data, int32(fromTimestamp), int32(untilTimestamp), prefix)
}

func (h *Handler) Reply(w http.ResponseWriter, r *http.Request, data *Data, from, until int32, prefix string) {
	start := time.Now()
	switch r.URL.Query().Get("format") {
	case "pickle":
		h.ReplyPickle(w, r, data, from, until, prefix)
	case "protobuf":
		h.ReplyProtobuf(w, r, data, from, until, prefix)
	}
	d := time.Since(start)
	log.FromContext(r.Context()).Debug("reply", zap.String("runtime", d.String()), zap.Duration("runtime_ns", d))
}

func (h *Handler) ReplyPickle(w http.ResponseWriter, r *http.Request, data *Data, from, until int32, prefix string) {
	var rollupTime time.Duration
	var pickleTime time.Duration

	points := data.Points

	defer func() {
		log.FromContext(r.Context()).Debug("rollup",
			zap.String("runtime", rollupTime.String()),
			zap.Duration("runtime_ns", rollupTime),
		)
		log.FromContext(r.Context()).Debug("pickle",
			zap.String("runtime", pickleTime.String()),
			zap.Duration("runtime_ns", pickleTime),
		)
	}()

	if len(points) == 0 {
		w.Write(pickle.EmptyList)
		return
	}

	writer := bufio.NewWriterSize(w, 1024*1024)
	p := pickle.NewWriter(writer)
	defer writer.Flush()

	p.List()

	writeMetric := func(points []point.Point) {
		rollupStart := time.Now()
		points, step := h.config.Rollup.RollupMetric(points)
		rollupTime += time.Since(rollupStart)

		pickleStart := time.Now()
		p.Dict()

		p.String("name")
		p.Bytes(data.Finder.Abs([]byte(points[0].Metric)))
		p.SetItem()

		p.String("step")
		p.Uint32(uint32(step))
		p.SetItem()

		start := from - (from % step)
		if start < from {
			start += step
		}
		end := until - (until % step)
		last := start - step

		p.String("values")
		p.List()
		for _, point := range points {
			if point.Time < start || point.Time > end {
				continue
			}

			if point.Time > last+step {
				p.AppendNulls(int(((point.Time - last) / step) - 1))
			}

			p.AppendFloat64(point.Value)

			last = point.Time
		}

		if end > last {
			p.AppendNulls(int((end - last) / step))
		}
		p.SetItem()

		p.String("start")
		p.Uint32(uint32(start))
		p.SetItem()

		p.String("end")
		p.Uint32(uint32(end))
		p.SetItem()

		p.Append()
		pickleTime += time.Since(pickleStart)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	p.Stop()
}

func (h *Handler) ReplyProtobuf(w http.ResponseWriter, r *http.Request, data *Data, from, until int32, prefix string) {
	points := data.Points

	if len(points) == 0 {
		return
	}

	var multiResponse carbonzipperpb.MultiFetchResponse

	writeMetric := func(points []point.Point) {
		points, step := h.config.Rollup.RollupMetric(points)

		var name string
		name = string(data.Finder.Abs([]byte(points[0].Metric)))

		start := from - (from % step)
		if start < from {
			start += step
		}
		stop := until - (until % step)
		count := ((stop - start) / step) + 1

		response := carbonzipperpb.FetchResponse{
			Name:      proto.String(name),
			StartTime: &start,
			StopTime:  &stop,
			StepTime:  &step,
			Values:    make([]float64, count),
			IsAbsent:  make([]bool, count),
		}

		var index int32
		// skip points before start
		for index = 0; index < int32(len(points)) && points[index].Time < start; index++ {
		}

		for i := int32(0); i < count; i++ {
			if index < int32(len(points)) && points[index].Time == start+step*i {
				response.Values[i] = points[index].Value
				response.IsAbsent[i] = false
				index++
			} else {
				response.Values[i] = 0
				response.IsAbsent[i] = true
			}
		}

		multiResponse.Metrics = append(multiResponse.Metrics, &response)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	body, _ := proto.Marshal(&multiResponse)
	w.Write(body)
}
<|endoftext|>


<commit_before>package render

import (
	"bufio"
	"bytes"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/uber-go/zap"

	"github.com/lomik/graphite-clickhouse/carbonzipperpb"
	"github.com/lomik/graphite-clickhouse/config"
	"github.com/lomik/graphite-clickhouse/find"
	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
	"github.com/lomik/graphite-clickhouse/helper/log"
	"github.com/lomik/graphite-clickhouse/helper/pickle"
	"github.com/lomik/graphite-clickhouse/helper/point"
)

type Handler struct {
	config *config.Config
}

func NewHandler(config *config.Config) *Handler {
	return &Handler{
		config: config,
	}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	logger := log.FromContext(r.Context())
	target := r.URL.Query().Get("target")

	if strings.IndexByte(target, '\'') > -1 { // sql injection dumb fix
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	var prefix string
	var err error

	if h.config.ClickHouse.ExtraPrefix != "" {
		prefix, target, err = find.RemoveExtraPrefix(h.config.ClickHouse.ExtraPrefix, target)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		if target == "" {
			h.Reply(w, r, make([]point.Point, 0), 0, 0, "")
			return
		}
	}

	fromTimestamp, err := strconv.ParseInt(r.URL.Query().Get("from"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	untilTimestamp, err := strconv.ParseInt(r.URL.Query().Get("until"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	var pathWhere string

	maxStep := int32(0)
	if find.HasWildcard(target) {
		// Search in small index table first
		treeWhere := find.MakeWhere(target, true)
		if treeWhere == "" {
			http.Error(w, "Bad or unsupported query", http.StatusBadRequest)
			return
		}

		treeData, err := clickhouse.Query(
			r.Context(),
			h.config.ClickHouse.Url,
			fmt.Sprintf("SELECT Path FROM %s WHERE %s GROUP BY Path", h.config.ClickHouse.TreeTable, treeWhere),
			h.config.ClickHouse.TreeTimeout.Value(),
		)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		listBuf := bytes.NewBuffer(nil)
		first := true
		for _, p := range strings.Split(string(treeData), "\n") {
			if p == "" {
				continue
			}
			step := h.config.Rollup.Step(p, int32(fromTimestamp))
			if step > maxStep {
				maxStep = step
			}

			if !first {
				listBuf.Write([]byte{','})
			}
			first = false

			listBuf.WriteString("'" + p + "'") // SQL-Injection
		}

		if listBuf.Len() == 0 {
			h.Reply(w, r, make([]point.Point, 0), 0, 0, "")
			return
		}

		pathWhere = fmt.Sprintf(
			"Path IN (%s)",
			string(listBuf.Bytes()),
		)

		// pathWhere = fmt.Sprintf(
		// 	"Path IN (SELECT Path FROM %s WHERE %s)",
		// 	h.config.ClickHouse.DataTable,
		// )
		// pathWhere = makeWhere(target, false)
	} else {
		pathWhere = fmt.Sprintf("Path = '%s'", target)
		maxStep = h.config.Rollup.Step(target, int32(fromTimestamp))
	}

	until := untilTimestamp - untilTimestamp % int64(maxStep) + int64(maxStep) - 1
	dateWhere := fmt.Sprintf(
		"(Date >='%s' AND Date <= '%s' AND Time >= %d AND Time <= %d)",
		time.Unix(fromTimestamp, 0).Format("2006-01-02"),
		time.Unix(untilTimestamp, 0).Format("2006-01-02"),
		fromTimestamp,
		until,
	)

	// @TODO: change format to RowBinary
	query := fmt.Sprintf(
		`
		SELECT
			Path, Time, Value, Timestamp
		FROM %s
		WHERE (%s) AND (%s)
		FORMAT RowBinary
		`,
		h.config.ClickHouse.DataTable,
		pathWhere,
		dateWhere,
	)

	body, err := clickhouse.Query(
		r.Context(),
		h.config.ClickHouse.Url,
		query,
		h.config.ClickHouse.DataTimeout.Value(),
	)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	parseStart := time.Now()

	data, err := DataParse(body)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	logger.Debug("parse", zap.Duration("time_ns", time.Since(parseStart)))

	sortStart := time.Now()
	sort.Sort(data)
	logger.Debug("sort", zap.Duration("time_ns", time.Since(sortStart)))

	data.Points = point.Uniq(data.Points)

	// pp.Println(points)
	h.Reply(w, r, data.Points, int32(fromTimestamp), int32(untilTimestamp), prefix)
}

func (h *Handler) Reply(w http.ResponseWriter, r *http.Request, points []point.Point, from, until int32, prefix string) {
	start := time.Now()
	switch r.URL.Query().Get("format") {
	case "pickle":
		h.ReplyPickle(w, r, points, from, until, prefix)
	case "protobuf":
		h.ReplyProtobuf(w, r, points, from, until, prefix)
	}
	log.FromContext(r.Context()).Debug("reply", zap.Duration("time_ns", time.Since(start)))
}

func (h *Handler) ReplyPickle(w http.ResponseWriter, r *http.Request, points []point.Point, from, until int32, prefix string) {
	var rollupTime time.Duration
	var pickleTime time.Duration

	defer func() {
		log.FromContext(r.Context()).Debug("rollup", zap.Duration("time_ns", rollupTime))
		log.FromContext(r.Context()).Debug("pickle", zap.Duration("time_ns", pickleTime))
	}()

	if len(points) == 0 {
		w.Write(pickle.EmptyList)
		return
	}

	writer := bufio.NewWriterSize(w, 1024*1024)
	p := pickle.NewWriter(writer)
	defer writer.Flush()

	p.List()

	writeMetric := func(points []point.Point) {
		rollupStart := time.Now()
		points, step := h.config.Rollup.RollupMetric(points)
		rollupTime += time.Since(rollupStart)

		pickleStart := time.Now()
		p.Dict()

		p.String("name")
		if prefix != "" {
			p.String(prefix + "." + points[0].Metric)
		} else {
			p.String(points[0].Metric)
		}
		p.SetItem()

		p.String("step")
		p.Uint32(uint32(step))
		p.SetItem()

		start := from - (from % step)
		if start < from {
			start += step
		}
		end := until - (until % step)
		last := start - step

		p.String("values")
		p.List()
		for _, point := range points {
			if point.Time < start || point.Time > end {
				continue
			}

			if point.Time > last+step {
				p.AppendNulls(int(((point.Time - last) / step) - 1))
			}

			p.AppendFloat64(point.Value)

			last = point.Time
		}

		if end > last {
			p.AppendNulls(int((end - last) / step))
		}
		p.SetItem()

		p.String("start")
		p.Uint32(uint32(start))
		p.SetItem()

		p.String("end")
		p.Uint32(uint32(end))
		p.SetItem()

		p.Append()
		pickleTime += time.Since(pickleStart)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	p.Stop()
}

func (h *Handler) ReplyProtobuf(w http.ResponseWriter, r *http.Request, points []point.Point, from, until int32, prefix string) {
	if len(points) == 0 {
		return
	}

	var multiResponse carbonzipperpb.MultiFetchResponse

	writeMetric := func(points []point.Point) {
		points, step := h.config.Rollup.RollupMetric(points)

		var name string

		if prefix != "" {
			name = prefix + "." + points[0].Metric
		} else {
			name = points[0].Metric
		}

		start := from - (from % step)
		if start < from {
			start += step
		}
		stop := until - (until % step)
		count := ((stop - start) / step) + 1

		response := carbonzipperpb.FetchResponse{
			Name:      proto.String(name),
			StartTime: &start,
			StopTime:  &stop,
			StepTime:  &step,
			Values:    make([]float64, count),
			IsAbsent:  make([]bool, count),
		}

		var index int32
		// skip points before start
		for index = 0; points[index].Time < start; index++ {
		}

		for i := int32(0); i < count; i++ {
			if index < int32(len(points)) && points[index].Time == start+step*i {
				response.Values[i] = points[index].Value
				response.IsAbsent[i] = false
				index++
			} else {
				response.Values[i] = 0
				response.IsAbsent[i] = true
			}
		}

		multiResponse.Metrics = append(multiResponse.Metrics, &response)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	body, _ := proto.Marshal(&multiResponse)
	w.Write(body)
}
<commit_msg>Fix "out of bounds" error in protobuf serialization<commit_after>package render

import (
	"bufio"
	"bytes"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/uber-go/zap"

	"github.com/lomik/graphite-clickhouse/carbonzipperpb"
	"github.com/lomik/graphite-clickhouse/config"
	"github.com/lomik/graphite-clickhouse/find"
	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
	"github.com/lomik/graphite-clickhouse/helper/log"
	"github.com/lomik/graphite-clickhouse/helper/pickle"
	"github.com/lomik/graphite-clickhouse/helper/point"
)

type Handler struct {
	config *config.Config
}

func NewHandler(config *config.Config) *Handler {
	return &Handler{
		config: config,
	}
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	logger := log.FromContext(r.Context())
	target := r.URL.Query().Get("target")

	if strings.IndexByte(target, '\'') > -1 { // sql injection dumb fix
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	var prefix string
	var err error

	if h.config.ClickHouse.ExtraPrefix != "" {
		prefix, target, err = find.RemoveExtraPrefix(h.config.ClickHouse.ExtraPrefix, target)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		if target == "" {
			h.Reply(w, r, make([]point.Point, 0), 0, 0, "")
			return
		}
	}

	fromTimestamp, err := strconv.ParseInt(r.URL.Query().Get("from"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	untilTimestamp, err := strconv.ParseInt(r.URL.Query().Get("until"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	var pathWhere string

	maxStep := int32(0)
	if find.HasWildcard(target) {
		// Search in small index table first
		treeWhere := find.MakeWhere(target, true)
		if treeWhere == "" {
			http.Error(w, "Bad or unsupported query", http.StatusBadRequest)
			return
		}

		treeData, err := clickhouse.Query(
			r.Context(),
			h.config.ClickHouse.Url,
			fmt.Sprintf("SELECT Path FROM %s WHERE %s GROUP BY Path", h.config.ClickHouse.TreeTable, treeWhere),
			h.config.ClickHouse.TreeTimeout.Value(),
		)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		listBuf := bytes.NewBuffer(nil)
		first := true
		for _, p := range strings.Split(string(treeData), "\n") {
			if p == "" {
				continue
			}
			step := h.config.Rollup.Step(p, int32(fromTimestamp))
			if step > maxStep {
				maxStep = step
			}

			if !first {
				listBuf.Write([]byte{','})
			}
			first = false

			listBuf.WriteString("'" + p + "'") // SQL-Injection
		}

		if listBuf.Len() == 0 {
			h.Reply(w, r, make([]point.Point, 0), 0, 0, "")
			return
		}

		pathWhere = fmt.Sprintf(
			"Path IN (%s)",
			string(listBuf.Bytes()),
		)

		// pathWhere = fmt.Sprintf(
		// 	"Path IN (SELECT Path FROM %s WHERE %s)",
		// 	h.config.ClickHouse.DataTable,
		// )
		// pathWhere = makeWhere(target, false)
	} else {
		pathWhere = fmt.Sprintf("Path = '%s'", target)
		maxStep = h.config.Rollup.Step(target, int32(fromTimestamp))
	}

	until := untilTimestamp - untilTimestamp % int64(maxStep) + int64(maxStep) - 1
	dateWhere := fmt.Sprintf(
		"(Date >='%s' AND Date <= '%s' AND Time >= %d AND Time <= %d)",
		time.Unix(fromTimestamp, 0).Format("2006-01-02"),
		time.Unix(untilTimestamp, 0).Format("2006-01-02"),
		fromTimestamp,
		until,
	)

	// @TODO: change format to RowBinary
	query := fmt.Sprintf(
		`
		SELECT
			Path, Time, Value, Timestamp
		FROM %s
		WHERE (%s) AND (%s)
		FORMAT RowBinary
		`,
		h.config.ClickHouse.DataTable,
		pathWhere,
		dateWhere,
	)

	body, err := clickhouse.Query(
		r.Context(),
		h.config.ClickHouse.Url,
		query,
		h.config.ClickHouse.DataTimeout.Value(),
	)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	parseStart := time.Now()

	data, err := DataParse(body)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	logger.Debug("parse", zap.Duration("time_ns", time.Since(parseStart)))

	sortStart := time.Now()
	sort.Sort(data)
	logger.Debug("sort", zap.Duration("time_ns", time.Since(sortStart)))

	data.Points = point.Uniq(data.Points)

	// pp.Println(points)
	h.Reply(w, r, data.Points, int32(fromTimestamp), int32(untilTimestamp), prefix)
}

func (h *Handler) Reply(w http.ResponseWriter, r *http.Request, points []point.Point, from, until int32, prefix string) {
	start := time.Now()
	switch r.URL.Query().Get("format") {
	case "pickle":
		h.ReplyPickle(w, r, points, from, until, prefix)
	case "protobuf":
		h.ReplyProtobuf(w, r, points, from, until, prefix)
	}
	log.FromContext(r.Context()).Debug("reply", zap.Duration("time_ns", time.Since(start)))
}

func (h *Handler) ReplyPickle(w http.ResponseWriter, r *http.Request, points []point.Point, from, until int32, prefix string) {
	var rollupTime time.Duration
	var pickleTime time.Duration

	defer func() {
		log.FromContext(r.Context()).Debug("rollup", zap.Duration("time_ns", rollupTime))
		log.FromContext(r.Context()).Debug("pickle", zap.Duration("time_ns", pickleTime))
	}()

	if len(points) == 0 {
		w.Write(pickle.EmptyList)
		return
	}

	writer := bufio.NewWriterSize(w, 1024*1024)
	p := pickle.NewWriter(writer)
	defer writer.Flush()

	p.List()

	writeMetric := func(points []point.Point) {
		rollupStart := time.Now()
		points, step := h.config.Rollup.RollupMetric(points)
		rollupTime += time.Since(rollupStart)

		pickleStart := time.Now()
		p.Dict()

		p.String("name")
		if prefix != "" {
			p.String(prefix + "." + points[0].Metric)
		} else {
			p.String(points[0].Metric)
		}
		p.SetItem()

		p.String("step")
		p.Uint32(uint32(step))
		p.SetItem()

		start := from - (from % step)
		if start < from {
			start += step
		}
		end := until - (until % step)
		last := start - step

		p.String("values")
		p.List()
		for _, point := range points {
			if point.Time < start || point.Time > end {
				continue
			}

			if point.Time > last+step {
				p.AppendNulls(int(((point.Time - last) / step) - 1))
			}

			p.AppendFloat64(point.Value)

			last = point.Time
		}

		if end > last {
			p.AppendNulls(int((end - last) / step))
		}
		p.SetItem()

		p.String("start")
		p.Uint32(uint32(start))
		p.SetItem()

		p.String("end")
		p.Uint32(uint32(end))
		p.SetItem()

		p.Append()
		pickleTime += time.Since(pickleStart)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	p.Stop()
}

func (h *Handler) ReplyProtobuf(w http.ResponseWriter, r *http.Request, points []point.Point, from, until int32, prefix string) {
	if len(points) == 0 {
		return
	}

	var multiResponse carbonzipperpb.MultiFetchResponse

	writeMetric := func(points []point.Point) {
		points, step := h.config.Rollup.RollupMetric(points)

		var name string

		if prefix != "" {
			name = prefix + "." + points[0].Metric
		} else {
			name = points[0].Metric
		}

		start := from - (from % step)
		if start < from {
			start += step
		}
		stop := until - (until % step)
		count := ((stop - start) / step) + 1

		response := carbonzipperpb.FetchResponse{
			Name:      proto.String(name),
			StartTime: &start,
			StopTime:  &stop,
			StepTime:  &step,
			Values:    make([]float64, count),
			IsAbsent:  make([]bool, count),
		}

		var index int32
		// skip points before start
		for index = 0; index < len(points) && points[index].Time < start; index++ {
		}

		for i := int32(0); i < count; i++ {
			if index < int32(len(points)) && points[index].Time == start+step*i {
				response.Values[i] = points[index].Value
				response.IsAbsent[i] = false
				index++
			} else {
				response.Values[i] = 0
				response.IsAbsent[i] = true
			}
		}

		multiResponse.Metrics = append(multiResponse.Metrics, &response)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	body, _ := proto.Marshal(&multiResponse)
	w.Write(body)
}
<|endoftext|>


<commit_before>package pi

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"github.com/gocarina/formdata"
	"github.com/gorilla/schema"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
)

var (
	decoderFormValues = schema.NewDecoder()
	errNoFiles        = fmt.Errorf("no files")
)

// J is an helper to write JSON.
// Example:
//		c.WriteJSON(pi.J{"status": "OK"})
//
type J map[string]interface{}

// RequestContext represents the context of the HTTP request.
// It is shared across interceptors and handler.
type RequestContext struct {
	W        http.ResponseWriter
	R        *http.Request
	RouteURL string
	Data     map[interface{}]interface{}
}

// newRequestContext returns a new RequestContext.
func newRequestContext(w http.ResponseWriter, r *http.Request, routeURL string) *RequestContext {
	return &RequestContext{
		W:        w,
		R:        r,
		RouteURL: routeURL,
		Data:     make(map[interface{}]interface{}),
	}
}

// WriteString writes the specified strings to the ResponseWriter.
func (c *RequestContext) WriteString(strings ...string) error {
	for _, s := range strings {
		if _, err := c.W.Write([]byte(s)); err != nil {
			return err
		}
	}
	return nil
}

// WriteJSON marshal the object to JSON and writes it via the ResponseWriter.
func (c *RequestContext) WriteJSON(object interface{}) error {
	c.W.Header().Set("Content-Type", "application/json; charset=utf-8")
	if debugMode {
		output, err := json.MarshalIndent(object, "", "  ")
		if err != nil {
			return err
		}
		writeDebug("WriteJSON", c.R.RemoteAddr, string(output))
		c.W.Write(output)
	} else {
		output, err := json.Marshal(object)
		if err != nil {
			return err
		}
		c.W.Write(output)
	}
	return nil
}

// WriteXML marshal the object to XML and writes it via the ResponseWriter.
func (c *RequestContext) WriteXML(object interface{}) error {
	c.W.Header().Set("Content-Type", "application/xml; charset=utf-8")
	if debugMode {
		output, err := xml.MarshalIndent(object, "", "  ")
		if err != nil {
			return err
		}
		writeDebug("WriteXML", c.R.RemoteAddr, string(output))
		c.W.Write(output)
	} else {
		output, err := xml.Marshal(object)
		if err != nil {
			return err
		}
		c.W.Write(output)
	}
	return nil
}

// WriteReader copy the reader to the ResponseWriter.
func (c *RequestContext) WriteReader(reader io.Reader) error {
	_, err := io.Copy(c.W, reader)
	return err
}

// SetStatusCode sets the response status code.
func (c *RequestContext) SetStatusCode(statusCode int) {
	c.W.WriteHeader(statusCode)
}

// GetBody return the body as a ReadCloser. It is the client responsibility to close the body.
func (c *RequestContext) GetBody() io.ReadCloser {
	return c.R.Body
}

// GetRawBody returns the body as a byte array, closing the body reader.
func (c *RequestContext) GetRawBody() ([]byte, error) {
	body := c.GetBody()
	defer body.Close()
	rawBody, err := ioutil.ReadAll(body)
	if err != nil {
		return nil, err
	}
	if debugMode {
		writeDebug("GetRawBody", c.R.RemoteAddr, fmt.Sprintf("got %s", string(rawBody)))
	}
	return rawBody, nil
}

// GetJSONObject call json.Unmarshal by sending the reference of the given object.
// For example:
//		func GetUser(c *pi.RequestContext) error {
// 			user := &User{}
// 			if err := c.GetJSONObject(user); err != nil {
//				return pi.NewError(400, err)
//			}
//			// Do something with the user...
//			return nil
//		}
//
func (c *RequestContext) GetJSONObject(object interface{}) error {
	rawBody, err := c.GetRawBody()
	if err != nil {
		return err
	}
	return json.Unmarshal(rawBody, &object)
}

// GetXMLObject call xml.Unmarshal by sending the reference of the given object.
func (c *RequestContext) GetXMLObject(object interface{}) error {
	rawBody, err := c.GetRawBody()
	if err != nil {
		return err
	}
	return xml.Unmarshal(rawBody, &object)
}

// GetFormObject call gorilla/schema.Decode to maps the form values of the request into the object.
// For more informations: http://www.gorillatoolkit.org/pkg/schema
func (c *RequestContext) GetFormObject(object interface{}) error {
	if err := c.R.ParseForm(); err != nil {
		return err
	}
	return decoderFormValues.Decode(object, c.R.PostForm)
}

// GetMultipartObject calls gocarina/formdata.Unmarshal to maps the multipart form values of the request into the object.
// It supports files through multipart.FileHeader.
func (c *RequestContext) GetMultipartObject(object interface{}) error {
	return formdata.Unmarshal(c.R, object)
}

// GetRouteExtraPath returns the extra path.
// For example:
// 		for route("/files"), "/files/home/user/.emacs" will return "/home/user/.emacs"
func (c *RequestContext) GetRouteExtraPath() (path string) {
	fullPath := c.R.URL.String()
	if len(fullPath) > len(c.RouteURL) {
		path = fullPath[len(c.RouteURL):]
	}
	return path
}

// GetURLParam returns an URL parameter.
// For example, given this URL:
//		/user?id=1234
//
//		fmt.Println(c.GetURLParam("id"))
//		// Outputs 1234
//
func (c *RequestContext) GetURLParam(param string) string {
	return c.R.FormValue(param)
}

// GetURLParamOrDefault returns an URL parameter or the defaultValue if the value is empty.
func (c *RequestContext) GetURLParamOrDefault(param, defaultValue string) string {
	value := c.GetURLParam(param)
	if value == "" {
		return defaultValue
	}
	return value
}

// GetRouteVariable returns a route variable.
// For example:
//		getUserByID := func(c *RequestContext) error {
//			id := c.GetRouteVariable("id")
//			// Do something with the ID.
//			return nil
//		}
//
//		p := pi.New()
//		p.Route("/user/{id}").Get(getUserByID)
//		p.ListenAndServe(":8080")
//
func (c *RequestContext) GetRouteVariable(key string) string {
	return c.R.URL.Query().Get(":" + key)
}

// GetFileHeaders returns an array of FileHeader.
// For example:
//		getImages := func(c *RequestContext) error {
//			files, err := c.GetFileHeaders("files")
//			if err != nil {
//				return pi.NewError(400, err)
//			}
//			// Handle files
//			return nil
//		}
//
//		p := pi.New()
//		p.Route("/images").Post(getImages)
//		p.ListenAndServe(":8080")
//
func (c *RequestContext) GetFileHeaders(key string) ([]*multipart.FileHeader, error) {
	if err := c.R.ParseMultipartForm(32 << 20); err != nil {
		return nil, err
	}
	if c.R.MultipartForm != nil && c.R.MultipartForm.File[key] != nil {
		if debugMode {
			writeDebug("GetFileHeaders", c.R.RemoteAddr, fmt.Sprintf("got %d files from key %s", len(c.R.MultipartForm.File[key]), key))
		}
		return c.R.MultipartForm.File[key], nil
	}
	if debugMode {
		writeDebug("GetFileHeaders", c.R.RemoteAddr, fmt.Sprintf("got 0 files from key %s", key))
	}
	return nil, errNoFiles
}

// GetHeader returns the first value of the header of the request associated with the given key.
func (c *RequestContext) GetHeader(key string) string {
	return c.R.Header.Get(key)
}

// GetHeaderOrDefault returns the value of the header of the given defaultValue if the value is empty.
func (c *RequestContext) GetHeaderOrDefault(key, defaultValue string) string {
	value := c.GetHeader(key)
	if value == "" {
		return defaultValue
	}
	return key
}

// SetHeader sets to the response the header entries associated with key to
// the single element value.  It replaces any existing
// values associated with key.
func (c *RequestContext) SetHeader(key, value string) {
	c.W.Header().Set(key, value)
}

// AddHeader adds to the response the key, value pair to the header.
// It appends to any existing values associated with key.
func (c *RequestContext) AddHeader(key, value string) {
	c.W.Header().Add(key, value)
}

// DeleteHeader deletes from the response the values associated with key.
func (c *RequestContext) DeleteHeader(key string) {
	c.W.Header().Del(key)
}

// GetHeaders returns the array of values associated with the key from the request.
func (c *RequestContext) GetHeaders(key string) []string {
	return c.R.Header[key]
}
<commit_msg>Added GetDefault and WriteDefault<commit_after>package pi

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"github.com/gocarina/formdata"
	"github.com/gorilla/schema"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
)

var (
	decoderFormValues = schema.NewDecoder()

	// ErrNoFiles is the error when no files is present in the value provided when
	// calling RequestContext.GetFileHeaders.
	ErrNoFiles = fmt.Errorf("no files")

	// ErrContentTypeNotSupported is the error when the format of the Content-Type is not supported
	ErrContentTypeNotSupported = fmt.Errorf("format not supported")

	// ContentTypeJSON is the default MIME for JSON data.
	ContentTypeJSON = "application/json"

	// ContentTypeXML is the default MIME for XML data.
	ContentTypeXML = "application/xml"

	// ContentTypeClassicForm is the default MIME for Form Encoded data.
	ContentTypeClassicForm = "application/x-www-form-urlencoded"

	// ContentTypeMultipart is the default MIME for Webform.
	ContentTypeMultipart = "multipart/form-data"

	// ContentTypeText is the default MIME for Textual data.
	ContentTypeText = "text/plain"
)

// J is an helper to write JSON.
// Example:
//		c.WriteJSON(pi.J{"status": "OK"})
//
type J map[string]interface{}

// RequestContext represents the context of the HTTP request.
// It is shared across interceptors and handler.
type RequestContext struct {
	W        http.ResponseWriter
	R        *http.Request
	RouteURL string
	Data     map[interface{}]interface{}
}

// newRequestContext returns a new RequestContext.
func newRequestContext(w http.ResponseWriter, r *http.Request, routeURL string) *RequestContext {
	return &RequestContext{
		W:        w,
		R:        r,
		RouteURL: routeURL,
		Data:     make(map[interface{}]interface{}),
	}
}

// WriteString writes the specified strings to the ResponseWriter.
func (c *RequestContext) WriteString(strings ...string) error {
	for _, s := range strings {
		if _, err := c.W.Write([]byte(s)); err != nil {
			return err
		}
	}
	return nil
}

// WriteJSON marshal the object to JSON and writes it via the ResponseWriter.
func (c *RequestContext) WriteJSON(object interface{}) error {
	c.W.Header().Set("Content-Type", "application/json; charset=utf-8")
	if debugMode {
		output, err := json.MarshalIndent(object, "", "  ")
		if err != nil {
			return err
		}
		writeDebug("WriteJSON", c.R.RemoteAddr, string(output))
		c.W.Write(output)
	} else {
		output, err := json.Marshal(object)
		if err != nil {
			return err
		}
		c.W.Write(output)
	}
	return nil
}

// WriteXML marshal the object to XML and writes it via the ResponseWriter.
func (c *RequestContext) WriteXML(object interface{}) error {
	c.W.Header().Set("Content-Type", "application/xml; charset=utf-8")
	if debugMode {
		output, err := xml.MarshalIndent(object, "", "  ")
		if err != nil {
			return err
		}
		writeDebug("WriteXML", c.R.RemoteAddr, string(output))
		c.W.Write(output)
	} else {
		output, err := xml.Marshal(object)
		if err != nil {
			return err
		}
		c.W.Write(output)
	}
	return nil
}

// WriteDefault writes the object to the caller according to the acceptable MIME in the Accept header value.
// If the MIME is not supported, it sends a 406 Not Acceptable request.
// text/plain uses the String method to be serialized.
// Mime supported for write:
//
//		application/json
//		application/xml
//		text/plain
//
// If no Accept header is present, it writes the object as JSON.
func (c *RequestContext) WriteDefault(object interface{}) error {
	acceptedContentType := c.GetAccepts()
	for _, contentType := range acceptedContentType {
		switch contentType {
		case ContentTypeText:
			return c.WriteString(fmt.Sprintf("%v", object))
		case ContentTypeJSON, "*/*":
			return c.WriteJSON(object)
		case ContentTypeXML:
			return c.WriteXML(object)
		}
	}
	return NewError(406, ErrContentTypeNotSupported)
}

// WriteReader copy the reader to the ResponseWriter.
func (c *RequestContext) WriteReader(reader io.Reader) error {
	_, err := io.Copy(c.W, reader)
	return err
}

// SetStatusCode sets the response status code.
func (c *RequestContext) SetStatusCode(statusCode int) {
	c.W.WriteHeader(statusCode)
}

// GetBody return the body as a ReadCloser. It is the client responsibility to close the body.
func (c *RequestContext) GetBody() io.ReadCloser {
	return c.R.Body
}

// GetRawBody returns the body as a byte array, closing the body reader.
func (c *RequestContext) GetRawBody() ([]byte, error) {
	body := c.GetBody()
	defer body.Close()
	rawBody, err := ioutil.ReadAll(body)
	if err != nil {
		return nil, err
	}
	if debugMode {
		writeDebug("GetRawBody", c.R.RemoteAddr, fmt.Sprintf("got %s", string(rawBody)))
	}
	return rawBody, nil
}

// GetJSONObject call json.Unmarshal by sending the reference of the given object.
// For example:
//		func GetUser(c *pi.RequestContext) error {
// 			user := &User{}
// 			if err := c.GetJSONObject(user); err != nil {
//				return pi.NewError(400, err)
//			}
//			// Do something with the user...
//			return nil
//		}
//
func (c *RequestContext) GetJSONObject(object interface{}) error {
	rawBody, err := c.GetRawBody()
	if err != nil {
		return err
	}
	return json.Unmarshal(rawBody, &object)
}

// GetXMLObject call xml.Unmarshal by sending the reference of the given object.
func (c *RequestContext) GetXMLObject(object interface{}) error {
	rawBody, err := c.GetRawBody()
	if err != nil {
		return err
	}
	return xml.Unmarshal(rawBody, &object)
}

// GetFormObject call gorilla/schema.Decode to maps the form values of the request into the object.
// For more informations: http://www.gorillatoolkit.org/pkg/schema
func (c *RequestContext) GetFormObject(object interface{}) error {
	if err := c.R.ParseForm(); err != nil {
		return err
	}
	return decoderFormValues.Decode(object, c.R.PostForm)
}

// GetMultipartObject calls gocarina/formdata.Unmarshal to maps the multipart form values of the request into the object.
// It supports files through multipart.FileHeader.
func (c *RequestContext) GetMultipartObject(object interface{}) error {
	return formdata.Unmarshal(c.R, object)
}

// GetDefaultObject calls one of a GetX method according to the Content-Type of the request.
// Content-Types supported:
// 		application/json
//		application/xml
//		application/x-www-form-urlencoded
//		multipart/form-data
//
func (c *RequestContext) GetDefaultObject(object interface{}) error {
	switch c.GetContentType() {
	case ContentTypeJSON:
		return c.GetJSONObject(object)
	case ContentTypeXML:
		return c.GetXMLObject(object)
	case ContentTypeClassicForm:
		return c.GetFormObject(object)
	case ContentTypeMultipart:
		return c.GetMultipartObject(object)
	}
	return ErrContentTypeNotSupported
}

// GetRouteExtraPath returns the extra path.
// For example:
// 		for route("/files"), "/files/home/user/.emacs" will return "/home/user/.emacs"
func (c *RequestContext) GetRouteExtraPath() (path string) {
	fullPath := c.R.URL.String()
	if len(fullPath) > len(c.RouteURL) {
		path = fullPath[len(c.RouteURL):]
	}
	return path
}

// GetURLParam returns an URL parameter.
// For example, given this URL:
//		/user?id=1234
//
//		fmt.Println(c.GetURLParam("id"))
//		// Outputs 1234
//
func (c *RequestContext) GetURLParam(param string) string {
	return c.R.FormValue(param)
}

// GetURLParamOrDefault returns an URL parameter or the defaultValue if the value is empty.
func (c *RequestContext) GetURLParamOrDefault(param, defaultValue string) string {
	value := c.GetURLParam(param)
	if value == "" {
		return defaultValue
	}
	return value
}

// GetRouteVariable returns a route variable.
// For example:
//		getUserByID := func(c *RequestContext) error {
//			id := c.GetRouteVariable("id")
//			// Do something with the ID.
//			return nil
//		}
//
//		p := pi.New()
//		p.Route("/user/{id}").Get(getUserByID)
//		p.ListenAndServe(":8080")
//
func (c *RequestContext) GetRouteVariable(key string) string {
	return c.R.URL.Query().Get(":" + key)
}

// GetFileHeaders returns an array of FileHeader.
// For example:
//		getImages := func(c *RequestContext) error {
//			files, err := c.GetFileHeaders("files")
//			if err != nil {
//				return pi.NewError(400, err)
//			}
//			// Handle files
//			return nil
//		}
//
//		p := pi.New()
//		p.Route("/images").Post(getImages)
//		p.ListenAndServe(":8080")
//
func (c *RequestContext) GetFileHeaders(key string) ([]*multipart.FileHeader, error) {
	if err := c.R.ParseMultipartForm(32 << 20); err != nil {
		return nil, err
	}
	if c.R.MultipartForm != nil && c.R.MultipartForm.File[key] != nil {
		if debugMode {
			writeDebug("GetFileHeaders", c.R.RemoteAddr, fmt.Sprintf("got %d files from key %s", len(c.R.MultipartForm.File[key]), key))
		}
		return c.R.MultipartForm.File[key], nil
	}
	if debugMode {
		writeDebug("GetFileHeaders", c.R.RemoteAddr, fmt.Sprintf("got 0 files from key %s", key))
	}
	return nil, ErrNoFiles
}

// GetHeader returns the first value of the header of the request associated with the given key.
func (c *RequestContext) GetHeader(key string) string {
	return c.R.Header.Get(key)
}

// GetContentType returns the Content-Type of the request.
func (c *RequestContext) GetContentType() string {
	return c.GetHeader("Content-Type")
}

// GetAccepts returns the MIME types acceptable by the caller.
func (c *RequestContext) GetAccepts() []string {
	return c.GetHeaders("Accept")
}

// GetHeaderOrDefault returns the value of the header of the given defaultValue if the value is empty.
func (c *RequestContext) GetHeaderOrDefault(key, defaultValue string) string {
	value := c.GetHeader(key)
	if value == "" {
		return defaultValue
	}
	return key
}

// SetHeader sets to the response the header entries associated with key to
// the single element value.  It replaces any existing
// values associated with key.
func (c *RequestContext) SetHeader(key, value string) {
	c.W.Header().Set(key, value)
}

// AddHeader adds to the response the key, value pair to the header.
// It appends to any existing values associated with key.
func (c *RequestContext) AddHeader(key, value string) {
	c.W.Header().Add(key, value)
}

// DeleteHeader deletes from the response the values associated with key.
func (c *RequestContext) DeleteHeader(key string) {
	c.W.Header().Del(key)
}

// GetHeaders returns the array of values associated with the key from the request.
func (c *RequestContext) GetHeaders(key string) []string {
	return c.R.Header[key]
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"io"
	"net"
	"reflect"
	"testing"
	"time"

	"github.com/davecgh/go-spew/spew"
)

// TestAddr tests the MsgAddr API.
func TestAddr(t *testing.T) {
	pver := ProtocolVersion

	// Ensure the command is expected value.
	wantCmd := "addr"
	msg := NewMsgAddr()
	if cmd := msg.Command(); cmd != wantCmd {
		t.Errorf("NewMsgAddr: wrong command - got %v want %v",
			cmd, wantCmd)
	}

	// Ensure max payload is expected value for latest protocol version.
	// Num addresses (varInt) + max allowed addresses.
	wantPayload := uint32(30009)
	maxPayload := msg.MaxPayloadLength(pver)
	if maxPayload != wantPayload {
		t.Errorf("MaxPayloadLength: wrong max payload length for "+
			"protocol version %d - got %v, want %v", pver,
			maxPayload, wantPayload)
	}

	// Ensure NetAddresses are added properly.
	tcpAddr := &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 9333}
	na := NewNetAddress(tcpAddr, SFNodeNetwork)
	err := msg.AddAddress(na)
	if err != nil {
		t.Errorf("AddAddress: %v", err)
	}
	if msg.AddrList[0] != na {
		t.Errorf("AddAddress: wrong address added - got %v, want %v",
			spew.Sprint(msg.AddrList[0]), spew.Sprint(na))
	}

	// Ensure the address list is cleared properly.
	msg.ClearAddresses()
	if len(msg.AddrList) != 0 {
		t.Errorf("ClearAddresses: address list is not empty - "+
			"got %v [%v], want %v", len(msg.AddrList),
			spew.Sprint(msg.AddrList[0]), 0)
	}

	// Ensure adding more than the max allowed addresses per message returns
	// error.
	for i := 0; i < MaxAddrPerMsg+1; i++ {
		err = msg.AddAddress(na)
	}
	if err == nil {
		t.Errorf("AddAddress: expected error on too many addresses " +
			"not received")
	}
	err = msg.AddAddresses(na)
	if err == nil {
		t.Errorf("AddAddresses: expected error on too many addresses " +
			"not received")
	}

	// Ensure max payload is expected value for protocol versions before
	// timestamp was added to NetAddress.
	// Num addresses (varInt) + max allowed addresses.
	pver = NetAddressTimeVersion - 1
	wantPayload = uint32(26009)
	maxPayload = msg.MaxPayloadLength(pver)
	if maxPayload != wantPayload {
		t.Errorf("MaxPayloadLength: wrong max payload length for "+
			"protocol version %d - got %v, want %v", pver,
			maxPayload, wantPayload)
	}

	// Ensure max payload is expected value for protocol versions before
	// multiple addresses were allowed.
	// Num addresses (varInt) + a single net addresses.
	pver = MultipleAddressVersion - 1
	wantPayload = uint32(35)
	maxPayload = msg.MaxPayloadLength(pver)
	if maxPayload != wantPayload {
		t.Errorf("MaxPayloadLength: wrong max payload length for "+
			"protocol version %d - got %v, want %v", pver,
			maxPayload, wantPayload)
	}
}

// TestAddrWire tests the MsgAddr wire encode and decode for various numbers
// of addresses and protocol versions.
func TestAddrWire(t *testing.T) {
	// A couple of NetAddresses to use for testing.
	na := &NetAddress{
		Timestamp: time.Unix(0x495fab29, 0), // 2009-01-03 12:15:05 -0600 CST
		Services:  SFNodeNetwork,
		IP:        net.ParseIP("127.0.0.1"),
		Port:      9333,
	}
	na2 := &NetAddress{
		Timestamp: time.Unix(0x495fab29, 0), // 2009-01-03 12:15:05 -0600 CST
		Services:  SFNodeNetwork,
		IP:        net.ParseIP("192.168.0.1"),
		Port:      9334,
	}

	// Empty address message.
	noAddr := NewMsgAddr()
	noAddrEncoded := []byte{
		0x00, // Varint for number of addresses
	}

	// Address message with multiple addresses.
	multiAddr := NewMsgAddr()
	multiAddr.AddAddresses(na, na2)
	multiAddrEncoded := []byte{
		0x02,                   // Varint for number of addresses
		0x29, 0xab, 0x5f, 0x49, // Timestamp
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // SFNodeNetwork
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0xff, 0xff, 0x7f, 0x00, 0x00, 0x01, // IP 127.0.0.1
		0x20, 0x8d, // Port 9333 in big-endian
		0x29, 0xab, 0x5f, 0x49, // Timestamp
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // SFNodeNetwork
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0xff, 0xff, 0xc0, 0xa8, 0x00, 0x01, // IP 192.168.0.1
		0x20, 0x8e, // Port 9334 in big-endian

	}

	tests := []struct {
		in   *MsgAddr        // Message to encode
		out  *MsgAddr        // Expected decoded message
		buf  []byte          // Wire encoding
		pver uint32          // Protocol version for wire encoding
		enc  MessageEncoding // Message encoding format
	}{
		// Latest protocol version with no addresses.
		{
			noAddr,
			noAddr,
			noAddrEncoded,
			ProtocolVersion,
			BaseEncoding,
		},

		// Latest protocol version with multiple addresses.
		{
			multiAddr,
			multiAddr,
			multiAddrEncoded,
			ProtocolVersion,
			BaseEncoding,
		},

		// Protocol version MultipleAddressVersion-1 with no addresses.
		{
			noAddr,
			noAddr,
			noAddrEncoded,
			MultipleAddressVersion - 1,
			BaseEncoding,
		},
	}

	t.Logf("Running %d tests", len(tests))
	for i, test := range tests {
		// Encode the message to wire format.
		var buf bytes.Buffer
		err := test.in.BtcEncode(&buf, test.pver, test.enc)
		if err != nil {
			t.Errorf("BtcEncode #%d error %v", i, err)
			continue
		}
		if !bytes.Equal(buf.Bytes(), test.buf) {
			t.Errorf("BtcEncode #%d\n got: %s want: %s", i,
				spew.Sdump(buf.Bytes()), spew.Sdump(test.buf))
			continue
		}

		// Decode the message from wire format.
		var msg MsgAddr
		rbuf := bytes.NewReader(test.buf)
		err = msg.BtcDecode(rbuf, test.pver, test.enc)
		if err != nil {
			t.Errorf("BtcDecode #%d error %v", i, err)
			continue
		}
		if !reflect.DeepEqual(&msg, test.out) {
			t.Errorf("BtcDecode #%d\n got: %s want: %s", i,
				spew.Sdump(msg), spew.Sdump(test.out))
			continue
		}
	}
}

// TestAddrWireErrors performs negative tests against wire encode and decode
// of MsgAddr to confirm error paths work correctly.
func TestAddrWireErrors(t *testing.T) {
	pver := ProtocolVersion
	pverMA := MultipleAddressVersion
	wireErr := &MessageError{}

	// A couple of NetAddresses to use for testing.
	na := &NetAddress{
		Timestamp: time.Unix(0x495fab29, 0), // 2009-01-03 12:15:05 -0600 CST
		Services:  SFNodeNetwork,
		IP:        net.ParseIP("127.0.0.1"),
		Port:      9333,
	}
	na2 := &NetAddress{
		Timestamp: time.Unix(0x495fab29, 0), // 2009-01-03 12:15:05 -0600 CST
		Services:  SFNodeNetwork,
		IP:        net.ParseIP("192.168.0.1"),
		Port:      9334,
	}

	// Address message with multiple addresses.
	baseAddr := NewMsgAddr()
	baseAddr.AddAddresses(na, na2)
	baseAddrEncoded := []byte{
		0x02,                   // Varint for number of addresses
		0x29, 0xab, 0x5f, 0x49, // Timestamp
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // SFNodeNetwork
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0xff, 0xff, 0x7f, 0x00, 0x00, 0x01, // IP 127.0.0.1
		0x20, 0x8d, // Port 9333 in big-endian
		0x29, 0xab, 0x5f, 0x49, // Timestamp
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // SFNodeNetwork
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0xff, 0xff, 0xc0, 0xa8, 0x00, 0x01, // IP 192.168.0.1
		0x20, 0x8e, // Port 9334 in big-endian

	}

	// Message that forces an error by having more than the max allowed
	// addresses.
	maxAddr := NewMsgAddr()
	for i := 0; i < MaxAddrPerMsg; i++ {
		maxAddr.AddAddress(na)
	}
	maxAddr.AddrList = append(maxAddr.AddrList, na)
	maxAddrEncoded := []byte{
		0xfd, 0x03, 0xe9, // Varint for number of addresses (1001)
	}

	tests := []struct {
		in       *MsgAddr        // Value to encode
		buf      []byte          // Wire encoding
		pver     uint32          // Protocol version for wire encoding
		enc      MessageEncoding // Message encoding format
		max      int             // Max size of fixed buffer to induce errors
		writeErr error           // Expected write error
		readErr  error           // Expected read error
	}{
		// Latest protocol version with intentional read/write errors.
		// Force error in addresses count
		{baseAddr, baseAddrEncoded, pver, BaseEncoding, 0, io.ErrShortWrite, io.EOF},
		// Force error in address list.
		{baseAddr, baseAddrEncoded, pver, BaseEncoding, 1, io.ErrShortWrite, io.EOF},
		// Force error with greater than max inventory vectors.
		{maxAddr, maxAddrEncoded, pver, BaseEncoding, 3, wireErr, wireErr},
		// Force error with greater than max inventory vectors for
		// protocol versions before multiple addresses were allowed.
		{maxAddr, maxAddrEncoded, pverMA - 1, BaseEncoding, 3, wireErr, wireErr},
	}

	t.Logf("Running %d tests", len(tests))
	for i, test := range tests {
		// Encode to wire format.
		w := newFixedWriter(test.max)
		err := test.in.BtcEncode(w, test.pver, test.enc)
		if reflect.TypeOf(err) != reflect.TypeOf(test.writeErr) {
			t.Errorf("BtcEncode #%d wrong error got: %v, want: %v",
				i, err, test.writeErr)
			continue
		}

		// For errors which are not of type MessageError, check them for
		// equality.
		if _, ok := err.(*MessageError); !ok {
			if err != test.writeErr {
				t.Errorf("BtcEncode #%d wrong error got: %v, "+
					"want: %v", i, err, test.writeErr)
				continue
			}
		}

		// Decode from wire format.
		var msg MsgAddr
		r := newFixedReader(test.max, test.buf)
		err = msg.BtcDecode(r, test.pver, test.enc)
		if reflect.TypeOf(err) != reflect.TypeOf(test.readErr) {
			t.Errorf("BtcDecode #%d wrong error got: %v, want: %v",
				i, err, test.readErr)
			continue
		}

		// For errors which are not of type MessageError, check them for
		// equality.
		if _, ok := err.(*MessageError); !ok {
			if err != test.readErr {
				t.Errorf("BtcDecode #%d wrong error got: %v, "+
					"want: %v", i, err, test.readErr)
				continue
			}
		}

	}
}
<commit_msg>test: fix msgaddr<commit_after>// Copyright (c) 2013-2016 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"bytes"
	"io"
	"net"
	"reflect"
	"testing"
	"time"

	"github.com/davecgh/go-spew/spew"
)

// TestAddr tests the MsgAddr API.
func TestAddr(t *testing.T) {
	pver := ProtocolVersion

	// Ensure the command is expected value.
	wantCmd := "addr"
	msg := NewMsgAddr()
	if cmd := msg.Command(); cmd != wantCmd {
		t.Errorf("NewMsgAddr: wrong command - got %v want %v",
			cmd, wantCmd)
	}

	// Ensure max payload is expected value for latest protocol version.
	// Num addresses (varInt) + max allowed addresses.
	wantPayload := uint32(30009)
	maxPayload := msg.MaxPayloadLength(pver)
	if maxPayload != wantPayload {
		t.Errorf("MaxPayloadLength: wrong max payload length for "+
			"protocol version %d - got %v, want %v", pver,
			maxPayload, wantPayload)
	}

	// Ensure NetAddresses are added properly.
	tcpAddr := &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 9333}
	na := NewNetAddress(tcpAddr, SFNodeNetwork)
	err := msg.AddAddress(na)
	if err != nil {
		t.Errorf("AddAddress: %v", err)
	}
	if msg.AddrList[0] != na {
		t.Errorf("AddAddress: wrong address added - got %v, want %v",
			spew.Sprint(msg.AddrList[0]), spew.Sprint(na))
	}

	// Ensure the address list is cleared properly.
	msg.ClearAddresses()
	if len(msg.AddrList) != 0 {
		t.Errorf("ClearAddresses: address list is not empty - "+
			"got %v [%v], want %v", len(msg.AddrList),
			spew.Sprint(msg.AddrList[0]), 0)
	}

	// Ensure adding more than the max allowed addresses per message returns
	// error.
	for i := 0; i < MaxAddrPerMsg+1; i++ {
		err = msg.AddAddress(na)
	}
	if err == nil {
		t.Errorf("AddAddress: expected error on too many addresses " +
			"not received")
	}
	err = msg.AddAddresses(na)
	if err == nil {
		t.Errorf("AddAddresses: expected error on too many addresses " +
			"not received")
	}

	// Ensure max payload is expected value for protocol versions before
	// timestamp was added to NetAddress.
	// Num addresses (varInt) + max allowed addresses.
	pver = NetAddressTimeVersion - 1
	wantPayload = uint32(26009)
	maxPayload = msg.MaxPayloadLength(pver)
	if maxPayload != wantPayload {
		t.Errorf("MaxPayloadLength: wrong max payload length for "+
			"protocol version %d - got %v, want %v", pver,
			maxPayload, wantPayload)
	}

	// Ensure max payload is expected value for protocol versions before
	// multiple addresses were allowed.
	// Num addresses (varInt) + a single net addresses.
	pver = MultipleAddressVersion - 1
	wantPayload = uint32(35)
	maxPayload = msg.MaxPayloadLength(pver)
	if maxPayload != wantPayload {
		t.Errorf("MaxPayloadLength: wrong max payload length for "+
			"protocol version %d - got %v, want %v", pver,
			maxPayload, wantPayload)
	}
}

// TestAddrWire tests the MsgAddr wire encode and decode for various numbers
// of addresses and protocol versions.
func TestAddrWire(t *testing.T) {
	// A couple of NetAddresses to use for testing.
	na := &NetAddress{
		Timestamp: time.Unix(0x495fab29, 0), // 2009-01-03 12:15:05 -0600 CST
		Services:  SFNodeNetwork,
		IP:        net.ParseIP("127.0.0.1"),
		Port:      9333,
	}
	na2 := &NetAddress{
		Timestamp: time.Unix(0x495fab29, 0), // 2009-01-03 12:15:05 -0600 CST
		Services:  SFNodeNetwork,
		IP:        net.ParseIP("192.168.0.1"),
		Port:      9334,
	}

	// Empty address message.
	noAddr := NewMsgAddr()
	noAddrEncoded := []byte{
		0x00, // Varint for number of addresses
	}

	// Address message with multiple addresses.
	multiAddr := NewMsgAddr()
	multiAddr.AddAddresses(na, na2)
	multiAddrEncoded := []byte{
		0x02,                   // Varint for number of addresses
		0x29, 0xab, 0x5f, 0x49, // Timestamp
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // SFNodeNetwork
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0xff, 0xff, 0x7f, 0x00, 0x00, 0x01, // IP 127.0.0.1
		0x24, 0x75, // Port 9333 in big-endian
		0x29, 0xab, 0x5f, 0x49, // Timestamp
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // SFNodeNetwork
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0xff, 0xff, 0xc0, 0xa8, 0x00, 0x01, // IP 192.168.0.1
		0x24, 0x76, // Port 9334 in big-endian

	}

	tests := []struct {
		in   *MsgAddr        // Message to encode
		out  *MsgAddr        // Expected decoded message
		buf  []byte          // Wire encoding
		pver uint32          // Protocol version for wire encoding
		enc  MessageEncoding // Message encoding format
	}{
		// Latest protocol version with no addresses.
		{
			noAddr,
			noAddr,
			noAddrEncoded,
			ProtocolVersion,
			BaseEncoding,
		},

		// Latest protocol version with multiple addresses.
		{
			multiAddr,
			multiAddr,
			multiAddrEncoded,
			ProtocolVersion,
			BaseEncoding,
		},

		// Protocol version MultipleAddressVersion-1 with no addresses.
		{
			noAddr,
			noAddr,
			noAddrEncoded,
			MultipleAddressVersion - 1,
			BaseEncoding,
		},
	}

	t.Logf("Running %d tests", len(tests))
	for i, test := range tests {
		// Encode the message to wire format.
		var buf bytes.Buffer
		err := test.in.BtcEncode(&buf, test.pver, test.enc)
		if err != nil {
			t.Errorf("BtcEncode #%d error %v", i, err)
			continue
		}
		if !bytes.Equal(buf.Bytes(), test.buf) {
			t.Errorf("BtcEncode #%d\n got: %s want: %s", i,
				spew.Sdump(buf.Bytes()), spew.Sdump(test.buf))
			continue
		}

		// Decode the message from wire format.
		var msg MsgAddr
		rbuf := bytes.NewReader(test.buf)
		err = msg.BtcDecode(rbuf, test.pver, test.enc)
		if err != nil {
			t.Errorf("BtcDecode #%d error %v", i, err)
			continue
		}
		if !reflect.DeepEqual(&msg, test.out) {
			t.Errorf("BtcDecode #%d\n got: %s want: %s", i,
				spew.Sdump(msg), spew.Sdump(test.out))
			continue
		}
	}
}

// TestAddrWireErrors performs negative tests against wire encode and decode
// of MsgAddr to confirm error paths work correctly.
func TestAddrWireErrors(t *testing.T) {
	pver := ProtocolVersion
	pverMA := MultipleAddressVersion
	wireErr := &MessageError{}

	// A couple of NetAddresses to use for testing.
	na := &NetAddress{
		Timestamp: time.Unix(0x495fab29, 0), // 2009-01-03 12:15:05 -0600 CST
		Services:  SFNodeNetwork,
		IP:        net.ParseIP("127.0.0.1"),
		Port:      9333,
	}
	na2 := &NetAddress{
		Timestamp: time.Unix(0x495fab29, 0), // 2009-01-03 12:15:05 -0600 CST
		Services:  SFNodeNetwork,
		IP:        net.ParseIP("192.168.0.1"),
		Port:      9334,
	}

	// Address message with multiple addresses.
	baseAddr := NewMsgAddr()
	baseAddr.AddAddresses(na, na2)
	baseAddrEncoded := []byte{
		0x02,                   // Varint for number of addresses
		0x29, 0xab, 0x5f, 0x49, // Timestamp
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // SFNodeNetwork
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0xff, 0xff, 0x7f, 0x00, 0x00, 0x01, // IP 127.0.0.1
		0x20, 0x8d, // Port 9333 in big-endian
		0x29, 0xab, 0x5f, 0x49, // Timestamp
		0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // SFNodeNetwork
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0xff, 0xff, 0xc0, 0xa8, 0x00, 0x01, // IP 192.168.0.1
		0x20, 0x8e, // Port 9334 in big-endian

	}

	// Message that forces an error by having more than the max allowed
	// addresses.
	maxAddr := NewMsgAddr()
	for i := 0; i < MaxAddrPerMsg; i++ {
		maxAddr.AddAddress(na)
	}
	maxAddr.AddrList = append(maxAddr.AddrList, na)
	maxAddrEncoded := []byte{
		0xfd, 0x03, 0xe9, // Varint for number of addresses (1001)
	}

	tests := []struct {
		in       *MsgAddr        // Value to encode
		buf      []byte          // Wire encoding
		pver     uint32          // Protocol version for wire encoding
		enc      MessageEncoding // Message encoding format
		max      int             // Max size of fixed buffer to induce errors
		writeErr error           // Expected write error
		readErr  error           // Expected read error
	}{
		// Latest protocol version with intentional read/write errors.
		// Force error in addresses count
		{baseAddr, baseAddrEncoded, pver, BaseEncoding, 0, io.ErrShortWrite, io.EOF},
		// Force error in address list.
		{baseAddr, baseAddrEncoded, pver, BaseEncoding, 1, io.ErrShortWrite, io.EOF},
		// Force error with greater than max inventory vectors.
		{maxAddr, maxAddrEncoded, pver, BaseEncoding, 3, wireErr, wireErr},
		// Force error with greater than max inventory vectors for
		// protocol versions before multiple addresses were allowed.
		{maxAddr, maxAddrEncoded, pverMA - 1, BaseEncoding, 3, wireErr, wireErr},
	}

	t.Logf("Running %d tests", len(tests))
	for i, test := range tests {
		// Encode to wire format.
		w := newFixedWriter(test.max)
		err := test.in.BtcEncode(w, test.pver, test.enc)
		if reflect.TypeOf(err) != reflect.TypeOf(test.writeErr) {
			t.Errorf("BtcEncode #%d wrong error got: %v, want: %v",
				i, err, test.writeErr)
			continue
		}

		// For errors which are not of type MessageError, check them for
		// equality.
		if _, ok := err.(*MessageError); !ok {
			if err != test.writeErr {
				t.Errorf("BtcEncode #%d wrong error got: %v, "+
					"want: %v", i, err, test.writeErr)
				continue
			}
		}

		// Decode from wire format.
		var msg MsgAddr
		r := newFixedReader(test.max, test.buf)
		err = msg.BtcDecode(r, test.pver, test.enc)
		if reflect.TypeOf(err) != reflect.TypeOf(test.readErr) {
			t.Errorf("BtcDecode #%d wrong error got: %v, want: %v",
				i, err, test.readErr)
			continue
		}

		// For errors which are not of type MessageError, check them for
		// equality.
		if _, ok := err.(*MessageError); !ok {
			if err != test.readErr {
				t.Errorf("BtcDecode #%d wrong error got: %v, "+
					"want: %v", i, err, test.readErr)
				continue
			}
		}

	}
}
<|endoftext|>


<commit_before>package cluster

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/db/query"
	"github.com/lxc/lxd/lxd/task"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	"github.com/pkg/errors"
)

// APIHeartbeatMember contains specific cluster node info.
type APIHeartbeatMember struct {
	ID            int64     // ID field value in nodes table.
	Address       string    // Host and Port of node.
	RaftID        uint64    // ID field value in raft_nodes table, zero if non-raft node.
	RaftRole      int       // Node role in the raft cluster, from the raft_nodes table
	Raft          bool      // Deprecated, use non-zero RaftID instead to indicate raft node.
	LastHeartbeat time.Time // Last time we received a successful response from node.
	Online        bool      // Calculated from offline threshold and LastHeatbeat time.
	updated       bool      // Has node been updated during this heartbeat run. Not sent to nodes.
}

// APIHeartbeatVersion contains max versions for all nodes in cluster.
type APIHeartbeatVersion struct {
	Schema        int
	APIExtensions int
}

// APIHeartbeat contains data sent to nodes in heartbeat.
type APIHeartbeat struct {
	sync.Mutex // Used to control access to Members maps.
	cluster    *db.Cluster
	Members    map[int64]APIHeartbeatMember
	Version    APIHeartbeatVersion
	Time       time.Time

	// Indicates if heartbeat contains a fresh set of node states.
	// This can be used to indicate to the receiving node that the state is fresh enough to
	// trigger node refresh activies (such as forkdns).
	FullStateList bool
}

// Update updates an existing APIHeartbeat struct with the raft and all node states supplied.
// If allNodes provided is an empty set then this is considered a non-full state list.
func (hbState *APIHeartbeat) Update(fullStateList bool, raftNodes []db.RaftNode, allNodes []db.NodeInfo, offlineThreshold time.Duration) {
	var maxSchemaVersion, maxAPIExtensionsVersion int

	if hbState.Members == nil {
		hbState.Members = make(map[int64]APIHeartbeatMember)
	}

	// If we've been supplied a fresh set of node states, this is a full state list.
	hbState.FullStateList = fullStateList

	raftNodeMap := make(map[string]db.RaftNode)

	// Convert raftNodes to a map keyed on address for lookups later.
	for _, raftNode := range raftNodes {
		raftNodeMap[raftNode.Address] = raftNode
	}

	// Add nodes (overwrites any nodes with same ID in map with fresh data).
	for _, node := range allNodes {
		member := APIHeartbeatMember{
			ID:            node.ID,
			Address:       node.Address,
			LastHeartbeat: node.Heartbeat,
			Online:        !node.Heartbeat.Before(time.Now().Add(-offlineThreshold)),
		}

		if raftNode, exists := raftNodeMap[member.Address]; exists {
			member.Raft = true // Deprecated
			member.RaftID = raftNode.ID
			member.RaftRole = int(raftNode.Role)
			delete(raftNodeMap, member.Address) // Used to check any remaining later.
		}

		// Add to the members map using the node ID (not the Raft Node ID).
		hbState.Members[node.ID] = member

		// Keep a record of highest APIExtensions and Schema version seen in all nodes.
		if node.APIExtensions > maxAPIExtensionsVersion {
			maxAPIExtensionsVersion = node.APIExtensions
		}

		if node.Schema > maxSchemaVersion {
			maxSchemaVersion = node.Schema
		}
	}

	hbState.Version = APIHeartbeatVersion{
		Schema:        maxSchemaVersion,
		APIExtensions: maxAPIExtensionsVersion,
	}

	if len(raftNodeMap) > 0 && hbState.cluster != nil {
		hbState.cluster.Transaction(func(tx *db.ClusterTx) error {
			for addr, raftNode := range raftNodeMap {
				_, err := tx.GetPendingNodeByAddress(addr)
				if err != nil {
					logger.Errorf("Unaccounted raft node(s) not found in 'nodes' table for heartbeat: %+v", raftNode)
				}
			}

			return nil
		})
	}

	return
}

// Send sends heartbeat requests to the nodes supplied and updates heartbeat state.
func (hbState *APIHeartbeat) Send(ctx context.Context, networkCert *shared.CertInfo, serverCert *shared.CertInfo, localAddress string, nodes []db.NodeInfo, spreadDuration time.Duration) {
	heartbeatsWg := sync.WaitGroup{}
	sendHeartbeat := func(nodeID int64, address string, spreadDuration time.Duration, heartbeatData *APIHeartbeat) {
		defer heartbeatsWg.Done()

		if spreadDuration > 0 {
			// Spread in time by waiting up to 3s less than the interval.
			spreadDurationMs := int(spreadDuration.Milliseconds())
			spreadRange := spreadDurationMs - 3000

			if spreadRange > 0 {
				time.Sleep(time.Duration(rand.Intn(spreadRange)) * time.Millisecond)
			}
		}

		logger.Debug("Sending heartbeat", log.Ctx{"address": address})

		// Update timestamp to current, used for time skew detection
		heartbeatData.Time = time.Now().UTC()

		err := HeartbeatNode(ctx, address, networkCert, serverCert, heartbeatData)
		if err == nil {
			heartbeatData.Lock()
			// Ensure only update nodes that exist in Members already.
			hbNode, existing := hbState.Members[nodeID]
			if !existing {
				return
			}

			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			heartbeatData.Members[nodeID] = hbNode
			heartbeatData.Unlock()
			logger.Debug("Successful heartbeat", log.Ctx{"address": address})
		} else {
			logger.Debug("Failed heartbeat", log.Ctx{"address": address, "err": err})
		}
	}

	for _, node := range nodes {
		// Special case for the local node - just record the time now.
		if node.Address == localAddress {
			hbState.Lock()
			hbNode := hbState.Members[node.ID]
			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			hbState.Members[node.ID] = hbNode
			hbState.Unlock()
			continue
		}

		// Parallelize the rest.
		heartbeatsWg.Add(1)
		go sendHeartbeat(node.ID, node.Address, spreadDuration, hbState)
	}

	heartbeatsWg.Wait()
}

// HeartbeatTask returns a task function that performs leader-initiated heartbeat
// checks against all LXD nodes in the cluster.
//
// It will update the heartbeat timestamp column of the nodes table
// accordingly, and also notify them of the current list of database nodes.
func HeartbeatTask(gateway *Gateway) (task.Func, task.Schedule) {
	// Since the database APIs are blocking we need to wrap the core logic
	// and run it in a goroutine, so we can abort as soon as the context expires.
	heartbeatWrapper := func(ctx context.Context) {
		ch := make(chan struct{})
		go func() {
			gateway.heartbeat(ctx, false)
			ch <- struct{}{}
		}()
		select {
		case <-ch:
		case <-ctx.Done():
		}
	}

	schedule := task.Every(time.Duration(heartbeatInterval) * time.Second)

	return heartbeatWrapper, schedule
}

// heartbeatInterval returns heartbeat interval to use.
func (g *Gateway) heartbeatInterval() time.Duration {
	threshold := g.HeartbeatOfflineThreshold
	if threshold <= 0 {
		threshold = db.DefaultOfflineThreshold
	}

	return threshold / 2
}

func (g *Gateway) heartbeat(ctx context.Context, initialHeartbeat bool) {
	// Avoid concurent heartbeat loops.
	// This is possible when both the task and the out of band heartbeat
	// round from a dqlite connection both kick in at the same time.
	g.heartbeatLock.Lock()
	defer g.heartbeatLock.Unlock()

	if g.Cluster == nil || g.server == nil || g.memoryDial != nil {
		// We're not a raft node or we're not clustered
		return
	}

	raftNodes, err := g.currentRaftNodes()
	if err != nil {
		if errors.Cause(err) == ErrNotLeader {
			return
		}

		logger.Error("Failed to get current raft members", log.Ctx{"err": err})
		return
	}

	if initialHeartbeat {
		logger.Debug("Starting heartbeat round (initial)")
	} else {
		logger.Debug("Starting heartbeat round")
	}

	// Replace the local raft_nodes table immediately because it
	// might miss a row containing ourselves, since we might have
	// been elected leader before the former leader had chance to
	// send us a fresh update through the heartbeat pool.
	logger.Debug("Heartbeat updating local raft members", log.Ctx{"members": raftNodes})
	err = g.db.Transaction(func(tx *db.NodeTx) error {
		return tx.ReplaceRaftNodes(raftNodes)
	})
	if err != nil {
		logger.Warn("Failed to replace local raft members", log.Ctx{"err": err})
		return
	}

	startTime := time.Now()

	var allNodes []db.NodeInfo
	var localAddress string // Address of this node
	var offlineThreshold time.Duration
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		allNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		localAddress, err = tx.GetLocalNodeAddress()
		if err != nil {
			return err
		}

		offlineThreshold, err = tx.GetNodeOfflineThreshold()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	// Cumulative set of node states (will be written back to database once done).
	hbState := &APIHeartbeat{cluster: g.Cluster}

	// If this leader node hasn't sent a heartbeat recently, then its node state records
	// are likely out of date, this can happen when a node becomes a leader.
	// Send stale set to all nodes in database to get a fresh set of active nodes.
	if initialHeartbeat {
		hbState.Update(false, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, false)

		// We have the latest set of node states now, lets send that state set to all nodes.
		hbState.FullStateList = true
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, false)
	} else {
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, true)
	}

	// Look for any new node which appeared since sending last heartbeat.
	var currentNodes []db.NodeInfo
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		currentNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	newNodes := []db.NodeInfo{}
	for _, currentNode := range currentNodes {
		existing := false
		for _, node := range allNodes {
			if node.Address == currentNode.Address && node.ID == currentNode.ID {
				existing = true
				break
			}
		}

		if !existing {
			// We found a new node
			allNodes = append(allNodes, currentNode)
			newNodes = append(newNodes, currentNode)
		}
	}

	// If any new nodes found, send heartbeat to just them (with full node state).
	if len(newNodes) > 0 {
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, newNodes, false)
	}

	// If the context has been cancelled, return immediately.
	err = ctx.Err()
	if err != nil {
		logger.Warn("Aborting heartbeat round", log.Ctx{"err": err})
		return
	}

	err = query.Retry(func() error {
		return g.Cluster.Transaction(func(tx *db.ClusterTx) error {
			for _, node := range hbState.Members {
				if !node.updated {
					continue
				}

				err := tx.SetNodeHeartbeat(node.Address, node.LastHeartbeat)
				if err != nil && errors.Cause(err) != db.ErrNoSuchObject {
					return errors.Wrapf(err, "Failed updating heartbeat time for member %q", node.Address)
				}
			}

			return nil
		})
	})
	if err != nil {
		logger.Error("Failed updating cluster heartbeats", log.Ctx{"err": err})
		return
	}

	// If full node state was sent and node refresh task is specified, run it async.
	if g.HeartbeatNodeHook != nil {
		go g.HeartbeatNodeHook(hbState)
	}

	duration := time.Now().Sub(startTime)
	if duration > (time.Duration(heartbeatInterval) * time.Second) {
		logger.Warn("Heartbeat round duration greater than heartbeat interval", log.Ctx{"duration": duration, "interval": heartbeatInterval})
	}

	// Update last leader heartbeat time so next time a full node state list can be sent (if not this time).
	logger.Debug("Completed heartbeat round", log.Ctx{"duration": duration})
}

// heartbeatInterval Number of seconds to wait between to heartbeat rounds.
const heartbeatInterval = 10

// HeartbeatNode performs a single heartbeat request against the node with the given address.
func HeartbeatNode(taskCtx context.Context, address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, heartbeatData *APIHeartbeat) error {
	logger.Debug("Sending heartbeat request", log.Ctx{"address": address})

	config, err := tlsClientConfig(networkCert, serverCert)
	if err != nil {
		return err
	}

	timeout := 2 * time.Second
	url := fmt.Sprintf("https://%s%s", address, databaseEndpoint)
	transport, cleanup := tlsTransport(config)
	defer cleanup()
	client := &http.Client{
		Transport: transport,
		Timeout:   timeout,
	}

	buffer := bytes.Buffer{}
	heartbeatData.Lock()
	err = json.NewEncoder(&buffer).Encode(heartbeatData)
	heartbeatData.Unlock()
	if err != nil {
		return err
	}

	request, err := http.NewRequest("PUT", url, bytes.NewReader(buffer.Bytes()))
	if err != nil {
		return err
	}
	setDqliteVersionHeader(request)

	// Use 1s later timeout to give HTTP client chance timeout with more useful info.
	ctx, cancel := context.WithTimeout(taskCtx, timeout+time.Second)
	defer cancel()
	request = request.WithContext(ctx)
	request.Close = true // Immediately close the connection after the request is done

	response, err := client.Do(request)
	if err != nil {
		return errors.Wrap(err, "Failed to send heartbeat request")
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("Heartbeat request failed with status code: %s", response.Status)
	}

	return nil
}
<commit_msg>lxd/cluster/heartbeat: Updates HeartbeatTask to use gateway.heartbeatInterval<commit_after>package cluster

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/db/query"
	"github.com/lxc/lxd/lxd/task"
	"github.com/lxc/lxd/shared"
	log "github.com/lxc/lxd/shared/log15"
	"github.com/lxc/lxd/shared/logger"
	"github.com/pkg/errors"
)

// APIHeartbeatMember contains specific cluster node info.
type APIHeartbeatMember struct {
	ID            int64     // ID field value in nodes table.
	Address       string    // Host and Port of node.
	RaftID        uint64    // ID field value in raft_nodes table, zero if non-raft node.
	RaftRole      int       // Node role in the raft cluster, from the raft_nodes table
	Raft          bool      // Deprecated, use non-zero RaftID instead to indicate raft node.
	LastHeartbeat time.Time // Last time we received a successful response from node.
	Online        bool      // Calculated from offline threshold and LastHeatbeat time.
	updated       bool      // Has node been updated during this heartbeat run. Not sent to nodes.
}

// APIHeartbeatVersion contains max versions for all nodes in cluster.
type APIHeartbeatVersion struct {
	Schema        int
	APIExtensions int
}

// APIHeartbeat contains data sent to nodes in heartbeat.
type APIHeartbeat struct {
	sync.Mutex // Used to control access to Members maps.
	cluster    *db.Cluster
	Members    map[int64]APIHeartbeatMember
	Version    APIHeartbeatVersion
	Time       time.Time

	// Indicates if heartbeat contains a fresh set of node states.
	// This can be used to indicate to the receiving node that the state is fresh enough to
	// trigger node refresh activies (such as forkdns).
	FullStateList bool
}

// Update updates an existing APIHeartbeat struct with the raft and all node states supplied.
// If allNodes provided is an empty set then this is considered a non-full state list.
func (hbState *APIHeartbeat) Update(fullStateList bool, raftNodes []db.RaftNode, allNodes []db.NodeInfo, offlineThreshold time.Duration) {
	var maxSchemaVersion, maxAPIExtensionsVersion int

	if hbState.Members == nil {
		hbState.Members = make(map[int64]APIHeartbeatMember)
	}

	// If we've been supplied a fresh set of node states, this is a full state list.
	hbState.FullStateList = fullStateList

	raftNodeMap := make(map[string]db.RaftNode)

	// Convert raftNodes to a map keyed on address for lookups later.
	for _, raftNode := range raftNodes {
		raftNodeMap[raftNode.Address] = raftNode
	}

	// Add nodes (overwrites any nodes with same ID in map with fresh data).
	for _, node := range allNodes {
		member := APIHeartbeatMember{
			ID:            node.ID,
			Address:       node.Address,
			LastHeartbeat: node.Heartbeat,
			Online:        !node.Heartbeat.Before(time.Now().Add(-offlineThreshold)),
		}

		if raftNode, exists := raftNodeMap[member.Address]; exists {
			member.Raft = true // Deprecated
			member.RaftID = raftNode.ID
			member.RaftRole = int(raftNode.Role)
			delete(raftNodeMap, member.Address) // Used to check any remaining later.
		}

		// Add to the members map using the node ID (not the Raft Node ID).
		hbState.Members[node.ID] = member

		// Keep a record of highest APIExtensions and Schema version seen in all nodes.
		if node.APIExtensions > maxAPIExtensionsVersion {
			maxAPIExtensionsVersion = node.APIExtensions
		}

		if node.Schema > maxSchemaVersion {
			maxSchemaVersion = node.Schema
		}
	}

	hbState.Version = APIHeartbeatVersion{
		Schema:        maxSchemaVersion,
		APIExtensions: maxAPIExtensionsVersion,
	}

	if len(raftNodeMap) > 0 && hbState.cluster != nil {
		hbState.cluster.Transaction(func(tx *db.ClusterTx) error {
			for addr, raftNode := range raftNodeMap {
				_, err := tx.GetPendingNodeByAddress(addr)
				if err != nil {
					logger.Errorf("Unaccounted raft node(s) not found in 'nodes' table for heartbeat: %+v", raftNode)
				}
			}

			return nil
		})
	}

	return
}

// Send sends heartbeat requests to the nodes supplied and updates heartbeat state.
func (hbState *APIHeartbeat) Send(ctx context.Context, networkCert *shared.CertInfo, serverCert *shared.CertInfo, localAddress string, nodes []db.NodeInfo, spreadDuration time.Duration) {
	heartbeatsWg := sync.WaitGroup{}
	sendHeartbeat := func(nodeID int64, address string, spreadDuration time.Duration, heartbeatData *APIHeartbeat) {
		defer heartbeatsWg.Done()

		if spreadDuration > 0 {
			// Spread in time by waiting up to 3s less than the interval.
			spreadDurationMs := int(spreadDuration.Milliseconds())
			spreadRange := spreadDurationMs - 3000

			if spreadRange > 0 {
				time.Sleep(time.Duration(rand.Intn(spreadRange)) * time.Millisecond)
			}
		}

		logger.Debug("Sending heartbeat", log.Ctx{"address": address})

		// Update timestamp to current, used for time skew detection
		heartbeatData.Time = time.Now().UTC()

		err := HeartbeatNode(ctx, address, networkCert, serverCert, heartbeatData)
		if err == nil {
			heartbeatData.Lock()
			// Ensure only update nodes that exist in Members already.
			hbNode, existing := hbState.Members[nodeID]
			if !existing {
				return
			}

			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			heartbeatData.Members[nodeID] = hbNode
			heartbeatData.Unlock()
			logger.Debug("Successful heartbeat", log.Ctx{"address": address})
		} else {
			logger.Debug("Failed heartbeat", log.Ctx{"address": address, "err": err})
		}
	}

	for _, node := range nodes {
		// Special case for the local node - just record the time now.
		if node.Address == localAddress {
			hbState.Lock()
			hbNode := hbState.Members[node.ID]
			hbNode.LastHeartbeat = time.Now()
			hbNode.Online = true
			hbNode.updated = true
			hbState.Members[node.ID] = hbNode
			hbState.Unlock()
			continue
		}

		// Parallelize the rest.
		heartbeatsWg.Add(1)
		go sendHeartbeat(node.ID, node.Address, spreadDuration, hbState)
	}

	heartbeatsWg.Wait()
}

// HeartbeatTask returns a task function that performs leader-initiated heartbeat
// checks against all LXD nodes in the cluster.
//
// It will update the heartbeat timestamp column of the nodes table
// accordingly, and also notify them of the current list of database nodes.
func HeartbeatTask(gateway *Gateway) (task.Func, task.Schedule) {
	// Since the database APIs are blocking we need to wrap the core logic
	// and run it in a goroutine, so we can abort as soon as the context expires.
	heartbeatWrapper := func(ctx context.Context) {
		ch := make(chan struct{})
		go func() {
			gateway.heartbeat(ctx, false)
			ch <- struct{}{}
		}()
		select {
		case <-ch:
		case <-ctx.Done():
		}
	}

	schedule := func() (time.Duration, error) {
		return task.Every(gateway.heartbeatInterval())()
	}

	return heartbeatWrapper, schedule
}

// heartbeatInterval returns heartbeat interval to use.
func (g *Gateway) heartbeatInterval() time.Duration {
	threshold := g.HeartbeatOfflineThreshold
	if threshold <= 0 {
		threshold = db.DefaultOfflineThreshold
	}

	return threshold / 2
}

func (g *Gateway) heartbeat(ctx context.Context, initialHeartbeat bool) {
	// Avoid concurent heartbeat loops.
	// This is possible when both the task and the out of band heartbeat
	// round from a dqlite connection both kick in at the same time.
	g.heartbeatLock.Lock()
	defer g.heartbeatLock.Unlock()

	if g.Cluster == nil || g.server == nil || g.memoryDial != nil {
		// We're not a raft node or we're not clustered
		return
	}

	raftNodes, err := g.currentRaftNodes()
	if err != nil {
		if errors.Cause(err) == ErrNotLeader {
			return
		}

		logger.Error("Failed to get current raft members", log.Ctx{"err": err})
		return
	}

	if initialHeartbeat {
		logger.Debug("Starting heartbeat round (initial)")
	} else {
		logger.Debug("Starting heartbeat round")
	}

	// Replace the local raft_nodes table immediately because it
	// might miss a row containing ourselves, since we might have
	// been elected leader before the former leader had chance to
	// send us a fresh update through the heartbeat pool.
	logger.Debug("Heartbeat updating local raft members", log.Ctx{"members": raftNodes})
	err = g.db.Transaction(func(tx *db.NodeTx) error {
		return tx.ReplaceRaftNodes(raftNodes)
	})
	if err != nil {
		logger.Warn("Failed to replace local raft members", log.Ctx{"err": err})
		return
	}

	startTime := time.Now()

	var allNodes []db.NodeInfo
	var localAddress string // Address of this node
	var offlineThreshold time.Duration
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		allNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		localAddress, err = tx.GetLocalNodeAddress()
		if err != nil {
			return err
		}

		offlineThreshold, err = tx.GetNodeOfflineThreshold()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	// Cumulative set of node states (will be written back to database once done).
	hbState := &APIHeartbeat{cluster: g.Cluster}

	// If this leader node hasn't sent a heartbeat recently, then its node state records
	// are likely out of date, this can happen when a node becomes a leader.
	// Send stale set to all nodes in database to get a fresh set of active nodes.
	if initialHeartbeat {
		hbState.Update(false, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, false)

		// We have the latest set of node states now, lets send that state set to all nodes.
		hbState.FullStateList = true
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, false)
	} else {
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, allNodes, true)
	}

	// Look for any new node which appeared since sending last heartbeat.
	var currentNodes []db.NodeInfo
	err = g.Cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		currentNodes, err = tx.GetNodes()
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		logger.Warn("Failed to get current cluster members", log.Ctx{"err": err})
		return
	}

	newNodes := []db.NodeInfo{}
	for _, currentNode := range currentNodes {
		existing := false
		for _, node := range allNodes {
			if node.Address == currentNode.Address && node.ID == currentNode.ID {
				existing = true
				break
			}
		}

		if !existing {
			// We found a new node
			allNodes = append(allNodes, currentNode)
			newNodes = append(newNodes, currentNode)
		}
	}

	// If any new nodes found, send heartbeat to just them (with full node state).
	if len(newNodes) > 0 {
		hbState.Update(true, raftNodes, allNodes, offlineThreshold)
		hbState.Send(ctx, g.networkCert, g.serverCert(), localAddress, newNodes, false)
	}

	// If the context has been cancelled, return immediately.
	err = ctx.Err()
	if err != nil {
		logger.Warn("Aborting heartbeat round", log.Ctx{"err": err})
		return
	}

	err = query.Retry(func() error {
		return g.Cluster.Transaction(func(tx *db.ClusterTx) error {
			for _, node := range hbState.Members {
				if !node.updated {
					continue
				}

				err := tx.SetNodeHeartbeat(node.Address, node.LastHeartbeat)
				if err != nil && errors.Cause(err) != db.ErrNoSuchObject {
					return errors.Wrapf(err, "Failed updating heartbeat time for member %q", node.Address)
				}
			}

			return nil
		})
	})
	if err != nil {
		logger.Error("Failed updating cluster heartbeats", log.Ctx{"err": err})
		return
	}

	// If full node state was sent and node refresh task is specified, run it async.
	if g.HeartbeatNodeHook != nil {
		go g.HeartbeatNodeHook(hbState)
	}

	duration := time.Now().Sub(startTime)
	if duration > (time.Duration(heartbeatInterval) * time.Second) {
		logger.Warn("Heartbeat round duration greater than heartbeat interval", log.Ctx{"duration": duration, "interval": heartbeatInterval})
	}

	// Update last leader heartbeat time so next time a full node state list can be sent (if not this time).
	logger.Debug("Completed heartbeat round", log.Ctx{"duration": duration})
}

// heartbeatInterval Number of seconds to wait between to heartbeat rounds.
const heartbeatInterval = 10

// HeartbeatNode performs a single heartbeat request against the node with the given address.
func HeartbeatNode(taskCtx context.Context, address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, heartbeatData *APIHeartbeat) error {
	logger.Debug("Sending heartbeat request", log.Ctx{"address": address})

	config, err := tlsClientConfig(networkCert, serverCert)
	if err != nil {
		return err
	}

	timeout := 2 * time.Second
	url := fmt.Sprintf("https://%s%s", address, databaseEndpoint)
	transport, cleanup := tlsTransport(config)
	defer cleanup()
	client := &http.Client{
		Transport: transport,
		Timeout:   timeout,
	}

	buffer := bytes.Buffer{}
	heartbeatData.Lock()
	err = json.NewEncoder(&buffer).Encode(heartbeatData)
	heartbeatData.Unlock()
	if err != nil {
		return err
	}

	request, err := http.NewRequest("PUT", url, bytes.NewReader(buffer.Bytes()))
	if err != nil {
		return err
	}
	setDqliteVersionHeader(request)

	// Use 1s later timeout to give HTTP client chance timeout with more useful info.
	ctx, cancel := context.WithTimeout(taskCtx, timeout+time.Second)
	defer cancel()
	request = request.WithContext(ctx)
	request.Close = true // Immediately close the connection after the request is done

	response, err := client.Do(request)
	if err != nil {
		return errors.Wrap(err, "Failed to send heartbeat request")
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("Heartbeat request failed with status code: %s", response.Status)
	}

	return nil
}
<|endoftext|>