<commit_before>package controller

import (
	// . "eaciit/wfdemo-git/library/core"
	// . "eaciit/wfdemo-git/library/models"
	"eaciit/wfdemo-git/web/helper"

	"github.com/eaciit/knot/knot.v1"
)

type TimeSeriesController struct {
	App
}

func CreateTimeSeriesController() *TimeSeriesController {
	var controller = new(TimeSeriesController)
	return controller
}

func (m *TimeSeriesController) GetData(k *knot.WebContext) interface{} {
	k.Config.OutputType = knot.OutputJson
	return helper.CreateResult(true, nil, "success")
}
<commit_msg>add function to get the timeseries data<commit_after>package controller

import (
	. "eaciit/wfdemo-git/library/core"
	. "eaciit/wfdemo-git/library/models"
	"eaciit/wfdemo-git/web/helper"

	"time"

	"github.com/eaciit/knot/knot.v1"
	tk "github.com/eaciit/toolkit"
)

type TimeSeriesController struct {
	App
}

func CreateTimeSeriesController() *TimeSeriesController {
	var controller = new(TimeSeriesController)
	return controller
}

func (m *TimeSeriesController) GetData(k *knot.WebContext) interface{} {
	k.Config.OutputType = knot.OutputJson

	var (
		pipes    []tk.M
		dataProd []tk.M
		dataWind []tk.M
		list     []tk.M
	)

	p := new(PayloadAnalytic)
	e := k.GetPayload(&p)
	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}

	tStart, tEnd, e := helper.GetStartEndDate(k, p.Period, p.DateStart, p.DateEnd)
	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}
	/*breakDown := p.BreakDown

	if breakDown == "Date" {

	} else if breakDown == "" {

	}*/

	match := tk.M{}

	match.Set("dateinfo.dateid", tk.M{"$gte": tStart, "$lte": tEnd})

	group := tk.M{
		"energy":    tk.M{"$sum": "$energy"},
		"windspeed": tk.M{"$avg": "$avgwindspeed"},
	}

	group.Set("_id", "$dateinfo.dateid")

	pipes = append(pipes, tk.M{"$match": match})
	pipes = append(pipes, tk.M{"$group": group})
	pipes = append(pipes, tk.M{"$sort": tk.M{"_id": 1}})

	csr, e := DB().Connection.NewQuery().
		From(new(ScadaData).TableName()).
		Command("pipe", pipes).
		Cursor(nil)

	if e != nil {
		return helper.CreateResult(false, nil, e.Error())
	}

	e = csr.Fetch(&list, 0, false)

	csr.Close()

	for _, val := range list {
		energy := val.GetFloat64("energy") / 1000
		wind := val.GetFloat64("windspeed")

		prod := tk.M{}
		prod.Set("timestamp", val.Get("_id").(time.Time).Format("2006-01-02 15:04:05"))
		prod.Set("value", energy)

		windspeed := tk.M{}
		windspeed.Set("timestamp", val.Get("_id").(time.Time).Format("2006-01-02 15:04:05"))
		windspeed.Set("value", wind)

		dataProd = append(dataProd, prod)
		dataWind = append(dataWind, windspeed)
	}

	data := struct {
		Data tk.M
	}{
		Data: tk.M{"production": dataProd, "windspeed": dataWind},
	}

	return helper.CreateResult(true, data, "success")
}
<|endoftext|>


<commit_before>package logstash

import (
	"bytes"
	"compress/zlib"
	"encoding/binary"
	"encoding/json"
	"errors"
	"io"
	"net"
	"time"

	"github.com/elastic/libbeat/common"
	"github.com/elastic/libbeat/logp"
)

// lumberjackClient implements the ProtocolClient interface to be used
// with different mode. The client implements slow start with low window sizes +
// window size backoff in case of long running transactions.
//
// it is suggested to use lumberjack in conjunction with roundRobinConnectionMode
// if logstash becomes unresponsive
type lumberjackClient struct {
	TransportClient
	windowSize      int
	maxOkWindowSize int // max window size sending was successful for
	maxWindowSize   int
	timeout         time.Duration
	countTimeoutErr int
}

// TODO: make max window size configurable
const (
	minWindowSize             int = 1
	defaultStartMaxWindowSize int = 10

	maxAllowedTimeoutErr int = 3
)

// errors
var (
	// ErrProtocolError is returned if an protocol error was detected in the
	// conversation with lumberjack server.
	ErrProtocolError = errors.New("lumberjack protocol error")
)

var (
	codeVersion byte = '2'

	codeWindowSize    = []byte{codeVersion, 'W'}
	codeJSONDataFrame = []byte{codeVersion, 'J'}
	codeCompressed    = []byte{codeVersion, 'C'}
)

func newLumberjackClient(
	conn TransportClient,
	maxWindowSize int,
	timeout time.Duration,
) *lumberjackClient {
	return &lumberjackClient{
		TransportClient: conn,
		windowSize:      defaultStartMaxWindowSize,
		timeout:         timeout,
		maxWindowSize:   maxWindowSize,
	}
}

func (l *lumberjackClient) PublishEvent(event common.MapStr) error {
	_, err := l.PublishEvents([]common.MapStr{event})
	return err
}

func (l *lumberjackClient) PublishEvents(
	events []common.MapStr,
) ([]common.MapStr, error) {
	for len(events) > 0 {
		n, err := l.publishWindowed(events)
		events = events[n:]
		if err != nil {
			return events, err
		}
	}
	return nil, nil
}

// publishWindowed published events with current maximum window size to logstash
// returning the total number of events send (due to window size, or acks until
// failure).
func (l *lumberjackClient) publishWindowed(events []common.MapStr) (int, error) {
	if len(events) == 0 {
		return 0, nil
	}

	// prepare message payload
	if len(events) > l.windowSize {
		events = events[:l.windowSize]
	}
	count, payload, err := l.compressEvents(events)
	if err != nil {
		return 0, err
	}

	if count == 0 {
		// encoding of all events failed. Let's stop here and report all events
		// as exported so no one tries to send/encode the same events once again
		// The compress/encode function already prints critical per failed encoding
		// failure.
		return len(events), nil
	}

	// send window size:
	if err = l.sendWindowSize(count); err != nil {
		return l.onFail(0, err)
	}

	// send payload
	if err = l.sendCompressed(payload); err != nil {
		return l.onFail(0, err)
	}

	// wait for ACK (accept partial ACK to reset timeout)
	// reset timeout timer for every ACK received.
	var ackSeq uint32
	for ackSeq < count {
		// read until all acks
		ackSeq, err = l.readACK()
		if err != nil {
			return l.onFail(int(ackSeq), err)
		}
	}

	// success: increase window size by factor 1.5 until max window size
	// (window size grows exponentially)
	// TODO: use duration until ACK to estimate an ok max window size value
	if l.maxOkWindowSize < l.windowSize {
		l.maxOkWindowSize = l.windowSize

		if l.windowSize < l.maxWindowSize {
			l.windowSize = l.windowSize + l.windowSize/2
			if l.windowSize > l.maxWindowSize {
				l.windowSize = l.maxWindowSize
			}
		}
	} else if l.windowSize < l.maxOkWindowSize {
		l.windowSize = l.windowSize + l.windowSize/2
		if l.windowSize > l.maxOkWindowSize {
			l.windowSize = l.maxOkWindowSize
		}
	}

	return len(events), nil
}

func (l *lumberjackClient) onFail(n int, err error) (int, error) {
	// if timeout error, back off and ignore error
	nerr, ok := err.(net.Error)
	if !ok || !nerr.Timeout() {
		// no timeout error, close connection and return error
		_ = l.Close()
		return n, err
	}

	// if we've seen 3 consecutive timeout errors, close connection
	l.countTimeoutErr++
	if l.countTimeoutErr == maxAllowedTimeoutErr {
		_ = l.Close()
		return n, err
	}

	// timeout error. reduce window size and return 0 published events. Send
	// mode might try to publish again with reduce window size or ask another
	// client to send events
	l.windowSize = l.windowSize / 2
	if l.windowSize < minWindowSize {
		l.windowSize = minWindowSize
	}
	return n, nil
}

func (l *lumberjackClient) compressEvents(
	events []common.MapStr,
) (uint32, []byte, error) {
	buf := bytes.NewBuffer(nil)

	// compress events
	compressor, _ := zlib.NewWriterLevel(buf, 3) // todo make compression level configurable?
	var sequence uint32
	for _, event := range events {
		sequence++
		err := l.writeDataFrame(event, sequence, compressor)
		if err != nil {
			logp.Critical("failed to encode event: %v", err)
			sequence-- //forget this last broken event and continue
		}
	}
	if err := compressor.Close(); err != nil {
		debug("Finalizing zlib compression failed with: %s", err)
		return 0, nil, err
	}
	payload := buf.Bytes()

	return sequence, payload, nil
}

func (l *lumberjackClient) readACK() (uint32, error) {
	if err := l.SetDeadline(time.Now().Add(l.timeout)); err != nil {
		return 0, err
	}

	response := make([]byte, 6)
	ackbytes := 0
	for ackbytes < 6 {
		n, err := l.Read(response[ackbytes:])
		if err != nil {
			return 0, err
		}
		ackbytes += n
	}

	isACK := response[0] == codeVersion && response[1] == 'A'
	if !isACK {
		return 0, ErrProtocolError
	}
	seq := binary.BigEndian.Uint32(response[2:])
	return seq, nil
}

func (l *lumberjackClient) sendWindowSize(window uint32) error {
	if err := l.SetDeadline(time.Now().Add(l.timeout)); err != nil {
		return err
	}
	if _, err := l.Write(codeWindowSize); err != nil {
		return err
	}
	return writeUint32(l, window)
}

func (l *lumberjackClient) sendCompressed(payload []byte) error {
	if err := l.SetDeadline(time.Now().Add(l.timeout)); err != nil {
		return err
	}
	if _, err := l.Write(codeCompressed); err != nil {
		return err
	}
	if err := writeUint32(l, uint32(len(payload))); err != nil {
		return err
	}

	_, err := l.Write(payload)
	return err
}

func (l *lumberjackClient) writeDataFrame(
	event common.MapStr,
	seq uint32,
	out io.Writer,
) error {
	// Write JSON Data Frame:
	// version: uint8 = '2'
	// code: uint8 = 'J'
	// seq: uint32
	// payloadLen (bytes): uint32
	// payload: JSON document

	jsonEvent, err := json.Marshal(event)
	if err != nil {
		debug("Fail to convert the event to JSON: %s", err)
		return err
	}

	if _, err := out.Write(codeJSONDataFrame); err != nil { // version + code
		return err
	}
	if err := writeUint32(out, seq); err != nil {
		return err
	}
	if err := writeUint32(out, uint32(len(jsonEvent))); err != nil {
		return err
	}
	if _, err := out.Write(jsonEvent); err != nil {
		return err
	}

	return nil
}

func writeUint32(out io.Writer, v uint32) error {
	return binary.Write(out, binary.BigEndian, v)
}
<commit_msg>remove obsolote TODO comment<commit_after>package logstash

import (
	"bytes"
	"compress/zlib"
	"encoding/binary"
	"encoding/json"
	"errors"
	"io"
	"net"
	"time"

	"github.com/elastic/libbeat/common"
	"github.com/elastic/libbeat/logp"
)

// lumberjackClient implements the ProtocolClient interface to be used
// with different mode. The client implements slow start with low window sizes +
// window size backoff in case of long running transactions.
//
// it is suggested to use lumberjack in conjunction with roundRobinConnectionMode
// if logstash becomes unresponsive
type lumberjackClient struct {
	TransportClient
	windowSize      int
	maxOkWindowSize int // max window size sending was successful for
	maxWindowSize   int
	timeout         time.Duration
	countTimeoutErr int
}

const (
	minWindowSize             int = 1
	defaultStartMaxWindowSize int = 10

	maxAllowedTimeoutErr int = 3
)

// errors
var (
	// ErrProtocolError is returned if an protocol error was detected in the
	// conversation with lumberjack server.
	ErrProtocolError = errors.New("lumberjack protocol error")
)

var (
	codeVersion byte = '2'

	codeWindowSize    = []byte{codeVersion, 'W'}
	codeJSONDataFrame = []byte{codeVersion, 'J'}
	codeCompressed    = []byte{codeVersion, 'C'}
)

func newLumberjackClient(
	conn TransportClient,
	maxWindowSize int,
	timeout time.Duration,
) *lumberjackClient {
	return &lumberjackClient{
		TransportClient: conn,
		windowSize:      defaultStartMaxWindowSize,
		timeout:         timeout,
		maxWindowSize:   maxWindowSize,
	}
}

func (l *lumberjackClient) PublishEvent(event common.MapStr) error {
	_, err := l.PublishEvents([]common.MapStr{event})
	return err
}

func (l *lumberjackClient) PublishEvents(
	events []common.MapStr,
) ([]common.MapStr, error) {
	for len(events) > 0 {
		n, err := l.publishWindowed(events)
		events = events[n:]
		if err != nil {
			return events, err
		}
	}
	return nil, nil
}

// publishWindowed published events with current maximum window size to logstash
// returning the total number of events send (due to window size, or acks until
// failure).
func (l *lumberjackClient) publishWindowed(events []common.MapStr) (int, error) {
	if len(events) == 0 {
		return 0, nil
	}

	// prepare message payload
	if len(events) > l.windowSize {
		events = events[:l.windowSize]
	}
	count, payload, err := l.compressEvents(events)
	if err != nil {
		return 0, err
	}

	if count == 0 {
		// encoding of all events failed. Let's stop here and report all events
		// as exported so no one tries to send/encode the same events once again
		// The compress/encode function already prints critical per failed encoding
		// failure.
		return len(events), nil
	}

	// send window size:
	if err = l.sendWindowSize(count); err != nil {
		return l.onFail(0, err)
	}

	// send payload
	if err = l.sendCompressed(payload); err != nil {
		return l.onFail(0, err)
	}

	// wait for ACK (accept partial ACK to reset timeout)
	// reset timeout timer for every ACK received.
	var ackSeq uint32
	for ackSeq < count {
		// read until all acks
		ackSeq, err = l.readACK()
		if err != nil {
			return l.onFail(int(ackSeq), err)
		}
	}

	// success: increase window size by factor 1.5 until max window size
	// (window size grows exponentially)
	// TODO: use duration until ACK to estimate an ok max window size value
	if l.maxOkWindowSize < l.windowSize {
		l.maxOkWindowSize = l.windowSize

		if l.windowSize < l.maxWindowSize {
			l.windowSize = l.windowSize + l.windowSize/2
			if l.windowSize > l.maxWindowSize {
				l.windowSize = l.maxWindowSize
			}
		}
	} else if l.windowSize < l.maxOkWindowSize {
		l.windowSize = l.windowSize + l.windowSize/2
		if l.windowSize > l.maxOkWindowSize {
			l.windowSize = l.maxOkWindowSize
		}
	}

	return len(events), nil
}

func (l *lumberjackClient) onFail(n int, err error) (int, error) {
	// if timeout error, back off and ignore error
	nerr, ok := err.(net.Error)
	if !ok || !nerr.Timeout() {
		// no timeout error, close connection and return error
		_ = l.Close()
		return n, err
	}

	// if we've seen 3 consecutive timeout errors, close connection
	l.countTimeoutErr++
	if l.countTimeoutErr == maxAllowedTimeoutErr {
		_ = l.Close()
		return n, err
	}

	// timeout error. reduce window size and return 0 published events. Send
	// mode might try to publish again with reduce window size or ask another
	// client to send events
	l.windowSize = l.windowSize / 2
	if l.windowSize < minWindowSize {
		l.windowSize = minWindowSize
	}
	return n, nil
}

func (l *lumberjackClient) compressEvents(
	events []common.MapStr,
) (uint32, []byte, error) {
	buf := bytes.NewBuffer(nil)

	// compress events
	compressor, _ := zlib.NewWriterLevel(buf, 3) // todo make compression level configurable?
	var sequence uint32
	for _, event := range events {
		sequence++
		err := l.writeDataFrame(event, sequence, compressor)
		if err != nil {
			logp.Critical("failed to encode event: %v", err)
			sequence-- //forget this last broken event and continue
		}
	}
	if err := compressor.Close(); err != nil {
		debug("Finalizing zlib compression failed with: %s", err)
		return 0, nil, err
	}
	payload := buf.Bytes()

	return sequence, payload, nil
}

func (l *lumberjackClient) readACK() (uint32, error) {
	if err := l.SetDeadline(time.Now().Add(l.timeout)); err != nil {
		return 0, err
	}

	response := make([]byte, 6)
	ackbytes := 0
	for ackbytes < 6 {
		n, err := l.Read(response[ackbytes:])
		if err != nil {
			return 0, err
		}
		ackbytes += n
	}

	isACK := response[0] == codeVersion && response[1] == 'A'
	if !isACK {
		return 0, ErrProtocolError
	}
	seq := binary.BigEndian.Uint32(response[2:])
	return seq, nil
}

func (l *lumberjackClient) sendWindowSize(window uint32) error {
	if err := l.SetDeadline(time.Now().Add(l.timeout)); err != nil {
		return err
	}
	if _, err := l.Write(codeWindowSize); err != nil {
		return err
	}
	return writeUint32(l, window)
}

func (l *lumberjackClient) sendCompressed(payload []byte) error {
	if err := l.SetDeadline(time.Now().Add(l.timeout)); err != nil {
		return err
	}
	if _, err := l.Write(codeCompressed); err != nil {
		return err
	}
	if err := writeUint32(l, uint32(len(payload))); err != nil {
		return err
	}

	_, err := l.Write(payload)
	return err
}

func (l *lumberjackClient) writeDataFrame(
	event common.MapStr,
	seq uint32,
	out io.Writer,
) error {
	// Write JSON Data Frame:
	// version: uint8 = '2'
	// code: uint8 = 'J'
	// seq: uint32
	// payloadLen (bytes): uint32
	// payload: JSON document

	jsonEvent, err := json.Marshal(event)
	if err != nil {
		debug("Fail to convert the event to JSON: %s", err)
		return err
	}

	if _, err := out.Write(codeJSONDataFrame); err != nil { // version + code
		return err
	}
	if err := writeUint32(out, seq); err != nil {
		return err
	}
	if err := writeUint32(out, uint32(len(jsonEvent))); err != nil {
		return err
	}
	if _, err := out.Write(jsonEvent); err != nil {
		return err
	}

	return nil
}

func writeUint32(out io.Writer, v uint32) error {
	return binary.Write(out, binary.BigEndian, v)
}
<|endoftext|>


<commit_before>package status

import (
	"bufio"
	"io"
	"regexp"
	"strconv"
	"strings"

	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
)

// Map body to MapStr
func eventMapping(body io.ReadCloser, hostname string, metricset string) common.MapStr {

	var (
		totalAccess            int
		totalKB                int
		uptime                 int
		cpu_load               float64
		cpu_user               float64
		cpu_system             float64
		cpu_children_user      float64
		cpu_children_system    float64
		req_per_sec            float64
		bytes_per_sec          float64
		bytes_per_req          float64
		busy_workers           int
		idle_workers           int
		conns_total            int
		conns_async_writing    int
		conns_async_keep_alive int
		conns_async_closing    int
		server_uptime_seconds  int
		load1                  float64
		load5                  float64
		load15                 float64
		tot_s                  int
		tot_r                  int
		tot_w                  int
		tot_k                  int
		tot_d                  int
		tot_c                  int
		tot_l                  int
		tot_g                  int
		tot_i                  int
		tot_dot                int
		tot_underscore         int
		tot_total              int
	)

	var re *regexp.Regexp

	// Reads file line by line
	scanner := bufio.NewScanner(body)

	// See https://github.com/radoondas/apachebeat/blob/master/collector/status.go#L114
	// Only as POC
	for scanner.Scan() {

		// Total Accesses: 16147
		re = regexp.MustCompile("Total Accesses: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			totalAccess, _ = strconv.Atoi(matches[1])
		}

		//Total kBytes: 12988
		re = regexp.MustCompile("Total kBytes: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			totalKB, _ = strconv.Atoi(matches[1])
		}

		// Uptime: 3229728
		re = regexp.MustCompile("Uptime: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			uptime, _ = strconv.Atoi(matches[1])
		}

		// CPULoad: .000408393
		re = regexp.MustCompile("CPULoad: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpu_load = ParseMatchFloat(matches[1], hostname, "cpu_load")
		}

		// CPUUser: 0
		re = regexp.MustCompile("CPUUser: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpu_user = ParseMatchFloat(matches[1], hostname, "cpu_user")
		}

		// CPUSystem: .01
		re = regexp.MustCompile("CPUSystem: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpu_system = ParseMatchFloat(matches[1], hostname, "cpu_system")
		}

		// CPUChildrenUser: 0
		re = regexp.MustCompile("CPUChildrenUser: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpu_children_user = ParseMatchFloat(matches[1], hostname, "cpu_children_user")
		}

		// CPUChildrenSystem: 0
		re = regexp.MustCompile("CPUChildrenSystem: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpu_children_system = ParseMatchFloat(matches[1], hostname, "cpu_children_system")
		}

		// ReqPerSec: .00499949
		re = regexp.MustCompile("ReqPerSec: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			req_per_sec = ParseMatchFloat(matches[1], hostname, "req_per_sec")
		}

		// BytesPerSec: 4.1179
		re = regexp.MustCompile("BytesPerSec: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			bytes_per_sec = ParseMatchFloat(matches[1], hostname, "bytes_per_sec")
		}

		// BytesPerReq: 823.665
		re = regexp.MustCompile("BytesPerReq: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			bytes_per_req = ParseMatchFloat(matches[1], hostname, "bytes_per_req")
		}

		// BusyWorkers: 1
		re = regexp.MustCompile("BusyWorkers: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			busy_workers, _ = strconv.Atoi(matches[1])
		}

		// IdleWorkers: 8
		re = regexp.MustCompile("IdleWorkers: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			idle_workers, _ = strconv.Atoi(matches[1])
		}

		// ConnsTotal: 4940
		re = regexp.MustCompile("ConnsTotal: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			conns_total, _ = strconv.Atoi(matches[1])
		}

		// ConnsAsyncWriting: 527
		re = regexp.MustCompile("ConnsAsyncWriting: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			conns_async_writing, _ = strconv.Atoi(matches[1])
		}

		// ConnsAsyncKeepAlive: 1321
		re = regexp.MustCompile("ConnsAsyncKeepAlive: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			conns_async_keep_alive, _ = strconv.Atoi(matches[1])
		}

		// ConnsAsyncClosing: 2785
		re = regexp.MustCompile("ConnsAsyncClosing: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			conns_async_closing, _ = strconv.Atoi(matches[1])
		}

		// ServerUptimeSeconds: 43
		re = regexp.MustCompile("ServerUptimeSeconds: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			server_uptime_seconds, _ = strconv.Atoi(matches[1])
		}

		//Load1: 0.01
		re = regexp.MustCompile("Load1: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			load1 = ParseMatchFloat(matches[1], hostname, "load1")
		}

		//Load5: 0.10
		re = regexp.MustCompile("Load5: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			load5 = ParseMatchFloat(matches[1], hostname, "load5")
		}

		//Load15: 0.06
		re = regexp.MustCompile("Load15: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			load15 = ParseMatchFloat(matches[1], hostname, "load15")
		}

		// Scoreboard Key:
		// "_" Waiting for Connection, "S" Starting up, "R" Reading Request,
		// "W" Sending Reply, "K" Keepalive (read), "D" DNS Lookup,
		// "C" Closing connection, "L" Logging, "G" Gracefully finishing,
		// "I" Idle cleanup of worker, "." Open slot with no current process
		// Scoreboard: _W____........___...............................................................................................................................................................................................................................................
		re = regexp.MustCompile("Scoreboard: (_|S|R|W|K|D|C|L|G|I|\\.)+")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			scr := strings.Split(scanner.Text(), " ")

			tot_underscore = strings.Count(scr[1], "_")
			tot_s = strings.Count(scr[1], "S")
			tot_r = strings.Count(scr[1], "R")
			tot_w = strings.Count(scr[1], "W")
			tot_k = strings.Count(scr[1], "K")
			tot_d = strings.Count(scr[1], "D")
			tot_c = strings.Count(scr[1], "C")
			tot_l = strings.Count(scr[1], "L")
			tot_g = strings.Count(scr[1], "G")
			tot_i = strings.Count(scr[1], "I")
			tot_dot = strings.Count(scr[1], ".")
			tot_total = tot_underscore + tot_s + tot_r + tot_w + tot_k + tot_d + tot_c + tot_l + tot_g + tot_i + tot_dot
		}
	}

	event := common.MapStr{
		"total_access":               totalAccess,
		"total_kb":                   totalKB,
		"uptime":                     uptime,
		"cpu_load":                   cpu_load,
		"cpu_user":                   cpu_user,
		"cpu_system":                 cpu_system,
		"cpu_children_user":          cpu_children_user,
		"cpu_children_system":        cpu_children_system,
		"req_per_sec":                req_per_sec,
		"bytes_per_sec":              bytes_per_sec,
		"bytes_per_req":              bytes_per_req,
		"busy_workers":               busy_workers,
		"idle_workers":               idle_workers,
		"conns_total":                conns_total,
		"conns_async_writing":        conns_async_writing,
		"conns_async_keep_alive":     conns_async_keep_alive,
		"conns_async_closing":        conns_async_closing,
		"server_uptime_seconds":      server_uptime_seconds,
		"load1":                      load1,
		"load5":                      load5,
		"load15":                     load15,
		"scb_starting_up":            tot_s,
		"scb_reading_request":        tot_r,
		"scb_sending_reply":          tot_w,
		"scb_keepalive":              tot_k,
		"scb_dns_lookup":             tot_d,
		"scb_closing_connection":     tot_c,
		"scb_logging":                tot_l,
		"scb_gracefully_finishing":   tot_g,
		"scb_idle_cleanup":           tot_i,
		"scb_open_slot":              tot_dot,
		"scb_waiting_for_connection": tot_underscore,
		"scb_total":                  tot_total,
		"hostname":                   hostname,
	}

	return event

}

func ParseMatchFloat(inputString, hostname, fieldName string) float64 {
	var parseString string
	if strings.HasPrefix(inputString, ".") {
		parseString = strings.Replace(inputString, ".", "0.", 1)
	} else {
		parseString = inputString
	}
	outputFloat, er := strconv.ParseFloat(parseString, 64)

	/* Do we need to log failure? */
	if er != nil {
		logp.Warn("Host: %s - cannot parse string %s: %s to float.", hostname, fieldName, inputString)
		return 0.0
	}
	return outputFloat
}
<commit_msg>Rename variables and output fields<commit_after>package status

import (
	"bufio"
	"io"
	"regexp"
	"strconv"
	"strings"

	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
)

// Map body to MapStr
func eventMapping(body io.ReadCloser, hostname string, metricset string) common.MapStr {

	var (
		totalAccesses       int
		totalKBytes         int
		uptime              int
		cpuLoad             float64
		cpuUser             float64
		cpuSystem           float64
		cpuChildrenUser     float64
		cpuChildrenSystem   float64
		reqPerSec           float64
		bytesPerSec         float64
		bytesPerReq         float64
		busyWorkers         int
		idleWorkers         int
		connsTotal          int
		connsAsyncWriting   int
		connsAsyncKeepAlive int
		connsAsyncClosing   int
		serverUptimeSeconds int
		load1               float64
		load5               float64
		load15              float64
		totalS              int
		totalR              int
		totalW              int
		totalK              int
		totalD              int
		totalC              int
		totalL              int
		totalG              int
		totalI              int
		totalDot            int
		totalUnderscore     int
		totalAll            int
	)

	var re *regexp.Regexp

	// Reads file line by line
	scanner := bufio.NewScanner(body)

	// See https://github.com/radoondas/apachebeat/blob/master/collector/status.go#L114
	// Only as POC
	for scanner.Scan() {

		// Total Accesses: 16147
		re = regexp.MustCompile("Total Accesses: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			totalAccesses, _ = strconv.Atoi(matches[1])
		}

		//Total kBytes: 12988
		re = regexp.MustCompile("Total kBytes: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			totalKBytes, _ = strconv.Atoi(matches[1])
		}

		// Uptime: 3229728
		re = regexp.MustCompile("Uptime: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			uptime, _ = strconv.Atoi(matches[1])
		}

		// CPULoad: .000408393
		re = regexp.MustCompile("CPULoad: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpuLoad = ParseMatchFloat(matches[1], hostname, "cpuLoad")
		}

		// CPUUser: 0
		re = regexp.MustCompile("CPUUser: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpuUser = ParseMatchFloat(matches[1], hostname, "cpuUser")
		}

		// CPUSystem: .01
		re = regexp.MustCompile("CPUSystem: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpuSystem = ParseMatchFloat(matches[1], hostname, "cpuSystem")
		}

		// CPUChildrenUser: 0
		re = regexp.MustCompile("CPUChildrenUser: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpuChildrenUser = ParseMatchFloat(matches[1], hostname, "cpuChildrenUser")
		}

		// CPUChildrenSystem: 0
		re = regexp.MustCompile("CPUChildrenSystem: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			cpuChildrenSystem = ParseMatchFloat(matches[1], hostname, "cpuChildrenSystem")
		}

		// ReqPerSec: .00499949
		re = regexp.MustCompile("ReqPerSec: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			reqPerSec = ParseMatchFloat(matches[1], hostname, "reqPerSec")
		}

		// BytesPerSec: 4.1179
		re = regexp.MustCompile("BytesPerSec: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			bytesPerSec = ParseMatchFloat(matches[1], hostname, "bytesPerSec")
		}

		// BytesPerReq: 823.665
		re = regexp.MustCompile("BytesPerReq: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			bytesPerReq = ParseMatchFloat(matches[1], hostname, "bytesPerReq")
		}

		// BusyWorkers: 1
		re = regexp.MustCompile("BusyWorkers: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			busyWorkers, _ = strconv.Atoi(matches[1])
		}

		// IdleWorkers: 8
		re = regexp.MustCompile("IdleWorkers: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			idleWorkers, _ = strconv.Atoi(matches[1])
		}

		// ConnsTotal: 4940
		re = regexp.MustCompile("ConnsTotal: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			connsTotal, _ = strconv.Atoi(matches[1])
		}

		// ConnsAsyncWriting: 527
		re = regexp.MustCompile("ConnsAsyncWriting: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			connsAsyncWriting, _ = strconv.Atoi(matches[1])
		}

		// ConnsAsyncKeepAlive: 1321
		re = regexp.MustCompile("ConnsAsyncKeepAlive: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			connsAsyncKeepAlive, _ = strconv.Atoi(matches[1])
		}

		// ConnsAsyncClosing: 2785
		re = regexp.MustCompile("ConnsAsyncClosing: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			connsAsyncClosing, _ = strconv.Atoi(matches[1])
		}

		// ServerUptimeSeconds: 43
		re = regexp.MustCompile("ServerUptimeSeconds: (\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			serverUptimeSeconds, _ = strconv.Atoi(matches[1])
		}

		//Load1: 0.01
		re = regexp.MustCompile("Load1: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			load1 = ParseMatchFloat(matches[1], hostname, "load1")
		}

		//Load5: 0.10
		re = regexp.MustCompile("Load5: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			load5 = ParseMatchFloat(matches[1], hostname, "load5")
		}

		//Load15: 0.06
		re = regexp.MustCompile("Load15: (\\d*.*\\d+)")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			load15 = ParseMatchFloat(matches[1], hostname, "load15")
		}

		// Scoreboard Key:
		// "_" Waiting for Connection, "S" Starting up, "R" Reading Request,
		// "W" Sending Reply, "K" Keepalive (read), "D" DNS Lookup,
		// "C" Closing connection, "L" Logging, "G" Gracefully finishing,
		// "I" Idle cleanup of worker, "." Open slot with no current process
		// Scoreboard: _W____........___...............................................................................................................................................................................................................................................
		re = regexp.MustCompile("Scoreboard: (_|S|R|W|K|D|C|L|G|I|\\.)+")
		if matches := re.FindStringSubmatch(scanner.Text()); matches != nil {
			scr := strings.Split(scanner.Text(), " ")

			totalUnderscore = strings.Count(scr[1], "_")
			totalS = strings.Count(scr[1], "S")
			totalR = strings.Count(scr[1], "R")
			totalW = strings.Count(scr[1], "W")
			totalK = strings.Count(scr[1], "K")
			totalD = strings.Count(scr[1], "D")
			totalC = strings.Count(scr[1], "C")
			totalL = strings.Count(scr[1], "L")
			totalG = strings.Count(scr[1], "G")
			totalI = strings.Count(scr[1], "I")
			totalDot = strings.Count(scr[1], ".")
			totalAll = totalUnderscore + totalS + totalR + totalW + totalK + totalD + totalC + totalL + totalG + totalI + totalDot
		}
	}

	event := common.MapStr{
		"totalAccesses":           totalAccesses,
		"totalKBytes":             totalKBytes,
		"uptime":                  uptime,
		"cpuLoad":                 cpuLoad,
		"cpuUser":                 cpuUser,
		"cpuSystem":               cpuSystem,
		"cpuChildrenUser":         cpuChildrenUser,
		"cpuChildrenSystem":       cpuChildrenSystem,
		"reqPerSec":               reqPerSec,
		"bytesPerSec":             bytesPerSec,
		"bytesPerReq":             bytesPerReq,
		"busyWorkers":             busyWorkers,
		"idleWorkers":             idleWorkers,
		"connsTotal":              connsTotal,
		"connsAsyncWriting":       connsAsyncWriting,
		"connsAsyncKeepAlive":     connsAsyncKeepAlive,
		"connsAsyncClosing":       connsAsyncClosing,
		"serverUptimeSeconds":     serverUptimeSeconds,
		"load1":                   load1,
		"load5":                   load5,
		"load15":                  load15,
		"scbStartingUp":           totalS,
		"scbReadingRequest":       totalR,
		"scbSendingReply":         totalW,
		"scbKeepalive":            totalK,
		"scbDnsLookup":            totalD,
		"scbClosingConnection":    totalC,
		"scbLogging":              totalL,
		"scbGracefullyFinishing":  totalG,
		"scbIdleCleanup":          totalI,
		"scbOpenSlot":             totalDot,
		"scbWaitingForConnection": totalUnderscore,
		"scbTotal":                totalAll,
		"hostname":                hostname,
	}

	return event

}

func ParseMatchFloat(inputString, hostname, fieldName string) float64 {
	var parseString string
	if strings.HasPrefix(inputString, ".") {
		parseString = strings.Replace(inputString, ".", "0.", 1)
	} else {
		parseString = inputString
	}
	outputFloat, er := strconv.ParseFloat(parseString, 64)

	/* Do we need to log failure? */
	if er != nil {
		logp.Warn("Host: %s - cannot parse string %s: %s to float.", hostname, fieldName, inputString)
		return 0.0
	}
	return outputFloat
}
<|endoftext|>


<commit_before>package memstream

import (
	"errors"
	"fmt"
	"io"
)

// MemoryStream is a memory-based, automatically resizing stream that can
// easily fill the role of any file-based IO.
type MemoryStream struct {
	buff []byte
	loc  int
}

// DefaultCapacity is the size in bytes of a new MemoryStream's backing buffer
const DefaultCapacity = 512

// New creates a new MemoryStream instance
func New() *MemoryStream {
	return NewCapacity(DefaultCapacity)
}

// NewCapacity starts the returned MemoryStream with the given capacity
func NewCapacity(cap int) *MemoryStream {
	return &MemoryStream{buff: make([]byte, 0, DefaultCapacity), loc: 0}
}

// Seek sets the offset for the next Read or Write to offset, interpreted
// according to whence: 0 means relative to the origin of the file, 1 means
// relative to the current offset, and 2 means relative to the end. Seek
// returns the new offset and an error, if any.
//
// Seeking to a negative offset is an error. Seeking to any positive offset is
// legal. If the location is beyond the end of the current length, the position
// will be placed at length.
func (m *MemoryStream) Seek(offset int64, whence int) (int64, error) {
	newLoc := m.loc
	switch whence {
	case 0:
		newLoc = int(offset)
	case 1:
		newLoc += whence
	case 2:
		newLoc = len(m.buff) - whence
	}

	if newLoc < 0 {
		return int64(m.loc), errors.New("Unable to seek to a location <0")
	}

	if newLoc > len(m.buff) {
		newLoc = len(m.buff)
	}

	m.loc = newLoc

	return int64(m.loc), nil
}

// Read puts up to len(p) bytes into p. Will return the number of bytes read.
func (m *MemoryStream) Read(p []byte) (n int, err error) {
	n = copy(p, m.buff[m.loc:len(m.buff)])
	m.loc += n

	if m.loc == len(m.buff) {
		return n, io.EOF
	}

	return n, nil
}

// Write writes the given bytes into the memory stream. If needed, the underlying
// buffer will be expanded to fit the new bytes.
func (m *MemoryStream) Write(p []byte) (n int, err error) {
	// Do we have space?
	if available := cap(m.buff) - m.loc; available < len(p) {
		// How much should we expand by?
		addCap := cap(m.buff)
		if addCap < len(p) {
			addCap = len(p)
		}

		newBuff := make([]byte, len(m.buff), cap(m.buff)+addCap)

		copy(newBuff, m.buff)

		fmt.Printf("Expanded to %v bytes from %v\n", cap(newBuff), cap(m.buff))

		m.buff = newBuff
	}

	// Write
	n = copy(m.buff[m.loc:cap(m.buff)], p)
	m.loc += n
	if len(m.buff) < m.loc {
		m.buff = m.buff[:m.loc]
	}

	return n, nil
}
<commit_msg>Package comment<commit_after>// Memstream is an expandable ReadWriteSeeker for Golang that works with an
// internally managed byte buffer. Operation is usage is intended to be seamless
// and smooth.
//
// In situations where the maximum read/write sizes are known, a fixed
// []byte/byte buffer will likely offer better performance.
package memstream

import (
	"errors"
	"fmt"
	"io"
)

// MemoryStream is a memory-based, automatically resizing stream that can
// easily fill the role of any file-based IO.
type MemoryStream struct {
	buff []byte
	loc  int
}

// DefaultCapacity is the size in bytes of a new MemoryStream's backing buffer
const DefaultCapacity = 512

// New creates a new MemoryStream instance
func New() *MemoryStream {
	return NewCapacity(DefaultCapacity)
}

// NewCapacity starts the returned MemoryStream with the given capacity
func NewCapacity(cap int) *MemoryStream {
	return &MemoryStream{buff: make([]byte, 0, DefaultCapacity), loc: 0}
}

// Seek sets the offset for the next Read or Write to offset, interpreted
// according to whence: 0 means relative to the origin of the file, 1 means
// relative to the current offset, and 2 means relative to the end. Seek
// returns the new offset and an error, if any.
//
// Seeking to a negative offset is an error. Seeking to any positive offset is
// legal. If the location is beyond the end of the current length, the position
// will be placed at length.
func (m *MemoryStream) Seek(offset int64, whence int) (int64, error) {
	newLoc := m.loc
	switch whence {
	case 0:
		newLoc = int(offset)
	case 1:
		newLoc += whence
	case 2:
		newLoc = len(m.buff) - whence
	}

	if newLoc < 0 {
		return int64(m.loc), errors.New("Unable to seek to a location <0")
	}

	if newLoc > len(m.buff) {
		newLoc = len(m.buff)
	}

	m.loc = newLoc

	return int64(m.loc), nil
}

// Read puts up to len(p) bytes into p. Will return the number of bytes read.
func (m *MemoryStream) Read(p []byte) (n int, err error) {
	n = copy(p, m.buff[m.loc:len(m.buff)])
	m.loc += n

	if m.loc == len(m.buff) {
		return n, io.EOF
	}

	return n, nil
}

// Write writes the given bytes into the memory stream. If needed, the underlying
// buffer will be expanded to fit the new bytes.
func (m *MemoryStream) Write(p []byte) (n int, err error) {
	// Do we have space?
	if available := cap(m.buff) - m.loc; available < len(p) {
		// How much should we expand by?
		addCap := cap(m.buff)
		if addCap < len(p) {
			addCap = len(p)
		}

		newBuff := make([]byte, len(m.buff), cap(m.buff)+addCap)

		copy(newBuff, m.buff)

		fmt.Printf("Expanded to %v bytes from %v\n", cap(newBuff), cap(m.buff))

		m.buff = newBuff
	}

	// Write
	n = copy(m.buff[m.loc:cap(m.buff)], p)
	m.loc += n
	if len(m.buff) < m.loc {
		m.buff = m.buff[:m.loc]
	}

	return n, nil
}
<|endoftext|>


<commit_before>// Package wikimg can pull the latest image URLs from Wikimedia Commons
// https://commons.wikimedia.org and map image colors to an xterm256 color
// palette.
package wikimg

import (
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/color"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"

	// We define which image formats we support by importing
	// decoder packages
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
)

var (
	// EndOfResults is returned by Next() when no more results are available
	EndOfResults = errors.New("end of results")

	// Canceled may be returned by Next() and FirstColor() when the client
	// closes the Cancel channel on a Puller
	Canceled = errors.New("wikimg: canceled image processing")
)

const (
	// queryURL is the API we are querying
	queryURL = "https://commons.wikimedia.org/w/api.php"

	// apiMax is the max results we can request from the API at one time
	apiMax = 500

	// cancelCheckpoint is the number of pixels between checking
	// whether the request was canceled when running FirstColor()
	cancelCheckpoint = 10000
)

// queryResp mirrors the JSON structure returned by queryURL, specifying only
// the info we're interested in.
type queryResp struct {

	// Continue contains strings we need to pass back into the API to
	// continue where we left off
	Continue struct {
		Continue   string
		AIContinue string
	}

	// Query contains the actual results
	Query struct {
		AllImages []struct {
			URL string
		}
	}
}

// Puller is an image puller that retrieves the most recent image URLs that
// have been uploaded to Wikimedia Commons https://commons.wikimedia.org
type Puller struct {
	// qr is the most recent response from the API
	qr *queryResp

	// i is the current index into qr.Query.AllImages
	i int

	// count is the total number of images we've collected
	count int

	// max is the maximum number of images we want to collect
	max int

	// Cancel is an optional channel. Setting this value on Puller
	// and closing the channel signals to the Puller that any
	// in process operations (i.e, retrieving an image or computing
	// its first color) should be canceled. Any future
	// calls to Next() or FirstColor() will return a Canceled
	// error.
	Cancel <-chan struct{}
}

// NewPuller creates a puller that can return at most max images
// when calls to Next() are made
func NewPuller(max int) *Puller {
	return &Puller{
		max: max,
	}
}

// Next returns the next most recent image URL. If no more results are
// available EndOfResults is returned as an error.
func (p *Puller) Next() (string, error) {
	// If we've exceeded that max we want to get, then stop
	if p.count >= p.max {
		return "", EndOfResults
	}

	// Ensure we haven't been canceled yet
	select {
	case <-p.Cancel:
		// If p.Cancel has been closed, this will be triggered
		return "", Canceled

	default:
		// Otherwise we'll just do nothing immediately
	}

	// If we're within the length of our current request,
	// return right away and increment our counters
	if p.qr != nil && p.i < len(p.qr.Query.AllImages) {
		img := p.qr.Query.AllImages[p.i].URL
		p.i++
		p.count++
		return img, nil
	}

	// Otherwise, we need to create a new request. Recreate our request params
	// and reset per-request counter.
	p.i = 0
	params := url.Values{}
	params.Set("action", "query")
	params.Set("format", "json")
	params.Set("list", "allimages")
	params.Set("aidir", "descending")
	params.Set("aisort", "timestamp")

	// 500 is the most allowed by the API per request, but we may want
	// less.
	if p.count+apiMax > p.max {
		params.Set("ailimit", strconv.Itoa(p.max-p.count))
	} else {
		params.Set("ailimit", strconv.Itoa(p.max))
	}

	// If we have a previous request with continue values, use them
	if p.qr != nil &&
		len(p.qr.Continue.Continue) > 0 &&
		len(p.qr.Continue.AIContinue) > 0 {
		params.Set("continue", p.qr.Continue.Continue)
		params.Set("aicontinue", p.qr.Continue.AIContinue)
	}

	// Call the wikimedia API
	resp, err := http.Get(queryURL + "?" + params.Encode())
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Read the contents of the response as bytes
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	// Parse the bytes into a struct
	p.qr = &queryResp{}
	err = json.Unmarshal(b, p.qr)
	if err != nil {
		return "", err
	}

	// If there's no more images, then return
	if len(p.qr.Query.AllImages) < 1 {
		return "", EndOfResults
	}

	// Return first value of the new request
	p.count++
	return p.qr.Query.AllImages[p.i].URL, nil
}

// FirstColor tries to return the first non-gray color in the image. A gray
// color is one that, when mapped to an xterm256 palette, has the same value
// for red, green and blue. Ok, we've defined "gray". What does "first" mean?
// We iterate through pixels starting with 0,0 (top left) and move to the
// bottom right one by one. In the worst case (a grayscale image), we iterate
// through every pixel, give up, and return the final pixel color even though
// it's gray. Both the xtermColor (an integer between 0-255) and a hex
// string (e.g., "#bb00cc") is returned.
func (p *Puller) FirstColor(imgURL string) (xtermColor int, hex string, err error) {
	// Create a request so we can use req.Cancel
	req, err := http.NewRequest("GET", imgURL, nil)
	if err != nil {
		return
	}

	// Set up cancellation pipeline, link request to puller
	req.Cancel = p.Cancel

	// Call the image server
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	// Decode into an object
	img, _, err := image.Decode(resp.Body)
	if err != nil {
		return
	}

	// Use our XTerm256 as a color.Palette so we can map the colors of the
	// image to our palette.
	pal := color.Palette(XTerm256)

	// Iterate through every pixel and try to find a color. If we don't
	// find a color (i.e., the image is grayscale) we'll default to the last
	// pixel in the image.
	rect := img.Bounds()
	i := 0
	for x := 0; x < rect.Dx(); x++ {
		for y := 0; y < rect.Dy(); y++ {

			// Check if p.Cancel has been closed once every cancelCheckpoint
			// iterations
			if i%cancelCheckpoint == 0 {
				select {

				case <-p.Cancel:
					// If p.Cancel has been closed, this will be triggered
					err = Canceled
					return

				default:
					// Otherwise we'll just do nothing immediately
				}
			}

			// xtermColor is the index in the palette which this
			// actual color maps to. It is also (by design) the
			// xterm256 value that maps to this color.
			xtermColor = pal.Index(img.At(x, y))

			// Get the color.RGBA value for this color. Not great to do a type
			// assertion here but easiest way to get 8-bit values without bit
			// fiddling.
			rgba, ok := pal[xtermColor].(color.RGBA)
			if !ok {
				err = errors.New("can't assert to color.RGBA")
				return
			}

			// Compute the hex value of the color
			hex = fmt.Sprintf("#%02x%02x%02x", rgba.R, rgba.G, rgba.B)

			// If any of the RGB values differ, it's a color, so we can
			// stop.
			if !(rgba.R == rgba.G && rgba.G == rgba.B) {
				break
			}
		}
	}

	return
}
<commit_msg>doc format fixes<commit_after>// Package wikimg can pull the latest image URLs from Wikimedia Commons
// https://commons.wikimedia.org and map image colors to an xterm256 color
// palette.
package wikimg

import (
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/color"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"

	// We define which image formats we support by importing decoder packages
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
)

var (
	// EndOfResults is returned by Next() when no more results are available
	EndOfResults = errors.New("end of results")

	// Canceled may be returned by Next() and FirstColor() when the client
	// closes the Cancel channel on a Puller
	Canceled = errors.New("wikimg: canceled image processing")
)

const (
	// queryURL is the API we are querying
	queryURL = "https://commons.wikimedia.org/w/api.php"

	// apiMax is the max results we can request from the API at one time
	apiMax = 500

	// cancelCheckpoint is the number of pixels between checking whether the
	// request was canceled when running FirstColor()
	cancelCheckpoint = 10000
)

// queryResp mirrors the JSON structure returned by queryURL, specifying only
// the info we're interested in.
type queryResp struct {

	// Continue contains strings we need to pass back into the API to
	// continue where we left off
	Continue struct {
		Continue   string
		AIContinue string
	}

	// Query contains the actual results
	Query struct {
		AllImages []struct {
			URL string
		}
	}
}

// Puller is an image puller that retrieves the most recent image URLs that
// have been uploaded to Wikimedia Commons https://commons.wikimedia.org
type Puller struct {
	// qr is the most recent response from the API
	qr *queryResp

	// i is the current index into qr.Query.AllImages
	i int

	// count is the total number of images we've collected
	count int

	// max is the maximum number of images we want to collect
	max int

	// Cancel is an optional channel. Setting this value on Puller
	// and closing the channel signals to the Puller that any
	// in process operations (i.e, retrieving an image or computing
	// its first color) should be canceled. Any future
	// calls to Next() or FirstColor() will return a Canceled
	// error.
	Cancel <-chan struct{}
}

// NewPuller creates a puller that can return at most max images when calls to
// Next() are made
func NewPuller(max int) *Puller {
	return &Puller{
		max: max,
	}
}

// Next returns the next most recent image URL. If no more results are
// available EndOfResults is returned as an error.
func (p *Puller) Next() (string, error) {
	// If we've exceeded that max we want to get, then stop
	if p.count >= p.max {
		return "", EndOfResults
	}

	// Ensure we haven't been canceled yet
	select {
	case <-p.Cancel:
		// If p.Cancel has been closed, this will be triggered
		return "", Canceled

	default:
		// Otherwise we'll just do nothing immediately
	}

	// If we're within the length of our current request,
	// return right away and increment our counters
	if p.qr != nil && p.i < len(p.qr.Query.AllImages) {
		img := p.qr.Query.AllImages[p.i].URL
		p.i++
		p.count++
		return img, nil
	}

	// Otherwise, we need to create a new request. Recreate our request params
	// and reset per-request counter.
	p.i = 0
	params := url.Values{}
	params.Set("action", "query")
	params.Set("format", "json")
	params.Set("list", "allimages")
	params.Set("aidir", "descending")
	params.Set("aisort", "timestamp")

	// 500 is the most allowed by the API per request, but we may want less.
	if p.count+apiMax > p.max {
		params.Set("ailimit", strconv.Itoa(p.max-p.count))
	} else {
		params.Set("ailimit", strconv.Itoa(p.max))
	}

	// If we have a previous request with continue values, use them
	if p.qr != nil &&
		len(p.qr.Continue.Continue) > 0 &&
		len(p.qr.Continue.AIContinue) > 0 {
		params.Set("continue", p.qr.Continue.Continue)
		params.Set("aicontinue", p.qr.Continue.AIContinue)
	}

	// Call the wikimedia API
	resp, err := http.Get(queryURL + "?" + params.Encode())
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Read the contents of the response as bytes
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	// Parse the bytes into a struct
	p.qr = &queryResp{}
	err = json.Unmarshal(b, p.qr)
	if err != nil {
		return "", err
	}

	// If there's no more images, then return
	if len(p.qr.Query.AllImages) < 1 {
		return "", EndOfResults
	}

	// Return first value of the new request
	p.count++
	return p.qr.Query.AllImages[p.i].URL, nil
}

// FirstColor tries to return the first non-gray color in the image. A gray
// color is one that, when mapped to an xterm256 palette, has the same value
// for red, green and blue. Ok, we've defined "gray". What does "first" mean?
// We iterate through pixels starting with 0,0 (top left) and move to the
// bottom right one by one. In the worst case (a grayscale image), we iterate
// through every pixel, give up, and return the final pixel color even though
// it's gray. Both the xtermColor (an integer between 0-255) and a hex string
// (e.g., "#bb00cc") is returned.
func (p *Puller) FirstColor(imgURL string) (xtermColor int, hex string, err error) {
	// Create a request so we can use req.Cancel
	req, err := http.NewRequest("GET", imgURL, nil)
	if err != nil {
		return
	}

	// Set up cancellation pipeline, link request to puller
	req.Cancel = p.Cancel

	// Call the image server
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	// Decode into an object
	img, _, err := image.Decode(resp.Body)
	if err != nil {
		return
	}

	// Use our XTerm256 as a color.Palette so we can map the colors of the
	// image to our palette.
	pal := color.Palette(XTerm256)

	// Iterate through every pixel and try to find a color. If we don't find a
	// color (i.e., the image is grayscale) we'll default to the last pixel in
	// the image.
	rect := img.Bounds()
	i := 0
	for x := 0; x < rect.Dx(); x++ {
		for y := 0; y < rect.Dy(); y++ {

			// Check if p.Cancel has been closed once every cancelCheckpoint
			// iterations
			if i%cancelCheckpoint == 0 {
				select {

				case <-p.Cancel:
					// If p.Cancel has been closed, this will be triggered
					err = Canceled
					return

				default:
					// Otherwise we'll just do nothing immediately
				}
			}

			// xtermColor is the index in the palette which this
			// actual color maps to. It is also (by design) the
			// xterm256 value that maps to this color.
			xtermColor = pal.Index(img.At(x, y))

			// Get the color.RGBA value for this color. Not great to do a type
			// assertion here but easiest way to get 8-bit values without bit
			// fiddling.
			rgba, ok := pal[xtermColor].(color.RGBA)
			if !ok {
				err = errors.New("can't assert to color.RGBA")
				return
			}

			// Compute the hex value of the color
			hex = fmt.Sprintf("#%02x%02x%02x", rgba.R, rgba.G, rgba.B)

			// If any of the RGB values differ, it's a color, so we can stop.
			if !(rgba.R == rgba.G && rgba.G == rgba.B) {
				break
			}
		}
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2017 Frédéric Guillot. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package worker // import "miniflux.app/worker"

import (
	"miniflux.app/logger"
	"miniflux.app/model"
	"miniflux.app/reader/feed"
)

// Worker refreshes a feed in the background.
type Worker struct {
	id          int
	feedHandler *feed.Handler
}

// Run wait for a job and refresh the given feed.
func (w *Worker) Run(c chan model.Job) {
	logger.Info("[Worker] #%d started", w.id)

	for {
		job := <-c
		logger.Debug("[Worker #%d] got userID=%d, feedID=%d", w.id, job.UserID, job.FeedID)

		err := w.feedHandler.RefreshFeed(job.UserID, job.FeedID)
		if err != nil {
			logger.Error("[Worker] %v", err)
		}
	}
}
<commit_msg>Change log level to debug when starting workers<commit_after>// Copyright 2017 Frédéric Guillot. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package worker // import "miniflux.app/worker"

import (
	"miniflux.app/logger"
	"miniflux.app/model"
	"miniflux.app/reader/feed"
)

// Worker refreshes a feed in the background.
type Worker struct {
	id          int
	feedHandler *feed.Handler
}

// Run wait for a job and refresh the given feed.
func (w *Worker) Run(c chan model.Job) {
	logger.Debug("[Worker] #%d started", w.id)

	for {
		job := <-c
		logger.Debug("[Worker #%d] got userID=%d, feedID=%d", w.id, job.UserID, job.FeedID)

		err := w.feedHandler.RefreshFeed(job.UserID, job.FeedID)
		if err != nil {
			logger.Error("[Worker] %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package main_test

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"time"

	. "github.com/alphagov/govuk_crawler_worker"

	. "github.com/alphagov/govuk_crawler_worker/http_crawler"
	. "github.com/alphagov/govuk_crawler_worker/queue"
	. "github.com/alphagov/govuk_crawler_worker/ttl_hash_set"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"os"

	"github.com/alphagov/govuk_crawler_worker/util"
	"github.com/fzzy/radix/redis"
	"github.com/streadway/amqp"
)

var _ = Describe("Workflow", func() {
	Describe("Acknowledging items", func() {
		amqpAddr := util.GetEnvDefault("AMQP_ADDRESS", "amqp://guest:guest@localhost:5672/")
		redisAddr := util.GetEnvDefault("REDIS_ADDRESS", "127.0.0.1:6379")
		exchangeName, queueName := "test-workflow-exchange", "test-workflow-queue"
		prefix := "govuk_mirror_crawler_workflow_test"

		var (
			queueManager    *QueueManager
			queueManagerErr error
			ttlHashSet      *TTLHashSet
			ttlHashSetErr   error
			mirrorRoot      string
		)

		BeforeEach(func() {
			mirrorRoot = os.Getenv("MIRROR_ROOT")

			ttlHashSet, ttlHashSetErr = NewTTLHashSet(prefix, redisAddr)
			Expect(ttlHashSetErr).To(BeNil())

			queueManager, queueManagerErr = NewQueueManager(
				amqpAddr,
				exchangeName,
				queueName)

			Expect(queueManagerErr).To(BeNil())
			Expect(queueManager).ToNot(BeNil())
		})

		AfterEach(func() {
			Expect(ttlHashSet.Close()).To(BeNil())
			Expect(purgeAllKeys(prefix, redisAddr)).To(BeNil())

			deleted, err := queueManager.Consumer.Channel.QueueDelete(queueName, false, false, true)
			Expect(err).To(BeNil())
			Expect(deleted).To(Equal(0))

			err = queueManager.Consumer.Channel.ExchangeDelete(exchangeName, false, true)
			Expect(err).To(BeNil())

			queueManager.Close()
			DeleteMirrorFilesFromDisk(mirrorRoot)
		})

		Describe("AcknowledgeItem", func() {
			It("should read from a channel and add URLs to the hash set", func() {
				url := "https://www.gov.uk/foo"

				exists, err := ttlHashSet.Exists(url)
				Expect(err).To(BeNil())
				Expect(exists).To(BeFalse())

				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())

				outbound := make(chan *CrawlerMessageItem, 1)

				err = queueManager.Publish("#", "text/plain", url)
				Expect(err).To(BeNil())

				for item := range deliveries {
					outbound <- NewCrawlerMessageItem(item, "www.gov.uk", []string{})
					break
				}

				Expect(len(outbound)).To(Equal(1))

				go AcknowledgeItem(outbound, ttlHashSet)
				time.Sleep(time.Millisecond)

				Expect(len(outbound)).To(Equal(0))

				exists, err = ttlHashSet.Exists(url)
				Expect(err).To(BeNil())
				Expect(exists).To(BeTrue())

				// Close the channel to stop the goroutine for AcknowledgeItem.
				close(outbound)
			})
		})

		Describe("CrawlURL", func() {
			var (
				crawler    *Crawler
				crawlerErr error
			)

			BeforeEach(func() {
				crawler, crawlerErr = NewCrawler("http://127.0.0.1/")

				Expect(crawlerErr).To(BeNil())
				Expect(crawler).ToNot(BeNil())
			})

			It("crawls a URL and assigns the body", func() {
				outbound := make(chan *CrawlerMessageItem, 1)

				body := `<a href="foo.com">bar</a>`
				server := testServer(200, body)

				deliveryItem := &amqp.Delivery{Body: []byte(server.URL)}
				outbound <- NewCrawlerMessageItem(*deliveryItem, "127.0.0.1", []string{})

				crawled := CrawlURL(outbound, crawler)

				Expect((<-crawled).HTMLBody[0:25]).To(Equal([]byte(body)))

				server.Close()
				close(outbound)
			})
		})

		Describe("WriteItemToDisk", func() {
			It("wrote something to disk", func() {
				url := "https://www.gov.uk/extract-some-urls"
				deliveryItem := &amqp.Delivery{Body: []byte(url)}
				item := NewCrawlerMessageItem(*deliveryItem, "www.gov.uk", []string{})
				item.HTMLBody = []byte(`<a href="https://www.gov.uk/some-url">a link</a>`)

				outbound := make(chan *CrawlerMessageItem, 1)
				extract := WriteItemToDisk(outbound)

				Expect(len(extract)).To(Equal(0))

				outbound <- item

				Expect(<-extract).To(Equal(item))

				filePath, _ := item.FilePath()
				fileStat, _ := os.Stat(filePath)

				Expect(fileStat).ToNot(BeNil())

				close(outbound)
			})
		})

		Describe("ExtractURLs", func() {
			It("extracts URLs from the HTML body and adds them to a new channel; acknowledging item", func() {
				url := "https://www.gov.uk/extract-some-urls"
				deliveryItem := &amqp.Delivery{Body: []byte(url)}
				item := NewCrawlerMessageItem(*deliveryItem, "www.gov.uk", []string{})
				item.HTMLBody = []byte(`<a href="https://www.gov.uk/some-url">a link</a>`)

				outbound := make(chan *CrawlerMessageItem, 1)
				publish, acknowledge := ExtractURLs(outbound)

				Expect(len(publish)).To(Equal(0))
				Expect(len(acknowledge)).To(Equal(0))

				outbound <- item

				Expect(<-publish).To(Equal("https://www.gov.uk/some-url"))
				Expect(<-acknowledge).To(Equal(item))

				close(outbound)
			})
		})

		Describe("PublishURLs", func() {
			It("doesn't publish URLs that have already been crawled", func() {
				url := "https://www.gov.uk/government/organisations"

				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())
				Expect(len(deliveries)).To(Equal(0))

				_, err = ttlHashSet.Add(url)
				Expect(err).To(BeNil())

				publish := make(chan string, 1)
				outbound := make(chan []byte, 1)

				go func() {
					for item := range deliveries {
						outbound <- item.Body
					}
				}()
				go PublishURLs(ttlHashSet, queueManager, publish)
				time.Sleep(time.Millisecond)

				publish <- url
				time.Sleep(time.Millisecond)

				Expect(len(publish)).To(Equal(0))
				Expect(len(outbound)).To(Equal(0))

				// Close the channel to stop the goroutine for PublishURLs.
				close(publish)
				close(outbound)
			})

			It("publishes URLs that haven't been crawled yet", func() {
				url := "https://www.gov.uk/government/foo"

				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())
				Expect(len(deliveries)).To(Equal(0))

				publish := make(chan string, 1)
				outbound := make(chan []byte, 1)

				go func() {
					for item := range deliveries {
						outbound <- item.Body
					}
				}()
				go PublishURLs(ttlHashSet, queueManager, publish)

				publish <- url

				Expect(<-outbound).To(Equal([]byte(url)))
				Expect(len(publish)).To(Equal(0))

				close(publish)
				close(outbound)
			})
		})

		Describe("ReadFromQueue", func() {
			It("provides a way of converting AMQP bodies to CrawlerMessageItems", func() {
				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())

				outbound := ReadFromQueue(deliveries, ttlHashSet, []string{})
				Expect(len(outbound)).To(Equal(0))

				url := "https://www.foo.com/bar"
				err = queueManager.Publish("#", "text/plain", url)
				Expect(err).To(BeNil())

				item := <-outbound
				Expect(string(item.Body)).To(Equal(url))

				close(outbound)
			})
		})
	})
})

func purgeAllKeys(prefix string, address string) error {
	client, err := redis.Dial("tcp", address)
	if err != nil {
		return err
	}

	keys, err := client.Cmd("KEYS", prefix+"*").List()
	if err != nil || len(keys) <= 0 {
		return err
	}

	reply := client.Cmd("DEL", keys)
	if reply.Err != nil {
		return reply.Err
	}

	return nil
}

func testServer(status int, body string) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(status)
		fmt.Fprintln(w, body)
	}))
}
<commit_msg>Test the content that is written to disk<commit_after>package main_test

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"time"

	. "github.com/alphagov/govuk_crawler_worker"

	. "github.com/alphagov/govuk_crawler_worker/http_crawler"
	. "github.com/alphagov/govuk_crawler_worker/queue"
	. "github.com/alphagov/govuk_crawler_worker/ttl_hash_set"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"os"

	"github.com/alphagov/govuk_crawler_worker/util"
	"github.com/fzzy/radix/redis"
	"github.com/streadway/amqp"
)

var _ = Describe("Workflow", func() {
	Describe("Acknowledging items", func() {
		amqpAddr := util.GetEnvDefault("AMQP_ADDRESS", "amqp://guest:guest@localhost:5672/")
		redisAddr := util.GetEnvDefault("REDIS_ADDRESS", "127.0.0.1:6379")
		exchangeName, queueName := "test-workflow-exchange", "test-workflow-queue"
		prefix := "govuk_mirror_crawler_workflow_test"

		var (
			queueManager    *QueueManager
			queueManagerErr error
			ttlHashSet      *TTLHashSet
			ttlHashSetErr   error
			mirrorRoot      string
		)

		BeforeEach(func() {
			mirrorRoot = os.Getenv("MIRROR_ROOT")

			ttlHashSet, ttlHashSetErr = NewTTLHashSet(prefix, redisAddr)
			Expect(ttlHashSetErr).To(BeNil())

			queueManager, queueManagerErr = NewQueueManager(
				amqpAddr,
				exchangeName,
				queueName)

			Expect(queueManagerErr).To(BeNil())
			Expect(queueManager).ToNot(BeNil())
		})

		AfterEach(func() {
			Expect(ttlHashSet.Close()).To(BeNil())
			Expect(purgeAllKeys(prefix, redisAddr)).To(BeNil())

			deleted, err := queueManager.Consumer.Channel.QueueDelete(queueName, false, false, true)
			Expect(err).To(BeNil())
			Expect(deleted).To(Equal(0))

			err = queueManager.Consumer.Channel.ExchangeDelete(exchangeName, false, true)
			Expect(err).To(BeNil())

			queueManager.Close()
			DeleteMirrorFilesFromDisk(mirrorRoot)
		})

		Describe("AcknowledgeItem", func() {
			It("should read from a channel and add URLs to the hash set", func() {
				url := "https://www.gov.uk/foo"

				exists, err := ttlHashSet.Exists(url)
				Expect(err).To(BeNil())
				Expect(exists).To(BeFalse())

				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())

				outbound := make(chan *CrawlerMessageItem, 1)

				err = queueManager.Publish("#", "text/plain", url)
				Expect(err).To(BeNil())

				for item := range deliveries {
					outbound <- NewCrawlerMessageItem(item, "www.gov.uk", []string{})
					break
				}

				Expect(len(outbound)).To(Equal(1))

				go AcknowledgeItem(outbound, ttlHashSet)
				time.Sleep(time.Millisecond)

				Expect(len(outbound)).To(Equal(0))

				exists, err = ttlHashSet.Exists(url)
				Expect(err).To(BeNil())
				Expect(exists).To(BeTrue())

				// Close the channel to stop the goroutine for AcknowledgeItem.
				close(outbound)
			})
		})

		Describe("CrawlURL", func() {
			var (
				crawler    *Crawler
				crawlerErr error
			)

			BeforeEach(func() {
				crawler, crawlerErr = NewCrawler("http://127.0.0.1/")

				Expect(crawlerErr).To(BeNil())
				Expect(crawler).ToNot(BeNil())
			})

			It("crawls a URL and assigns the body", func() {
				outbound := make(chan *CrawlerMessageItem, 1)

				body := `<a href="foo.com">bar</a>`
				server := testServer(200, body)

				deliveryItem := &amqp.Delivery{Body: []byte(server.URL)}
				outbound <- NewCrawlerMessageItem(*deliveryItem, "127.0.0.1", []string{})

				crawled := CrawlURL(outbound, crawler)

				Expect((<-crawled).HTMLBody[0:25]).To(Equal([]byte(body)))

				server.Close()
				close(outbound)
			})
		})

		Describe("WriteItemToDisk", func() {
			It("wrote the item to disk", func() {
				url := "https://www.gov.uk/extract-some-urls"
				deliveryItem := &amqp.Delivery{Body: []byte(url)}
				item := NewCrawlerMessageItem(*deliveryItem, "www.gov.uk", []string{})
				item.HTMLBody = []byte(`<a href="https://www.gov.uk/some-url">a link</a>`)

				outbound := make(chan *CrawlerMessageItem, 1)
				extract := WriteItemToDisk(outbound)

				Expect(len(extract)).To(Equal(0))

				outbound <- item

				Expect(<-extract).To(Equal(item))

				filePath, _ := item.FilePath()
				fileContent, err := ioutil.ReadFile(filePath)

				Expect(err).To(BeNil())
				Expect(fileContent).To(Equal(item.HTMLBody))

				close(outbound)
			})
		})

		Describe("ExtractURLs", func() {
			It("extracts URLs from the HTML body and adds them to a new channel; acknowledging item", func() {
				url := "https://www.gov.uk/extract-some-urls"
				deliveryItem := &amqp.Delivery{Body: []byte(url)}
				item := NewCrawlerMessageItem(*deliveryItem, "www.gov.uk", []string{})
				item.HTMLBody = []byte(`<a href="https://www.gov.uk/some-url">a link</a>`)

				outbound := make(chan *CrawlerMessageItem, 1)
				publish, acknowledge := ExtractURLs(outbound)

				Expect(len(publish)).To(Equal(0))
				Expect(len(acknowledge)).To(Equal(0))

				outbound <- item

				Expect(<-publish).To(Equal("https://www.gov.uk/some-url"))
				Expect(<-acknowledge).To(Equal(item))

				close(outbound)
			})
		})

		Describe("PublishURLs", func() {
			It("doesn't publish URLs that have already been crawled", func() {
				url := "https://www.gov.uk/government/organisations"

				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())
				Expect(len(deliveries)).To(Equal(0))

				_, err = ttlHashSet.Add(url)
				Expect(err).To(BeNil())

				publish := make(chan string, 1)
				outbound := make(chan []byte, 1)

				go func() {
					for item := range deliveries {
						outbound <- item.Body
					}
				}()
				go PublishURLs(ttlHashSet, queueManager, publish)
				time.Sleep(time.Millisecond)

				publish <- url
				time.Sleep(time.Millisecond)

				Expect(len(publish)).To(Equal(0))
				Expect(len(outbound)).To(Equal(0))

				// Close the channel to stop the goroutine for PublishURLs.
				close(publish)
				close(outbound)
			})

			It("publishes URLs that haven't been crawled yet", func() {
				url := "https://www.gov.uk/government/foo"

				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())
				Expect(len(deliveries)).To(Equal(0))

				publish := make(chan string, 1)
				outbound := make(chan []byte, 1)

				go func() {
					for item := range deliveries {
						outbound <- item.Body
					}
				}()
				go PublishURLs(ttlHashSet, queueManager, publish)

				publish <- url

				Expect(<-outbound).To(Equal([]byte(url)))
				Expect(len(publish)).To(Equal(0))

				close(publish)
				close(outbound)
			})
		})

		Describe("ReadFromQueue", func() {
			It("provides a way of converting AMQP bodies to CrawlerMessageItems", func() {
				deliveries, err := queueManager.Consume()
				Expect(err).To(BeNil())

				outbound := ReadFromQueue(deliveries, ttlHashSet, []string{})
				Expect(len(outbound)).To(Equal(0))

				url := "https://www.foo.com/bar"
				err = queueManager.Publish("#", "text/plain", url)
				Expect(err).To(BeNil())

				item := <-outbound
				Expect(string(item.Body)).To(Equal(url))

				close(outbound)
			})
		})
	})
})

func purgeAllKeys(prefix string, address string) error {
	client, err := redis.Dial("tcp", address)
	if err != nil {
		return err
	}

	keys, err := client.Cmd("KEYS", prefix+"*").List()
	if err != nil || len(keys) <= 0 {
		return err
	}

	reply := client.Cmd("DEL", keys)
	if reply.Err != nil {
		return reply.Err
	}

	return nil
}

func testServer(status int, body string) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(status)
		fmt.Fprintln(w, body)
	}))
}
<|endoftext|>


<commit_before>package fetch

import (
	"fmt"
	"os"
	"testing"
)

func TestGetMaxWidth(t *testing.T) {
	os.Setenv("VIP_MAX_WIDTH", "500")
	if width := getMaxWidth(); width != 500 {
		t.Fail()
	}

	os.Setenv("VIP_MAX_WIDTH", "1024")
	if width := getMaxWidth(); width != 1024 {
		t.Fail()
	}

	// Test default value
	os.Setenv("VIP_MAX_WIDTH", "")
	if width := getMaxWidth(); width != 720 {
		t.Fail()
	}
}

func TestNeedsRotation(t *testing.T) {
	for i := 1; i <= 8; i++ {
		filename := fmt.Sprintf("f%d-exif.jpg", i)
		f, err := os.Open(fmt.Sprintf("../test/%s", filename))
		if err != nil {
			t.Errorf("Could not open %s.", filename)
		}

		rotate, angle := needsRotation(f)

		fmt.Println(rotate, angle)

		switch i {
		case 6:
			if angle != 90 && rotate != true {
				t.Errorf("Expected true, 90; got %d, %t", rotate, angle)
			}
		case 3:
			if angle != 180 && rotate != true {
				t.Errorf("Expected true, 180; got %d, %t", rotate, angle)
			}
		case 8:
			if angle != 270 && rotate != true {
				t.Errorf("Expected true, 270; got %d, %t", rotate, angle)
			}
		default:
			if angle != 0 && rotate != false {
				t.Errorf("Expected false, 0; got %d, %t", rotate, angle)
			}
		}
	}

}
<commit_msg>Test proves both files can be tested w/o error<commit_after>package fetch

import (
	"fmt"
	"os"
	"testing"
)

func TestGetMaxWidth(t *testing.T) {
	os.Setenv("VIP_MAX_WIDTH", "500")
	if width := getMaxWidth(); width != 500 {
		t.Fail()
	}

	os.Setenv("VIP_MAX_WIDTH", "1024")
	if width := getMaxWidth(); width != 1024 {
		t.Fail()
	}

	// Test default value
	os.Setenv("VIP_MAX_WIDTH", "")
	if width := getMaxWidth(); width != 720 {
		t.Fail()
	}
}

func TestNeedsRotation(t *testing.T) {
	for i := 1; i <= 8; i++ {
		filename := fmt.Sprintf("f%d-exif.jpg", i)
		f, err := os.Open(fmt.Sprintf("../test/%s", filename))
		if err != nil {
			t.Errorf("Could not open %s.", filename)
		}

		rotate, angle := needsRotation(f)

		fmt.Println(rotate, angle)

		switch i {
		case 6:
			if angle != 90 && rotate != true {
				t.Errorf("Expected true, 90; got %d, %t", rotate, angle)
			}
		case 3:
			if angle != 180 && rotate != true {
				t.Errorf("Expected true, 180; got %d, %t", rotate, angle)
			}
		case 8:
			if angle != 270 && rotate != true {
				t.Errorf("Expected true, 270; got %d, %t", rotate, angle)
			}
		default:
			if angle != 0 && rotate != false {
				t.Errorf("Expected false, 0; got %d, %t", rotate, angle)
			}
		}
	}
}

func TestNeedsRotationAltFiles(t *testing.T) {
	filenames := map[int]string{
		1: "awesome.jpeg",
		2: "exif_test_img.jpg",
	}

	for key, filename := range filenames {
		f, err := os.Open(fmt.Sprintf("../test/%s", filename))
		if err != nil {
			t.Errorf("Could not open %s.", filename)
		}

		rotate, angle := needsRotation(f)

		fmt.Println(rotate, angle)

		switch key {
		case 1:
			if angle != 0 && rotate != false {
				t.Errorf("Expected true, 90; got %d, %t", rotate, angle)
			}
		case 2:
			if angle != 90 && rotate != true {
				t.Errorf("Expected true, 90; got %d, %t", rotate, angle)
			}
		}
	}
}
<|endoftext|>


<commit_before>package gogopro

import (
	"io/ioutil"
)

type Power struct {
	APIRequester *APIRequester
}

func (p *Power) Init() *Power {
	return p
}

func CreatePower(APIRequester *APIRequester) *Power {
	power := &Power{}
	power.APIRequester = APIRequester
	return power
}

func (p *Power) GetPowerStatus() (string, error) {
	resp, err := p.APIRequester.get("/bacpac/se")
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if body[len(body)-1] == 0 {
		return "off", nil
	}
	return "on", nil
}
<commit_msg>Move status command to generic object<commit_after>package gogopro

import ()

type Power struct {
	APIRequester   *APIRequester
	StatusCommands map[string]StatusCommand
}

func (p *Power) Init() *Power {
	return p
}

func CreatePower(APIRequester *APIRequester) *Power {
	power := &Power{}
	power.APIRequester = APIRequester
	statusCommands := CreateStatusCommands()
	power.StatusCommands = statusCommands
	return power
}

func CreateStatusCommands() map[string]StatusCommand {
	sc := make(map[string]StatusCommand)
	sc["power"] = StatusCommand{Endpoint: "/bacpac/se", ResultByte: -1,
		Translaters: []StatusTranslater{
			StatusTranslater{
				Result:         0,
				ExpectedReturn: "off"},
			StatusTranslater{
				Result:         1,
				ExpectedReturn: "on"}}}
	return sc
}

func (p *Power) GetPowerStatus() (string, error) {
	result, err := p.StatusCommands["power"].RunStatusCommand(p.APIRequester)
	if err != nil {
		return "", err
	}
	return result, nil
}
<|endoftext|>


<commit_before>package indexer

import (
	"bytes"
	"encoding/json"
	"io"
	"math"
	"sort"
	"strings"

	"github.com/dchest/stemmer/porter2"

	"github.com/dchest/kkr/search/indexer/tokenizer"
)

type Index struct {
	Docs       []*Document              `json:"docs"`
	Words      map[string][]interface{} `json:"words"`
	wordsToDoc map[string]map[*Document]int

	ContentWordWeight      float64 `json:"-"`
	HTMLTitleWeight        float64 `json:"-"`
	HTMLURLComponentWeight float64 `json:"-"`
}

type Document struct {
	URL   string `json:"u"`
	Title string `json:"t"`

	numWords  int `json:"-"`
	selfIndex int `json:"-"`
}

func New() *Index {
	return &Index{
		Docs:                   make([]*Document, 0),
		Words:                  make(map[string][]interface{}),
		wordsToDoc:             make(map[string]map[*Document]int), // word => doc => weight
		ContentWordWeight:      1,
		HTMLTitleWeight:        3,
		HTMLURLComponentWeight: 3,
	}
}

func (n *Index) WriteJSON(w io.Writer) error {
	// Calculate minWeight (for normalization later),
	// and add word frequencies to doc.
	minWeight := math.MaxInt
	for _, m := range n.wordsToDoc {
		for doc, weight := range m {
			if weight < minWeight {
				minWeight = weight
			}
			doc.numWords += 1
		}
	}

	// Make JSON smaller:
	// Sort docs by most frequent, so that smaller
	// doc ids are taken for most frequent ones.
	sort.Slice(n.Docs, func(i, j int) bool {
		return n.Docs[i].numWords > n.Docs[j].numWords
	})

	// Add indexes to docs
	for i, d := range n.Docs {
		d.selfIndex = i
	}

	minWeight -= 1
	for word, m := range n.wordsToDoc {
		for doc, weight := range m {
			// Normalize weight.
			normWeight := weight - minWeight
			if normWeight == 1 {
				n.Words[word] = append(n.Words[word], doc.selfIndex)
			} else {
				n.Words[word] = append(n.Words[word], [2]interface{}{doc.selfIndex, normWeight})
			}
		}
	}
	return json.NewEncoder(w).Encode(n)
}

func (n *Index) addWord(word string, doc *Document, weight float64) {
	m := n.wordsToDoc[word]
	if m == nil {
		m = make(map[*Document]int)
		n.wordsToDoc[word] = m
	}
	m[doc] += int(weight * 100)
}

func (n *Index) newDocument(url, title string) *Document {
	doc := &Document{URL: url, Title: title}
	n.Docs = append(n.Docs, doc)
	return doc
}

func stem(word string) string {
	if strings.ContainsAny(word, "0123456789") {
		return word // don't stem words with digits
	}
	return porter2.Stemmer.Stem(word)
}

func (n *Index) addString(doc *Document, text string, wordWeight float64) {
	wordcnt := make(map[string]float64)
	tk := tokenizer.Words(text)
	for tk.Next() {
		w := normalizeWord(tk.Token())
		if len(w) < 1 || isStopWord(w) {
			continue
		}
		wordcnt[stem(w)] += wordWeight
	}
	// Scale each word weight and add it
	for w, c := range wordcnt {
		scaled := c / float64(len(wordcnt))
		n.addWord(w, doc, scaled)
	}
}

func (n *Index) AddText(url, title string, r io.Reader) error {
	var b bytes.Buffer
	if _, err := io.Copy(&b, r); err != nil {
		return err
	}
	n.addString(n.newDocument(url, title), b.String(), 1)
	return nil
}

func (n *Index) AddHTML(url string, r io.Reader) (indexed bool, err error) {
	title, content, indexable, err := parseHTML(r)
	if err != nil {
		return false, err
	}
	if !indexable {
		return false, nil
	}
	doc := n.newDocument(url, title)

	// Adjust word weight according to document level
	url = strings.TrimPrefix(url, "http://")
	url = strings.TrimPrefix(url, "https://")
	url = strings.TrimPrefix(url, "www.")
	url = strings.TrimSuffix(url, "/")

	level := float64(strings.Count(url, "/"))
	if level == 0 {
		level = 1
	}

	n.addString(doc, title, n.HTMLTitleWeight/level)
	n.addString(doc, content, 0.5+0.5*(n.ContentWordWeight/level))
	// Add URL components.
	url = strings.ReplaceAll(url, "/", " ")
	url = strings.ReplaceAll(url, "_", " ")
	url = strings.ReplaceAll(url, "-", " ")
	n.addString(doc, url, n.HTMLURLComponentWeight/level)
	return true, nil
}
<commit_msg>Sort doc indexes in search index for better compression.<commit_after>package indexer

import (
	"bytes"
	"encoding/json"
	"io"
	"math"
	"sort"
	"strings"

	"github.com/dchest/stemmer/porter2"

	"github.com/dchest/kkr/search/indexer/tokenizer"
)

type Index struct {
	Docs       []*Document              `json:"docs"`
	Words      map[string][]interface{} `json:"words"`
	wordsToDoc map[string]map[*Document]int

	ContentWordWeight      float64 `json:"-"`
	HTMLTitleWeight        float64 `json:"-"`
	HTMLURLComponentWeight float64 `json:"-"`
}

type Document struct {
	URL   string `json:"u"`
	Title string `json:"t"`

	numWords  int `json:"-"`
	selfIndex int `json:"-"`
}

func New() *Index {
	return &Index{
		Docs:                   make([]*Document, 0),
		Words:                  make(map[string][]interface{}),
		wordsToDoc:             make(map[string]map[*Document]int), // word => doc => weight
		ContentWordWeight:      1,
		HTMLTitleWeight:        3,
		HTMLURLComponentWeight: 3,
	}
}

type docIndexAndWeight [2]int

func (n *Index) WriteJSON(w io.Writer) error {
	// Calculate minWeight (for normalization later),
	// and add word frequencies to doc.
	minWeight := math.MaxInt
	for _, m := range n.wordsToDoc {
		for doc, weight := range m {
			if weight < minWeight {
				minWeight = weight
			}
			doc.numWords += 1
		}
	}

	// Make JSON smaller:
	// Sort docs by most frequent, so that smaller
	// doc ids are taken for most frequent ones.
	sort.Slice(n.Docs, func(i, j int) bool {
		return n.Docs[i].numWords > n.Docs[j].numWords
	})

	// Add indexes to docs
	for i, d := range n.Docs {
		d.selfIndex = i
	}

	minWeight -= 1
	for word, m := range n.wordsToDoc {
		for doc, weight := range m {
			// Normalize weight.
			normWeight := weight - minWeight
			if normWeight == 1 {
				n.Words[word] = append(n.Words[word], doc.selfIndex)
			} else {
				n.Words[word] = append(n.Words[word], docIndexAndWeight{doc.selfIndex, normWeight})
			}
		}
		// Sort doc indexes for better compression.
		w := n.Words[word]
		getDocIndex := func(i int) int {
			switch v := w[i].(type) {
			case docIndexAndWeight:
				return v[0]
			case int:
				return v
			default:
				panic("wrong type in index")
			}
		}
		sort.SliceStable(w, func(i, j int) bool {
			return getDocIndex(i) < getDocIndex(j)
		})
	}
	return json.NewEncoder(w).Encode(n)
}

func (n *Index) addWord(word string, doc *Document, weight float64) {
	m := n.wordsToDoc[word]
	if m == nil {
		m = make(map[*Document]int)
		n.wordsToDoc[word] = m
	}
	m[doc] += int(weight * 100)
}

func (n *Index) newDocument(url, title string) *Document {
	doc := &Document{URL: url, Title: title}
	n.Docs = append(n.Docs, doc)
	return doc
}

func stem(word string) string {
	if strings.ContainsAny(word, "0123456789") {
		return word // don't stem words with digits
	}
	return porter2.Stemmer.Stem(word)
}

func (n *Index) addString(doc *Document, text string, wordWeight float64) {
	wordcnt := make(map[string]float64)
	tk := tokenizer.Words(text)
	for tk.Next() {
		w := normalizeWord(tk.Token())
		if len(w) < 1 || isStopWord(w) {
			continue
		}
		wordcnt[stem(w)] += wordWeight
	}
	// Scale each word weight and add it
	for w, c := range wordcnt {
		scaled := c / float64(len(wordcnt))
		n.addWord(w, doc, scaled)
	}
}

func (n *Index) AddText(url, title string, r io.Reader) error {
	var b bytes.Buffer
	if _, err := io.Copy(&b, r); err != nil {
		return err
	}
	n.addString(n.newDocument(url, title), b.String(), 1)
	return nil
}

func (n *Index) AddHTML(url string, r io.Reader) (indexed bool, err error) {
	title, content, indexable, err := parseHTML(r)
	if err != nil {
		return false, err
	}
	if !indexable {
		return false, nil
	}
	doc := n.newDocument(url, title)

	// Adjust word weight according to document level
	url = strings.TrimPrefix(url, "http://")
	url = strings.TrimPrefix(url, "https://")
	url = strings.TrimPrefix(url, "www.")
	url = strings.TrimSuffix(url, "/")

	level := float64(strings.Count(url, "/"))
	if level == 0 {
		level = 1
	}

	n.addString(doc, title, n.HTMLTitleWeight/level)
	n.addString(doc, content, 0.5+0.5*(n.ContentWordWeight/level))
	// Add URL components.
	url = strings.ReplaceAll(url, "/", " ")
	url = strings.ReplaceAll(url, "_", " ")
	url = strings.ReplaceAll(url, "-", " ")
	n.addString(doc, url, n.HTMLURLComponentWeight/level)
	return true, nil
}
<|endoftext|>


<commit_before>package github

import (
	"io"

	"github.com/docker/docker/pkg/jsonmessage"
	"github.com/docker/docker/pkg/term"
	"github.com/ejholmes/hookshot/events"
	"github.com/remind101/empire"
	"github.com/remind101/pkg/trace"
	"github.com/remind101/tugboat"
	"golang.org/x/net/context"
)

// deployer represents something that can deploy a github deployment.
type deployer interface {
	// Deploy performs the deployment, writing output to w.
	Deploy(context.Context, events.Deployment, io.Writer) error
}

// newDeployer is a factory method that generates a composed deployer instance
// depending on the options.
func newDeployer(e *empire.Empire, opts Options) deployer {
	var d deployer
	d = &tracedDeployer{newEmpireDeployer(e, opts.ImageTemplate)}
	d = &prettyDeployer{deployer: d}

	if opts.TugboatURL != "" {
		d = newTugboatDeployer(d, opts.TugboatURL)
	}

	return &asyncDeployer{
		deployer: &tracedDeployer{deployer: d},
	}
}

// prettyDeployer is a deployer implementation that converts the raw json stream
// into pretty output using the jsonmessage package.
type prettyDeployer struct {
	deployer
}

func (d *prettyDeployer) Deploy(ctx context.Context, p events.Deployment, out io.Writer) error {
	r, w := io.Pipe()

	errCh := make(chan error, 1)
	go func() {
		err := d.deployer.Deploy(ctx, p, w)
		errCh <- err
		w.Close()
	}()

	outFd, _ := term.GetFdInfo(out)
	if err := jsonmessage.DisplayJSONMessagesStream(r, out, outFd, false); err != nil {
		w.Close()
		return err
	}

	err := <-errCh
	return err
}

// empireDeployer is a deployer implementation that uses the Deploy method in
// Empire to perform the deployment.
type empireDeployer struct {
	*empire.Empire
	imageTmpl string
}

// newEmpireDeployer returns a new empireDeployer instance.
func newEmpireDeployer(e *empire.Empire, imageTmpl string) *empireDeployer {
	return &empireDeployer{
		Empire:    e,
		imageTmpl: imageTmpl,
	}
}

func (d *empireDeployer) Deploy(ctx context.Context, p events.Deployment, w io.Writer) error {
	img, err := Image(d.imageTmpl, p)
	if err != nil {
		return err
	}

	_, err = d.Empire.Deploy(ctx, empire.DeploymentsCreateOpts{
		Image:  img,
		Output: w,
		User:   &empire.User{Name: p.Deployment.Creator.Login},
	})

	return err
}

// tugboatDeployer is an implementtion of the deployer interface that sends logs
// and updates the status of the deployment within a Tugboat instance.
type tugboatDeployer struct {
	deployer
	client *tugboat.Client
}

func newTugboatDeployer(d deployer, url string) *tugboatDeployer {
	c := tugboat.NewClient(nil)
	c.URL = url
	return &tugboatDeployer{
		deployer: d,
		client:   c,
	}
}

func (d *tugboatDeployer) Deploy(ctx context.Context, p events.Deployment, out io.Writer) error {
	opts := tugboat.NewDeployOptsFromEvent(p)

	// Perform the deployment, wrapped in Deploy. This will automatically
	// write hte logs to tugboat and update the deployment status when this
	// function returns.
	_, err := d.client.Deploy(ctx, opts, provider(func(ctx context.Context, _ *tugboat.Deployment, w io.Writer) error {
		// Write logs to both tugboat as well as the writer we were
		// provided (probably stdout).
		w = io.MultiWriter(w, out)

		if err := d.deployer.Deploy(ctx, p, w); err != nil {
			// If we got a deployment error, write the error to
			// tugboat and return tugboat.ErrFailed to indicate the
			// failure.
			io.WriteString(w, err.Error())
			return tugboat.ErrFailed
		}

		return nil
	}))

	return err
}

// provider implements the tugboat.Provider interface.
type provider func(context.Context, *tugboat.Deployment, io.Writer) error

func (fn provider) Name() string {
	return "empire"
}

func (fn provider) Deploy(ctx context.Context, d *tugboat.Deployment, w io.Writer) error {
	return fn(ctx, d, w)
}

// tracedDeployer is an implementation of the deployer interface that calls
// trace.Trace.
type tracedDeployer struct {
	deployer
}

func (d *tracedDeployer) Deploy(ctx context.Context, p events.Deployment, w io.Writer) (err error) {
	ctx, done := trace.Trace(ctx)
	err = d.deployer.Deploy(ctx, p, w)
	done(err, "Deploy",
		"repository", p.Repository.FullName,
		"creator", p.Deployment.Creator.Login,
		"ref", p.Deployment.Ref,
		"sha", p.Deployment.Sha,
	)
	return err
}

// asyncDeployer is a deployer implementation that runs the deployment in a
// goroutine.
type asyncDeployer struct {
	// Should use a tracedDeployer so errors are propagated.
	deployer *tracedDeployer
}

func (d *asyncDeployer) Deploy(ctx context.Context, p events.Deployment, w io.Writer) error {
	go d.deployer.Deploy(ctx, p, w)
	return nil
}
<commit_msg>Trace error after prettifying logs.<commit_after>package github

import (
	"io"

	"github.com/docker/docker/pkg/jsonmessage"
	"github.com/docker/docker/pkg/term"
	"github.com/ejholmes/hookshot/events"
	"github.com/remind101/empire"
	"github.com/remind101/pkg/trace"
	"github.com/remind101/tugboat"
	"golang.org/x/net/context"
)

// deployer represents something that can deploy a github deployment.
type deployer interface {
	// Deploy performs the deployment, writing output to w.
	Deploy(context.Context, events.Deployment, io.Writer) error
}

// newDeployer is a factory method that generates a composed deployer instance
// depending on the options.
func newDeployer(e *empire.Empire, opts Options) deployer {
	var d deployer
	d = newEmpireDeployer(e, opts.ImageTemplate)
	d = &prettyDeployer{deployer: d}
	d = &tracedDeployer{d}

	if opts.TugboatURL != "" {
		d = newTugboatDeployer(d, opts.TugboatURL)
	}

	return &asyncDeployer{
		deployer: &tracedDeployer{deployer: d},
	}
}

// prettyDeployer is a deployer implementation that converts the raw json stream
// into pretty output using the jsonmessage package.
type prettyDeployer struct {
	deployer
}

func (d *prettyDeployer) Deploy(ctx context.Context, p events.Deployment, out io.Writer) error {
	r, w := io.Pipe()

	errCh := make(chan error, 1)
	go func() {
		err := d.deployer.Deploy(ctx, p, w)
		errCh <- err
		w.Close()
	}()

	outFd, _ := term.GetFdInfo(out)
	if err := jsonmessage.DisplayJSONMessagesStream(r, out, outFd, false); err != nil {
		w.Close()
		return err
	}

	err := <-errCh
	return err
}

// empireDeployer is a deployer implementation that uses the Deploy method in
// Empire to perform the deployment.
type empireDeployer struct {
	*empire.Empire
	imageTmpl string
}

// newEmpireDeployer returns a new empireDeployer instance.
func newEmpireDeployer(e *empire.Empire, imageTmpl string) *empireDeployer {
	return &empireDeployer{
		Empire:    e,
		imageTmpl: imageTmpl,
	}
}

func (d *empireDeployer) Deploy(ctx context.Context, p events.Deployment, w io.Writer) error {
	img, err := Image(d.imageTmpl, p)
	if err != nil {
		return err
	}

	_, err = d.Empire.Deploy(ctx, empire.DeploymentsCreateOpts{
		Image:  img,
		Output: w,
		User:   &empire.User{Name: p.Deployment.Creator.Login},
	})

	return err
}

// tugboatDeployer is an implementtion of the deployer interface that sends logs
// and updates the status of the deployment within a Tugboat instance.
type tugboatDeployer struct {
	deployer
	client *tugboat.Client
}

func newTugboatDeployer(d deployer, url string) *tugboatDeployer {
	c := tugboat.NewClient(nil)
	c.URL = url
	return &tugboatDeployer{
		deployer: d,
		client:   c,
	}
}

func (d *tugboatDeployer) Deploy(ctx context.Context, p events.Deployment, out io.Writer) error {
	opts := tugboat.NewDeployOptsFromEvent(p)

	// Perform the deployment, wrapped in Deploy. This will automatically
	// write hte logs to tugboat and update the deployment status when this
	// function returns.
	_, err := d.client.Deploy(ctx, opts, provider(func(ctx context.Context, _ *tugboat.Deployment, w io.Writer) error {
		// Write logs to both tugboat as well as the writer we were
		// provided (probably stdout).
		w = io.MultiWriter(w, out)

		if err := d.deployer.Deploy(ctx, p, w); err != nil {
			// If we got a deployment error, write the error to
			// tugboat and return tugboat.ErrFailed to indicate the
			// failure.
			io.WriteString(w, err.Error())
			return tugboat.ErrFailed
		}

		return nil
	}))

	return err
}

// provider implements the tugboat.Provider interface.
type provider func(context.Context, *tugboat.Deployment, io.Writer) error

func (fn provider) Name() string {
	return "empire"
}

func (fn provider) Deploy(ctx context.Context, d *tugboat.Deployment, w io.Writer) error {
	return fn(ctx, d, w)
}

// tracedDeployer is an implementation of the deployer interface that calls
// trace.Trace.
type tracedDeployer struct {
	deployer
}

func (d *tracedDeployer) Deploy(ctx context.Context, p events.Deployment, w io.Writer) (err error) {
	ctx, done := trace.Trace(ctx)
	err = d.deployer.Deploy(ctx, p, w)
	done(err, "Deploy",
		"repository", p.Repository.FullName,
		"creator", p.Deployment.Creator.Login,
		"ref", p.Deployment.Ref,
		"sha", p.Deployment.Sha,
	)
	return err
}

// asyncDeployer is a deployer implementation that runs the deployment in a
// goroutine.
type asyncDeployer struct {
	// Should use a tracedDeployer so errors are propagated.
	deployer *tracedDeployer
}

func (d *asyncDeployer) Deploy(ctx context.Context, p events.Deployment, w io.Writer) error {
	go d.deployer.Deploy(ctx, p, w)
	return nil
}
<|endoftext|>