<commit_before>package main

import (
	"fmt"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

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

type Exporter struct {
	lastScrapeErrors                  prometheus.Gauge
	totalScrapes, totalScrapeFailures prometheus.Counter
}

const PROM_INSTANCE_UP_STATE_METRIC = "instance_up"

func NewExporter() (*Exporter, error) {
	return &Exporter{
		lastScrapeErrors: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: opts.PrometheusNamespace,
			Name:      "exporter_last_scrape_errors",
			Help:      "Last scrape error count for all monitored hosts / metrics",
		}),
		totalScrapes: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: opts.PrometheusNamespace,
			Name:      "exporter_total_scrapes",
			Help:      "Total scrape attempts.",
		}),
		totalScrapeFailures: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: opts.PrometheusNamespace,
			Name:      "exporter_total_scrape_failures",
			Help:      "Number of errors while executing metric queries",
		}),
	}, nil
}

// Not really needed for scraping to work
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
}

func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	var lastScrapeErrors float64

	e.totalScrapes.Add(1)
	ch <- e.totalScrapes

	monitoredDatabases := getMonitoredDatabasesSnapshot()
	if len(monitoredDatabases) == 0 {
		log.Warning("No dbs configured for monitoring. Check config")
		ch <- e.totalScrapeFailures
		e.lastScrapeErrors.Set(0)
		ch <- e.lastScrapeErrors
		return
	}
	for name, md := range monitoredDatabases {
		setInstanceUpDownState(ch, md) // makes easier to differentiate between PG instance / machine  failures
		// https://prometheus.io/docs/instrumenting/writing_exporters/#failed-scrapes

		for metric, interval := range md.Metrics {
			if metric == "change_events" {
				log.Warningf("[%s] Skipping change_events metric as host state is not supported for Prometheus currently", md.DBUniqueName)
				continue
			}
			if interval > 0 {
				log.Debugf("scraping [%s:%s]...", md.DBUniqueName, metric)
				metricStoreMessages, err := FetchMetrics(
					MetricFetchMessage{DBUniqueName: name, DBUniqueNameOrig: md.DBUniqueNameOrig, MetricName: metric, DBType: md.DBType, Interval: time.Second * time.Duration(interval)},
					nil,
					nil,
					CONTEXT_PROMETHEUS_SCRAPE)
				if err != nil {
					log.Errorf("failed to scrape [%s:%s]: %v", name, metric, err)
					e.totalScrapeFailures.Add(1)
					lastScrapeErrors++
					continue
				}
				if len(metricStoreMessages) > 0 {
					promMetrics := MetricStoreMessageToPromMetrics(metricStoreMessages[0])
					for _, pm := range promMetrics { // collect & send later in batch? capMetricChan = 1000 limit in prometheus code
						ch <- pm
					}
				}
			}
		}
	}
	ch <- e.totalScrapeFailures
	e.lastScrapeErrors.Set(lastScrapeErrors)
	ch <- e.lastScrapeErrors

	atomic.StoreInt64(&lastSuccessfulDatastoreWriteTimeEpoch, time.Now().Unix())
}

func setInstanceUpDownState(ch chan<- prometheus.Metric, md MonitoredDatabase) {
	log.Debugf("checking availability of configured DB [%s:%s]...", md.DBUniqueName, PROM_INSTANCE_UP_STATE_METRIC)
	vme, err := DBGetPGVersion(md.DBUniqueName, md.DBType, true)
	data := make(map[string]interface{})
	if err != nil {
		data[PROM_INSTANCE_UP_STATE_METRIC] = 0
		log.Errorf("[%s:%s] could not determine instance version, reporting as 'down': %v", md.DBUniqueName, PROM_INSTANCE_UP_STATE_METRIC, err)
		//return
	} else {
		data[PROM_INSTANCE_UP_STATE_METRIC] = 1
	}
	data[EPOCH_COLUMN_NAME] = time.Now().UnixNano()

	pm := MetricStoreMessageToPromMetrics(MetricStoreMessage{
		DBUniqueName:     md.DBUniqueName,
		DBType:           md.DBType,
		MetricName:       PROM_INSTANCE_UP_STATE_METRIC,
		CustomTags:       md.CustomTags,
		Data:             []map[string]interface{}{data},
		RealDbname:       vme.RealDbname,
		SystemIdentifier: vme.SystemIdentifier,
	})

	if len(pm) > 0 {
		ch <- pm[0]
	} else {
		log.Errorf("Could not formulate an instance state report - should not happen")
	}
}

func getMonitoredDatabasesSnapshot() map[string]MonitoredDatabase {
	mdSnap := make(map[string]MonitoredDatabase)

	if monitored_db_cache != nil {
		monitored_db_cache_lock.RLock()
		defer monitored_db_cache_lock.RUnlock()

		for _, row := range monitored_db_cache {
			mdSnap[row.DBUniqueName] = row
		}
	}

	return mdSnap
}

func MetricStoreMessageToPromMetrics(msg MetricStoreMessage) []prometheus.Metric {
	promMetrics := make([]prometheus.Metric, 0)

	var epoch_time time.Time
	var epoch_ns int64

	if len(msg.Data) == 0 {
		return promMetrics
	}

	epoch_ns, ok := (msg.Data[0][EPOCH_COLUMN_NAME]).(int64)
	if !ok {
		if msg.MetricName != "pgbouncer_stats" {
			log.Warning("No timestamp_ns found, (gatherer) server time will be used. measurement:", msg.MetricName)
		}
		epoch_time = time.Now()
	} else {
		epoch_time = time.Unix(0, epoch_ns)
	}

	for _, dr := range msg.Data {
		labels := make(map[string]string)
		fields := make(map[string]float64)
		labels["dbname"] = msg.DBUniqueName

		for k, v := range dr {
			if v == nil || v == "" || k == EPOCH_COLUMN_NAME {
				continue // not storing NULLs. epoch checked/assigned once
			}

			if strings.HasPrefix(k, "tag_") {
				tag := k[4:]
				labels[tag] = fmt.Sprintf("%v", v)
			} else {
				dataType := reflect.TypeOf(v).String()
				if dataType == "float64" || dataType == "float32" || dataType == "int64" || dataType == "int32" || dataType == "int" {
					f, err := strconv.ParseFloat(fmt.Sprintf("%v", v), 64)
					if err != nil {
						log.Debugf("Skipping scraping column %s of [%s:%s]: %v", k, msg.DBUniqueName, msg.MetricName, err)
					}
					fields[k] = f
				} else if dataType == "bool" {
					if v.(bool) {
						fields[k] = 1
					} else {
						fields[k] = 0
					}
				} else {
					log.Debugf("Skipping scraping column %s of [%s:%s], unsupported datatype: %s", k, msg.DBUniqueName, msg.MetricName, dataType)
					continue
				}
			}
		}
		if msg.CustomTags != nil {
			for k, v := range msg.CustomTags {
				labels[k] = fmt.Sprintf("%v", v)
			}
		}

		if addRealDbname && opts.RealDbnameField != "" && msg.RealDbname != "" {
			labels[opts.RealDbnameField] = msg.RealDbname
		}
		if addSystemIdentifier && opts.SystemIdentifierField != "" && msg.SystemIdentifier != "" {
			labels[opts.SystemIdentifierField] = msg.SystemIdentifier
		}

		label_keys := make([]string, 0)
		label_values := make([]string, 0)
		for k, v := range labels {
			label_keys = append(label_keys, k)
			label_values = append(label_values, v)
		}

		for field, value := range fields {
			skip := false
			fieldPromDataType := prometheus.CounterValue

			if msg.MetricDefinitionDetails.ColumnAttrs.PrometheusAllGaugeColumns {
				fieldPromDataType = prometheus.GaugeValue
			} else {
				for _, gaugeColumns := range msg.MetricDefinitionDetails.ColumnAttrs.PrometheusGaugeColumns {
					if gaugeColumns == field {
						fieldPromDataType = prometheus.GaugeValue
						break
					}
				}
			}
			if msg.MetricName == PROM_INSTANCE_UP_STATE_METRIC {
				fieldPromDataType = prometheus.GaugeValue
			}

			for _, ignoredColumns := range msg.MetricDefinitionDetails.ColumnAttrs.PrometheusIgnoredColumns {
				if ignoredColumns == field {
					skip = true
					break
				}
			}
			if skip {
				continue
			}
			var desc *prometheus.Desc
			if opts.PrometheusNamespace != "" {
				if msg.MetricName == PROM_INSTANCE_UP_STATE_METRIC { // handle the special "instance_up" check
					desc = prometheus.NewDesc(fmt.Sprintf("%s_%s", opts.PrometheusNamespace, msg.MetricName),
						msg.MetricName, label_keys, nil)
				} else {
					desc = prometheus.NewDesc(fmt.Sprintf("%s_%s_%s", opts.PrometheusNamespace, msg.MetricName, field),
						msg.MetricName, label_keys, nil)
				}
			} else {
				if msg.MetricName == PROM_INSTANCE_UP_STATE_METRIC { // handle the special "instance_up" check
					desc = prometheus.NewDesc(fmt.Sprintf("%s", field), msg.MetricName, label_keys, nil)
				} else {
					desc = prometheus.NewDesc(fmt.Sprintf("%s_%s", msg.MetricName, field), msg.MetricName, label_keys, nil)
				}
			}
			m := prometheus.MustNewConstMetric(desc, fieldPromDataType, value, label_values...)
			promMetrics = append(promMetrics, prometheus.NewMetricWithTimestamp(epoch_time, m))
		}
	}
	return promMetrics
}

func StartPrometheusExporter(port int64) {
	listenLoops := 0
	promExporter, err := NewExporter()
	if err != nil {
		log.Fatal(err)
	}

	prometheus.MustRegister(promExporter)

	var promServer = &http.Server{Addr: fmt.Sprintf("%s:%d", opts.PrometheusListenAddr, opts.PrometheusPort), Handler: promhttp.Handler()}

	for { // ListenAndServe call should not normally return, but looping just in case
		log.Infof("starting Prometheus exporter on %s:%d ...", opts.PrometheusListenAddr, opts.PrometheusPort)
		err = promServer.ListenAndServe()
		if listenLoops == 0 {
			log.Fatal("Prometheus listener failure:", err)
		} else {
			log.Error("Prometheus listener failure:", err)
		}
		time.Sleep(time.Second * 5)
	}
}
<commit_msg>Gatherer: fix a Prometheus mode bug when "instance_up" metric was explicitly specified<commit_after>package main

import (
	"fmt"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

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

type Exporter struct {
	lastScrapeErrors                  prometheus.Gauge
	totalScrapes, totalScrapeFailures prometheus.Counter
}

const PROM_INSTANCE_UP_STATE_METRIC = "instance_up"

func NewExporter() (*Exporter, error) {
	return &Exporter{
		lastScrapeErrors: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: opts.PrometheusNamespace,
			Name:      "exporter_last_scrape_errors",
			Help:      "Last scrape error count for all monitored hosts / metrics",
		}),
		totalScrapes: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: opts.PrometheusNamespace,
			Name:      "exporter_total_scrapes",
			Help:      "Total scrape attempts.",
		}),
		totalScrapeFailures: prometheus.NewCounter(prometheus.CounterOpts{
			Namespace: opts.PrometheusNamespace,
			Name:      "exporter_total_scrape_failures",
			Help:      "Number of errors while executing metric queries",
		}),
	}, nil
}

// Not really needed for scraping to work
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
}

func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	var lastScrapeErrors float64

	e.totalScrapes.Add(1)
	ch <- e.totalScrapes

	monitoredDatabases := getMonitoredDatabasesSnapshot()
	if len(monitoredDatabases) == 0 {
		log.Warning("No dbs configured for monitoring. Check config")
		ch <- e.totalScrapeFailures
		e.lastScrapeErrors.Set(0)
		ch <- e.lastScrapeErrors
		return
	}
	for name, md := range monitoredDatabases {
		setInstanceUpDownState(ch, md) // makes easier to differentiate between PG instance / machine  failures
		// https://prometheus.io/docs/instrumenting/writing_exporters/#failed-scrapes

		for metric, interval := range md.Metrics {
			if metric == SPECIAL_METRIC_CHANGE_EVENTS {
				log.Warningf("[%s] Skipping change_events metric as host state is not supported for Prometheus currently", md.DBUniqueName)
				continue
			}
			if metric == PROM_INSTANCE_UP_STATE_METRIC {
				continue // always included in Prometheus case
			}
			if interval > 0 {
				log.Debugf("scraping [%s:%s]...", md.DBUniqueName, metric)
				metricStoreMessages, err := FetchMetrics(
					MetricFetchMessage{DBUniqueName: name, DBUniqueNameOrig: md.DBUniqueNameOrig, MetricName: metric, DBType: md.DBType, Interval: time.Second * time.Duration(interval)},
					nil,
					nil,
					CONTEXT_PROMETHEUS_SCRAPE)
				if err != nil {
					log.Errorf("failed to scrape [%s:%s]: %v", name, metric, err)
					e.totalScrapeFailures.Add(1)
					lastScrapeErrors++
					continue
				}
				if len(metricStoreMessages) > 0 {
					promMetrics := MetricStoreMessageToPromMetrics(metricStoreMessages[0])
					for _, pm := range promMetrics { // collect & send later in batch? capMetricChan = 1000 limit in prometheus code
						ch <- pm
					}
				}
			}
		}
	}
	ch <- e.totalScrapeFailures
	e.lastScrapeErrors.Set(lastScrapeErrors)
	ch <- e.lastScrapeErrors

	atomic.StoreInt64(&lastSuccessfulDatastoreWriteTimeEpoch, time.Now().Unix())
}

func setInstanceUpDownState(ch chan<- prometheus.Metric, md MonitoredDatabase) {
	log.Debugf("checking availability of configured DB [%s:%s]...", md.DBUniqueName, PROM_INSTANCE_UP_STATE_METRIC)
	vme, err := DBGetPGVersion(md.DBUniqueName, md.DBType, true)
	data := make(map[string]interface{})
	if err != nil {
		data[PROM_INSTANCE_UP_STATE_METRIC] = 0
		log.Errorf("[%s:%s] could not determine instance version, reporting as 'down': %v", md.DBUniqueName, PROM_INSTANCE_UP_STATE_METRIC, err)
		//return
	} else {
		data[PROM_INSTANCE_UP_STATE_METRIC] = 1
	}
	data[EPOCH_COLUMN_NAME] = time.Now().UnixNano()

	pm := MetricStoreMessageToPromMetrics(MetricStoreMessage{
		DBUniqueName:     md.DBUniqueName,
		DBType:           md.DBType,
		MetricName:       PROM_INSTANCE_UP_STATE_METRIC,
		CustomTags:       md.CustomTags,
		Data:             []map[string]interface{}{data},
		RealDbname:       vme.RealDbname,
		SystemIdentifier: vme.SystemIdentifier,
	})

	if len(pm) > 0 {
		ch <- pm[0]
	} else {
		log.Errorf("Could not formulate an instance state report - should not happen")
	}
}

func getMonitoredDatabasesSnapshot() map[string]MonitoredDatabase {
	mdSnap := make(map[string]MonitoredDatabase)

	if monitored_db_cache != nil {
		monitored_db_cache_lock.RLock()
		defer monitored_db_cache_lock.RUnlock()

		for _, row := range monitored_db_cache {
			mdSnap[row.DBUniqueName] = row
		}
	}

	return mdSnap
}

func MetricStoreMessageToPromMetrics(msg MetricStoreMessage) []prometheus.Metric {
	promMetrics := make([]prometheus.Metric, 0)

	var epoch_time time.Time
	var epoch_ns int64

	if len(msg.Data) == 0 {
		return promMetrics
	}

	epoch_ns, ok := (msg.Data[0][EPOCH_COLUMN_NAME]).(int64)
	if !ok {
		if msg.MetricName != "pgbouncer_stats" {
			log.Warning("No timestamp_ns found, (gatherer) server time will be used. measurement:", msg.MetricName)
		}
		epoch_time = time.Now()
	} else {
		epoch_time = time.Unix(0, epoch_ns)
	}

	for _, dr := range msg.Data {
		labels := make(map[string]string)
		fields := make(map[string]float64)
		labels["dbname"] = msg.DBUniqueName

		for k, v := range dr {
			if v == nil || v == "" || k == EPOCH_COLUMN_NAME {
				continue // not storing NULLs. epoch checked/assigned once
			}

			if strings.HasPrefix(k, "tag_") {
				tag := k[4:]
				labels[tag] = fmt.Sprintf("%v", v)
			} else {
				dataType := reflect.TypeOf(v).String()
				if dataType == "float64" || dataType == "float32" || dataType == "int64" || dataType == "int32" || dataType == "int" {
					f, err := strconv.ParseFloat(fmt.Sprintf("%v", v), 64)
					if err != nil {
						log.Debugf("Skipping scraping column %s of [%s:%s]: %v", k, msg.DBUniqueName, msg.MetricName, err)
					}
					fields[k] = f
				} else if dataType == "bool" {
					if v.(bool) {
						fields[k] = 1
					} else {
						fields[k] = 0
					}
				} else {
					log.Debugf("Skipping scraping column %s of [%s:%s], unsupported datatype: %s", k, msg.DBUniqueName, msg.MetricName, dataType)
					continue
				}
			}
		}
		if msg.CustomTags != nil {
			for k, v := range msg.CustomTags {
				labels[k] = fmt.Sprintf("%v", v)
			}
		}

		if addRealDbname && opts.RealDbnameField != "" && msg.RealDbname != "" {
			labels[opts.RealDbnameField] = msg.RealDbname
		}
		if addSystemIdentifier && opts.SystemIdentifierField != "" && msg.SystemIdentifier != "" {
			labels[opts.SystemIdentifierField] = msg.SystemIdentifier
		}

		label_keys := make([]string, 0)
		label_values := make([]string, 0)
		for k, v := range labels {
			label_keys = append(label_keys, k)
			label_values = append(label_values, v)
		}

		for field, value := range fields {
			skip := false
			fieldPromDataType := prometheus.CounterValue

			if msg.MetricDefinitionDetails.ColumnAttrs.PrometheusAllGaugeColumns {
				fieldPromDataType = prometheus.GaugeValue
			} else {
				for _, gaugeColumns := range msg.MetricDefinitionDetails.ColumnAttrs.PrometheusGaugeColumns {
					if gaugeColumns == field {
						fieldPromDataType = prometheus.GaugeValue
						break
					}
				}
			}
			if msg.MetricName == PROM_INSTANCE_UP_STATE_METRIC {
				fieldPromDataType = prometheus.GaugeValue
			}

			for _, ignoredColumns := range msg.MetricDefinitionDetails.ColumnAttrs.PrometheusIgnoredColumns {
				if ignoredColumns == field {
					skip = true
					break
				}
			}
			if skip {
				continue
			}
			var desc *prometheus.Desc
			if opts.PrometheusNamespace != "" {
				if msg.MetricName == PROM_INSTANCE_UP_STATE_METRIC { // handle the special "instance_up" check
					desc = prometheus.NewDesc(fmt.Sprintf("%s_%s", opts.PrometheusNamespace, msg.MetricName),
						msg.MetricName, label_keys, nil)
				} else {
					desc = prometheus.NewDesc(fmt.Sprintf("%s_%s_%s", opts.PrometheusNamespace, msg.MetricName, field),
						msg.MetricName, label_keys, nil)
				}
			} else {
				if msg.MetricName == PROM_INSTANCE_UP_STATE_METRIC { // handle the special "instance_up" check
					desc = prometheus.NewDesc(fmt.Sprintf("%s", field), msg.MetricName, label_keys, nil)
				} else {
					desc = prometheus.NewDesc(fmt.Sprintf("%s_%s", msg.MetricName, field), msg.MetricName, label_keys, nil)
				}
			}
			m := prometheus.MustNewConstMetric(desc, fieldPromDataType, value, label_values...)
			promMetrics = append(promMetrics, prometheus.NewMetricWithTimestamp(epoch_time, m))
		}
	}
	return promMetrics
}

func StartPrometheusExporter(port int64) {
	listenLoops := 0
	promExporter, err := NewExporter()
	if err != nil {
		log.Fatal(err)
	}

	prometheus.MustRegister(promExporter)

	var promServer = &http.Server{Addr: fmt.Sprintf("%s:%d", opts.PrometheusListenAddr, opts.PrometheusPort), Handler: promhttp.Handler()}

	for { // ListenAndServe call should not normally return, but looping just in case
		log.Infof("starting Prometheus exporter on %s:%d ...", opts.PrometheusListenAddr, opts.PrometheusPort)
		err = promServer.ListenAndServe()
		if listenLoops == 0 {
			log.Fatal("Prometheus listener failure:", err)
		} else {
			log.Error("Prometheus listener failure:", err)
		}
		time.Sleep(time.Second * 5)
	}
}
<|endoftext|>


<commit_before>package sse

import (
	"errors"
	"io"
	"net/http"
	"sync"
	"time"
)

type ReadyState byte

const (
	AllowedContentType = "text/event-stream"

	// Connecting while trying to establish connection with the stream.
	Connecting ReadyState = iota - 1
	// Open after connection is established with the server.
	Open
	// Closing after Close is invoked.
	Closing
	// Closed after the connection is closed.
	Closed

	defaultRetry = time.Duration(1000)
)

var (
	ErrContentType = errors.New("eventsource: the content type of the stream is not allowed")

	// Map used by decoders to be able to change the retry time of the eventSource when
	// a retry event is received.
	globalDecoderMap = map[Decoder]*eventSource{}
	retryMux         = sync.RWMutex{}
)

type (
	// EventSource connects and processes events from an SSE stream.
	EventSource interface {
		URL() (url string)
		ReadyState() (state ReadyState)
		LastEventID() (id string)
		Events() (events <-chan Event)
		Close()
	}
	eventSource struct {
		url          string
		resp         *http.Response
		out          chan Event
		closeOutOnce chan bool

		// Last recorded event ID
		lastEventID    string
		lastEventIDMux sync.RWMutex

		// Status of the event stream.
		readyState    ReadyState
		readyStateMux sync.RWMutex

		// Reconnection waiting time in milliseconds
		retry time.Duration
	}
)

// NewEventSource constructs returns an EventSource that satisfies the HTML5 EventSource specification.
func NewEventSource(url string) (EventSource, error) {
	es := eventSource{
		url:          url,
		out:          make(chan Event),
		closeOutOnce: make(chan bool),
		retry:        defaultRetry,
	}

	// Ensure the output channel is closed only once.
	go es.closeOnce()

	return &es, es.connect()
}

// connect does a connection attempt, if the operation fails, attempt reconnecting
// according to the spec.
func (es *eventSource) connect() (err error) {
	es.setReadyState(Connecting)
	err = es.connectOnce()
	if err != nil {
		err = es.reconnect()
	}
	return
}

// reconnect to the stream several until the operation succeeds or the conditions
// to retry no longer hold true.
func (es *eventSource) reconnect() (err error) {
	es.setReadyState(Connecting)
	for es.mustReconnect(err) {
		time.Sleep(es.retry * time.Millisecond)
		err = es.connectOnce()
	}
	if err != nil {
		es.Close()
	}
	return
}

// Attempts to connect and updates internal status depending on the outcome.
func (es *eventSource) connectOnce() error {
	resp, err := http.Get(es.url)
	if err != nil {
		es.resp = nil
		return err
	}
	if resp.Header.Get("Content-Type") != AllowedContentType {
		return ErrContentType
	}
	es.resp = resp
	es.setReadyState(Open)
	go es.consume()
	return err
}

// Method consume() must be called once connect() succeeds.
// It parses the input reader and assigns the event output channel accordingly.
func (es *eventSource) consume() {
	d := NewDecoder(es.resp.Body)

	retryMux.Lock()
	globalDecoderMap[d] = es
	retryMux.Unlock()

	for {
		ev, err := d.Decode()
		if err != nil {
			if es.mustReconnect(err) {
				err = es.reconnect()
				return
			}
			es.Close()
			return
		}
		id := ev.ID()
		if id != "" {
			es.setLastEventID(id)
		}
		es.out <- ev
	}
}

// Clients will reconnect if the connection is closed;
// a client can be told to stop reconnecting using the HTTP 204 No Content response code.
func (es *eventSource) mustReconnect(err error) bool {
	switch err {
	case ErrContentType:
		return false
	case io.ErrUnexpectedEOF:
		return true
	}
	if es.resp != nil && es.resp.StatusCode == http.StatusNoContent {
		return false
	}
	return true
}

// Returns the event source URL.
func (es *eventSource) URL() string {
	return es.url
}

// Returns the event source connection state, either connecting, open or closed.
func (es *eventSource) ReadyState() ReadyState {
	es.readyStateMux.RLock()
	defer es.readyStateMux.RUnlock()
	return es.readyState
}

func (es *eventSource) setReadyState(newState ReadyState) {
	es.readyStateMux.Lock()
	defer es.readyStateMux.Unlock()
	es.readyState = newState
}

// Returns the last event source Event id.
func (es *eventSource) LastEventID() string {
	es.lastEventIDMux.RLock()
	defer es.lastEventIDMux.RUnlock()
	return es.lastEventID
}

func (es *eventSource) setLastEventID(id string) {
	es.lastEventIDMux.Lock()
	defer es.lastEventIDMux.Unlock()
	es.lastEventID = id
}

// Returns the channel of events. Events will be queued in the channel as they
// are received.
func (es *eventSource) Events() <-chan Event {
	return es.out
}

// Closes the event source.
// After closing the event source, it cannot be reused again.
func (es *eventSource) Close() {
	if es.ReadyState() == Closed {
		return
	}
	es.setReadyState(Closing)
	if es.resp != nil {
		es.resp.Body.Close()
	}
	es.sendClose()
	es.setReadyState(Closed)
}

func (es *eventSource) sendClose() {
	select {
	case es.closeOutOnce <- true:
	default:
	}
}

func (es *eventSource) closeOnce() {
	select {
	case <-es.closeOutOnce:
		close(es.out)
	}
}
<commit_msg>eventsource: update the condition to do nothing if the state is closing.<commit_after>package sse

import (
	"errors"
	"io"
	"net/http"
	"sync"
	"time"
)

type ReadyState byte

const (
	AllowedContentType = "text/event-stream"

	// Connecting while trying to establish connection with the stream.
	Connecting ReadyState = iota - 1
	// Open after connection is established with the server.
	Open
	// Closing after Close is invoked.
	Closing
	// Closed after the connection is closed.
	Closed

	defaultRetry = time.Duration(1000)
)

var (
	ErrContentType = errors.New("eventsource: the content type of the stream is not allowed")

	// Map used by decoders to be able to change the retry time of the eventSource when
	// a retry event is received.
	globalDecoderMap = map[Decoder]*eventSource{}
	retryMux         = sync.RWMutex{}
)

type (
	// EventSource connects and processes events from an SSE stream.
	EventSource interface {
		URL() (url string)
		ReadyState() (state ReadyState)
		LastEventID() (id string)
		Events() (events <-chan Event)
		Close()
	}
	eventSource struct {
		url          string
		resp         *http.Response
		out          chan Event
		closeOutOnce chan bool

		// Last recorded event ID
		lastEventID    string
		lastEventIDMux sync.RWMutex

		// Status of the event stream.
		readyState    ReadyState
		readyStateMux sync.RWMutex

		// Reconnection waiting time in milliseconds
		retry time.Duration
	}
)

// NewEventSource constructs returns an EventSource that satisfies the HTML5 EventSource specification.
func NewEventSource(url string) (EventSource, error) {
	es := eventSource{
		url:          url,
		out:          make(chan Event),
		closeOutOnce: make(chan bool),
		retry:        defaultRetry,
	}

	// Ensure the output channel is closed only once.
	go es.closeOnce()

	return &es, es.connect()
}

// connect does a connection attempt, if the operation fails, attempt reconnecting
// according to the spec.
func (es *eventSource) connect() (err error) {
	es.setReadyState(Connecting)
	err = es.connectOnce()
	if err != nil {
		err = es.reconnect()
	}
	return
}

// reconnect to the stream several until the operation succeeds or the conditions
// to retry no longer hold true.
func (es *eventSource) reconnect() (err error) {
	es.setReadyState(Connecting)
	for es.mustReconnect(err) {
		time.Sleep(es.retry * time.Millisecond)
		err = es.connectOnce()
	}
	if err != nil {
		es.Close()
	}
	return
}

// Attempts to connect and updates internal status depending on the outcome.
func (es *eventSource) connectOnce() error {
	resp, err := http.Get(es.url)
	if err != nil {
		es.resp = nil
		return err
	}
	if resp.Header.Get("Content-Type") != AllowedContentType {
		return ErrContentType
	}
	es.resp = resp
	es.setReadyState(Open)
	go es.consume()
	return err
}

// Method consume() must be called once connect() succeeds.
// It parses the input reader and assigns the event output channel accordingly.
func (es *eventSource) consume() {
	d := NewDecoder(es.resp.Body)

	retryMux.Lock()
	globalDecoderMap[d] = es
	retryMux.Unlock()

	for {
		ev, err := d.Decode()
		if err != nil {
			if es.mustReconnect(err) {
				err = es.reconnect()
				return
			}
			es.Close()
			return
		}
		id := ev.ID()
		if id != "" {
			es.setLastEventID(id)
		}
		es.out <- ev
	}
}

// Clients will reconnect if the connection is closed;
// a client can be told to stop reconnecting using the HTTP 204 No Content response code.
func (es *eventSource) mustReconnect(err error) bool {
	switch err {
	case ErrContentType:
		return false
	case io.ErrUnexpectedEOF:
		return true
	}
	if es.resp != nil && es.resp.StatusCode == http.StatusNoContent {
		return false
	}
	return true
}

// Returns the event source URL.
func (es *eventSource) URL() string {
	return es.url
}

// Returns the event source connection state, either connecting, open or closed.
func (es *eventSource) ReadyState() ReadyState {
	es.readyStateMux.RLock()
	defer es.readyStateMux.RUnlock()
	return es.readyState
}

func (es *eventSource) setReadyState(newState ReadyState) {
	es.readyStateMux.Lock()
	defer es.readyStateMux.Unlock()
	es.readyState = newState
}

// Returns the last event source Event id.
func (es *eventSource) LastEventID() string {
	es.lastEventIDMux.RLock()
	defer es.lastEventIDMux.RUnlock()
	return es.lastEventID
}

func (es *eventSource) setLastEventID(id string) {
	es.lastEventIDMux.Lock()
	defer es.lastEventIDMux.Unlock()
	es.lastEventID = id
}

// Returns the channel of events. Events will be queued in the channel as they
// are received.
func (es *eventSource) Events() <-chan Event {
	return es.out
}

// Closes the event source.
// After closing the event source, it cannot be reused again.
func (es *eventSource) Close() {
	state := es.ReadyState()
	if state == Closed || state == Closing {
		return
	}
	es.setReadyState(Closing)
	if es.resp != nil {
		es.resp.Body.Close()
	}
	es.sendClose()
	es.setReadyState(Closed)
}

func (es *eventSource) sendClose() {
	select {
	case es.closeOutOnce <- true:
	default:
	}
}

func (es *eventSource) closeOnce() {
	select {
	case <-es.closeOutOnce:
		close(es.out)
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/cilium/cilium/pkg/defaults"
	"github.com/cilium/cilium/pkg/k8s"
	clientset "github.com/cilium/cilium/pkg/k8s/client/clientset/versioned"
	"github.com/cilium/cilium/pkg/kvstore"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/version"

	gops "github.com/google/gops/agent"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var (
	log = logging.DefaultLogger.WithField(logfields.LogSubsys, "cilium-operator")

	rootCmd = &cobra.Command{
		Use:   "cilium-operator",
		Short: "Run the cilium-operator",
		Run: func(cmd *cobra.Command, args []string) {
			runOperator(cmd)
		},
	}

	k8sAPIServer        string
	k8sKubeConfigPath   string
	kvStore             string
	kvStoreOpts         = make(map[string]string)
	shutdownSignal      = make(chan bool, 1)
	synchronizeServices bool
	enableCepGC         bool

	ciliumK8sClient clientset.Interface
)

func main() {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-signals
		shutdownSignal <- true
	}()

	// Open socket for using gops to get stacktraces of the agent.
	if err := gops.Listen(gops.Options{}); err != nil {
		errorString := fmt.Sprintf("unable to start gops: %s", err)
		fmt.Println(errorString)
		os.Exit(-1)
	}

	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)

	flags := rootCmd.Flags()
	flags.Bool("version", false, "Print version information")
	flags.Int(option.ClusterIDName, 0, "Unique identifier of the cluster")
	option.BindEnv(option.ClusterIDName)
	flags.String(option.ClusterName, defaults.ClusterName, "Name of the cluster")
	option.BindEnv(option.ClusterName)
	flags.BoolP("debug", "D", false, "Enable debugging mode")
	flags.StringVar(&k8sAPIServer, "k8s-api-server", "", "Kubernetes api address server (for https use --k8s-kubeconfig-path instead)")
	flags.StringVar(&k8sKubeConfigPath, "k8s-kubeconfig-path", "", "Absolute path of the kubernetes kubeconfig file")
	flags.StringVar(&kvStore, "kvstore", "", "Key-value store type")
	flags.Var(option.NewNamedMapOptions("kvstore-opts", &kvStoreOpts, nil), "kvstore-opt", "Key-value store options")

	flags.BoolVar(&synchronizeServices, "synchronize-k8s-services", true, "Synchronize Kubernetes services to kvstore")
	flags.BoolVar(&enableCepGC, "cilium-endpoint-gc", true, "Enable CiliumEndpoint garbage collector")

	flags.IntVar(&unmanagedKubeDnsWatcherInterval, "unmanaged-pod-watcher-interval", 15, "Interval to check for unmanaged kube-dns pods (0 to disable)")

	// We need to obtain from Cilium ConfigMap if the CiliumEndpointCRD option
	// is enabled or disabled. This option is marked as hidden because the
	// Cilium Endpoint CRD controller is not in this program and by having it
	// being printed by operator --help could confuse users.
	flags.Bool(option.DisableCiliumEndpointCRDName, false, "")
	flags.MarkHidden(option.DisableCiliumEndpointCRDName)
	option.BindEnv(option.DisableCiliumEndpointCRDName)

	viper.BindPFlags(flags)
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if viper.GetBool("version") {
		fmt.Printf("Cilium %s\n", version.Version)
		os.Exit(0)
	}

	option.Config.ClusterName = viper.GetString(option.ClusterName)
	option.Config.ClusterID = viper.GetInt(option.ClusterIDName)
	option.Config.DisableCiliumEndpointCRD = viper.GetBool(option.DisableCiliumEndpointCRDName)

	viper.SetEnvPrefix("cilium")
	viper.SetConfigName("cilium-operator")
}

func runOperator(cmd *cobra.Command) {
	logging.SetupLogging([]string{}, map[string]string{}, "cilium-operator", viper.GetBool("debug"))

	// Make sure glog (used by the client-go dependency) logs to stderr, as it
	// will try to log to directories that may not exist in the cilium-operator
	// container and cause the cilium-operator to exit.
	flag.Set("logtostderr", "true")

	log.Infof("Cilium Operator %s", version.Version)

	if err := kvstore.Setup(kvStore, kvStoreOpts); err != nil {
		log.WithError(err).WithFields(logrus.Fields{
			"kvstore": kvStore,
			"address": kvStoreOpts[fmt.Sprintf("%s.address", kvStore)],
		}).Fatal("Unable to setup kvstore")
	}

	k8s.Configure(k8sAPIServer, k8sKubeConfigPath)
	if err := k8s.Init(); err != nil {
		log.WithError(err).Fatal("Unable to connect to Kubernetes apiserver")
	}

	restConfig, err := k8s.CreateConfig()
	if err != nil {
		log.WithError(err).Fatal("Unable to create rest configuration")
	}

	ciliumK8sClient, err = clientset.NewForConfig(restConfig)
	if err != nil {
		log.WithError(err).Fatal("Unable to create cilium network policy client")
	}

	if synchronizeServices {
		startSynchronizingServices()
	}

	if enableCepGC {
		enableCiliumEndpointSyncGC()
	}

	if option.Config.DisableCiliumEndpointCRD {
		log.Infof("KubeDNS unmanaged pods controller disabled as %q option is set to 'disabled' in Cilium ConfigMap", option.DisableCiliumEndpointCRDName)
	} else {
		enableUnmanagedKubeDNSController()
	}

	err = enableCNPWatcher()
	if err != nil {
		log.WithError(err).WithField("subsys", "CNPWatcher").Fatal(
			"Cannot connect to Kubernetes apiserver ")
	}

	for {
		select {
		case <-shutdownSignal:
			// graceful exit
			log.Info("Received termination signal. Shutting down")
			return
		case <-time.After(time.Second):
		}
	}
}
<commit_msg>cilium-operator.Dockerfile: set `klog` logging values from cilium-operator<commit_after>// Copyright 2018 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/cilium/cilium/pkg/defaults"
	"github.com/cilium/cilium/pkg/k8s"
	clientset "github.com/cilium/cilium/pkg/k8s/client/clientset/versioned"
	"github.com/cilium/cilium/pkg/kvstore"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/version"

	gops "github.com/google/gops/agent"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"k8s.io/klog"
)

var (
	log = logging.DefaultLogger.WithField(logfields.LogSubsys, "cilium-operator")

	rootCmd = &cobra.Command{
		Use:   "cilium-operator",
		Short: "Run the cilium-operator",
		Run: func(cmd *cobra.Command, args []string) {
			runOperator(cmd)
		},
	}

	k8sAPIServer        string
	k8sKubeConfigPath   string
	kvStore             string
	kvStoreOpts         = make(map[string]string)
	shutdownSignal      = make(chan bool, 1)
	synchronizeServices bool
	enableCepGC         bool

	ciliumK8sClient clientset.Interface
)

func main() {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-signals
		shutdownSignal <- true
	}()

	// Open socket for using gops to get stacktraces of the agent.
	if err := gops.Listen(gops.Options{}); err != nil {
		errorString := fmt.Sprintf("unable to start gops: %s", err)
		fmt.Println(errorString)
		os.Exit(-1)
	}

	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)

	flags := rootCmd.Flags()
	flags.Bool("version", false, "Print version information")
	flags.Int(option.ClusterIDName, 0, "Unique identifier of the cluster")
	option.BindEnv(option.ClusterIDName)
	flags.String(option.ClusterName, defaults.ClusterName, "Name of the cluster")
	option.BindEnv(option.ClusterName)
	flags.BoolP("debug", "D", false, "Enable debugging mode")
	flags.StringVar(&k8sAPIServer, "k8s-api-server", "", "Kubernetes api address server (for https use --k8s-kubeconfig-path instead)")
	flags.StringVar(&k8sKubeConfigPath, "k8s-kubeconfig-path", "", "Absolute path of the kubernetes kubeconfig file")
	flags.StringVar(&kvStore, "kvstore", "", "Key-value store type")
	flags.Var(option.NewNamedMapOptions("kvstore-opts", &kvStoreOpts, nil), "kvstore-opt", "Key-value store options")

	flags.BoolVar(&synchronizeServices, "synchronize-k8s-services", true, "Synchronize Kubernetes services to kvstore")
	flags.BoolVar(&enableCepGC, "cilium-endpoint-gc", true, "Enable CiliumEndpoint garbage collector")

	flags.IntVar(&unmanagedKubeDnsWatcherInterval, "unmanaged-pod-watcher-interval", 15, "Interval to check for unmanaged kube-dns pods (0 to disable)")

	// We need to obtain from Cilium ConfigMap if the CiliumEndpointCRD option
	// is enabled or disabled. This option is marked as hidden because the
	// Cilium Endpoint CRD controller is not in this program and by having it
	// being printed by operator --help could confuse users.
	flags.Bool(option.DisableCiliumEndpointCRDName, false, "")
	flags.MarkHidden(option.DisableCiliumEndpointCRDName)
	option.BindEnv(option.DisableCiliumEndpointCRDName)

	viper.BindPFlags(flags)

	// Make sure that klog logging variables are initialized so that we can
	// update them from this file.
	klog.InitFlags(nil)

	// Make sure klog (used by the client-go dependency) logs to stderr, as it
	// will try to log to directories that may not exist in the cilium-operator
	// container (/tmp) and cause the cilium-operator to exit.
	flag.Set("logtostderr", "true")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if viper.GetBool("version") {
		fmt.Printf("Cilium %s\n", version.Version)
		os.Exit(0)
	}

	option.Config.ClusterName = viper.GetString(option.ClusterName)
	option.Config.ClusterID = viper.GetInt(option.ClusterIDName)
	option.Config.DisableCiliumEndpointCRD = viper.GetBool(option.DisableCiliumEndpointCRDName)

	viper.SetEnvPrefix("cilium")
	viper.SetConfigName("cilium-operator")
}

func runOperator(cmd *cobra.Command) {
	logging.SetupLogging([]string{}, map[string]string{}, "cilium-operator", viper.GetBool("debug"))

	log.Infof("Cilium Operator %s", version.Version)

	if err := kvstore.Setup(kvStore, kvStoreOpts); err != nil {
		log.WithError(err).WithFields(logrus.Fields{
			"kvstore": kvStore,
			"address": kvStoreOpts[fmt.Sprintf("%s.address", kvStore)],
		}).Fatal("Unable to setup kvstore")
	}

	k8s.Configure(k8sAPIServer, k8sKubeConfigPath)
	if err := k8s.Init(); err != nil {
		log.WithError(err).Fatal("Unable to connect to Kubernetes apiserver")
	}

	restConfig, err := k8s.CreateConfig()
	if err != nil {
		log.WithError(err).Fatal("Unable to create rest configuration")
	}

	ciliumK8sClient, err = clientset.NewForConfig(restConfig)
	if err != nil {
		log.WithError(err).Fatal("Unable to create cilium network policy client")
	}

	if synchronizeServices {
		startSynchronizingServices()
	}

	if enableCepGC {
		enableCiliumEndpointSyncGC()
	}

	if option.Config.DisableCiliumEndpointCRD {
		log.Infof("KubeDNS unmanaged pods controller disabled as %q option is set to 'disabled' in Cilium ConfigMap", option.DisableCiliumEndpointCRDName)
	} else {
		enableUnmanagedKubeDNSController()
	}

	err = enableCNPWatcher()
	if err != nil {
		log.WithError(err).WithField("subsys", "CNPWatcher").Fatal(
			"Cannot connect to Kubernetes apiserver ")
	}

	for {
		select {
		case <-shutdownSignal:
			// graceful exit
			log.Info("Received termination signal. Shutting down")
			return
		case <-time.After(time.Second):
		}
	}
}
<|endoftext|>


<commit_before>// Copyright © 2016 Mark Bates <mark@markbates.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/gobuffalo/velvet"
	"github.com/spf13/cobra"
)

var outputBinName string
var zipBin bool

type builder struct {
	cleanup       []string
	original_main []byte
	workDir       string
}

func (b *builder) clean(name ...string) string {
	path := filepath.Join(name...)
	b.cleanup = append(b.cleanup, path)
	return path
}

func (b *builder) exec(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	fmt.Printf("--> running %s\n", strings.Join(cmd.Args, " "))
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func (b *builder) execQuiet(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	return cmd.Run()
}

func (b *builder) buildWebpack() error {
	_, err := os.Stat("webpack.config.js")
	if err == nil {
		// build webpack
		return b.exec("webpack")
	}
	return nil
}

func (b *builder) buildAPack() error {
	err := os.MkdirAll(b.clean("a"), 0766)
	if err != nil {
		return err
	}
	err = b.buildAInit()
	if err != nil {
		return err
	}
	err = b.buildDatabase()
	if err != nil {
		return err
	}
	return nil
}

func (b *builder) buildAInit() error {
	a, err := os.Create(b.clean("a", "a.go"))
	if err != nil {
		return err
	}
	a.WriteString(aGo)
	return nil
}

func (b *builder) buildDatabase() error {
	bb := &bytes.Buffer{}
	dgo, err := os.Create(b.clean("a", "database.go"))
	if err != nil {
		return err
	}
	_, err = os.Stat("database.yml")
	if err == nil {
		// copy the database.yml file to the migrations folder so it's available through rice
		d, err := os.Open("database.yml")
		if err != nil {
			return err
		}
		_, err = io.Copy(bb, d)
		if err != nil {
			return err
		}
	}
	dgo.WriteString("package a\n")
	dgo.WriteString(fmt.Sprintf("var DB_CONFIG = `%s`", bb.String()))
	return nil
}

func (b *builder) buildRiceZip() error {
	defer os.Chdir(b.workDir)
	_, err := exec.LookPath("rice")
	if err == nil {
		paths := map[string]bool{}
		// if rice exists, try and build some cleanup:
		err = filepath.Walk(b.workDir, func(path string, info os.FileInfo, err error) error {
			if info.IsDir() {
				base := filepath.Base(path)
				if base == "node_modules" || base == ".git" || base == "bin" {
					return filepath.SkipDir
				}
			} else {
				err = os.Chdir(filepath.Dir(path))
				if err != nil {
					return err
				}

				s, err := ioutil.ReadFile(path)
				if err != nil {
					return err
				}
				rx := regexp.MustCompile("(rice.FindBox|rice.MustFindBox)")
				if rx.Match(s) {
					gopath := filepath.Join(os.Getenv("GOPATH"), "src")
					pkg := filepath.Dir(strings.Replace(path, gopath+"/", "", -1))
					paths[pkg] = true
				}
			}
			return nil
		})
		if err != nil {
			return err
		}
		if len(paths) != 0 {
			args := []string{"append", "--exec", filepath.Join(b.workDir, outputBinName)}
			for k, _ := range paths {
				args = append(args, "-i", k)
			}
			return b.exec("rice", args...)
		}
		// rice append --exec example
	}
	return nil
}
func (b *builder) buildRiceEmbedded() error {
	defer os.Chdir(b.workDir)
	_, err := exec.LookPath("rice")
	if err == nil {
		// if rice exists, try and build some cleanup:
		err = filepath.Walk(b.workDir, func(path string, info os.FileInfo, err error) error {
			if info.IsDir() {
				base := filepath.Base(path)
				if base == "node_modules" || base == ".git" {
					return filepath.SkipDir
				}
				err = os.Chdir(path)
				if err != nil {
					return err
				}
				err = b.execQuiet("rice", "embed-go")
				if err == nil {
					bp := filepath.Join(path, "rice-box.go")
					_, err := os.Stat(bp)
					if err == nil {
						fmt.Printf("--> built rice box %s\n", bp)
						b.clean(bp)
					}
				}
			}
			return nil
		})
		if err != nil {
			return err
		}
		// rice append --exec example
	}
	return nil
}

func (b *builder) buildMain() error {
	new_main := strings.Replace(string(b.original_main), "func main()", "func original_main()", 1)
	maingo, err := os.Create("main.go")
	if err != nil {
		return err
	}
	_, err = maingo.WriteString(new_main)
	if err != nil {
		return err
	}

	root, err := rootPath("")
	if err != nil {
		return err
	}
	ctx := velvet.NewContext()
	ctx.Set("root", root)
	ctx.Set("modelsPack", filepath.Join(packagePath(root), "models"))
	ctx.Set("aPack", filepath.Join(packagePath(root), "a"))
	ctx.Set("name", filepath.Base(root))
	s, err := velvet.Render(buildMainTmpl, ctx)
	if err != nil {
		return err
	}
	f, err := os.Create(b.clean("buffalo_build_main.go"))
	if err != nil {
		return err
	}
	f.WriteString(s)

	return nil
}

func (b *builder) cleanupBuild() {
	fmt.Println("--> cleaning up build")
	for _, b := range b.cleanup {
		fmt.Printf("--> cleaning up %s\n", b)
		os.RemoveAll(b)
	}
	maingo, _ := os.Create("main.go")
	maingo.Write(b.original_main)
}

func (b *builder) run() error {
	err := b.buildMain()
	if err != nil {
		return err
	}

	err = b.buildWebpack()
	if err != nil {
		return err
	}

	err = b.buildAPack()
	if err != nil {
		return err
	}

	err = b.buildMain()
	if err != nil {
		return err
	}

	if zipBin {
		err = b.buildBin()
		if err != nil {
			return err
		}
		return b.buildRiceZip()
	}

	err = b.buildRiceEmbedded()
	if err != nil {
		return err
	}
	return b.buildBin()
}

func (b *builder) buildBin() error {
	buildArgs := []string{"build", "-v", "-o", outputBinName}
	_, err := exec.LookPath("git")
	buildTime := fmt.Sprintf("\"%s\"", time.Now().Format(time.RFC3339))
	version := buildTime
	if err == nil {
		cmd := exec.Command("git", "rev-parse", "--short", "HEAD")
		out := &bytes.Buffer{}
		cmd.Stdout = out
		err = cmd.Run()
		if err == nil && out.String() != "" {
			version = strings.TrimSpace(out.String())
		}
	}
	buildArgs = append(buildArgs, "-ldflags", fmt.Sprintf("-X main.version=%s -X main.buildTime=%s", version, buildTime))

	return b.exec("go", buildArgs...)
}

// buildCmd represents the build command
var buildCmd = &cobra.Command{
	Use:   "build",
	Short: "Builds a Buffalo binary, including bundling of assets (go.rice & webpack)",
	RunE: func(cc *cobra.Command, args []string) error {
		original_main := &bytes.Buffer{}
		maingo, err := os.Open("main.go")
		_, err = original_main.ReadFrom(maingo)
		if err != nil {
			return err
		}
		maingo.Close()
		pwd, _ := os.Getwd()
		b := builder{
			cleanup:       []string{},
			original_main: original_main.Bytes(),
			workDir:       pwd,
		}
		defer b.cleanupBuild()

		return b.run()
	},
}

func init() {
	RootCmd.AddCommand(buildCmd)
	pwd, _ := os.Getwd()
	buildCmd.Flags().StringVarP(&outputBinName, "output", "o", filepath.Join("bin", filepath.Base(pwd)), "set the name of the binary")
	buildCmd.Flags().BoolVarP(&zipBin, "zip", "z", false, "zips the assets to the binary, this requires zip installed")
}

var buildMainTmpl = `package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	rice "github.com/GeertJohan/go.rice"
	_ "{{aPack}}"
	"{{modelsPack}}"
)

var version = "unknown"
var buildTime = "unknown"
var migrationBox *rice.Box

func main() {
	args := os.Args
	if len(args) == 1 {
		original_main()
	}
	c := args[1]
	switch c {
	case "migrate":
		migrate()
	case "start", "run", "serve":
		printVersion()
		original_main()
	case "version":
		printVersion()
	default:
		log.Fatalf("Could not find a command named: %s", c)
	}
}

func printVersion() {
	fmt.Printf("{{name}} version %s (%s)\n\n", version, buildTime)
}

func migrate() {
	var err error
	migrationBox, err = rice.FindBox("./migrations")
	if err != nil {
		fmt.Println("--> No migrations found.")
		return
	}
	fmt.Println("--> Running migrations")
	path, err := unpackMigrations()
	if err != nil {
		log.Fatalf("Failed to unpack migrations: %s", err)
	}
	defer os.RemoveAll(path)

	models.DB.MigrateUp(path)
}

func unpackMigrations() (string, error) {
	dir, err := ioutil.TempDir("", "{{name}}-migrations")
	if err != nil {
		log.Fatalf("Unable to create temp directory: %s", err)
	}

	migrationBox.Walk("", func(path string, fi os.FileInfo, e error) error {
		if !fi.IsDir() {
			content := migrationBox.MustBytes(path)
			file := filepath.Join(dir, path)
			if err := ioutil.WriteFile(file, content, 0666); err != nil {
				log.Fatalf("Failed to write migration to disk: %s", err)
			}
		}
		return e
	})

	return dir, nil
}`

var aGo = `package a

import (
	"log"
	"os"
)

func init() {
	dropDatabaseYml()
}

func dropDatabaseYml() {
	if DB_CONFIG != "" {

		_, err := os.Stat("database.yml")
		if err == nil {
			// yaml already exists, don't do anything
			return
		}
		f, err := os.Create("database.yml")
		if err != nil {
			log.Fatal(err)
		}
		_, err = f.WriteString(DB_CONFIG)
		if err != nil {
			log.Fatal(err)
		}
	}
}`
<commit_msg>added a couple of aliases for the build command<commit_after>// Copyright © 2016 Mark Bates <mark@markbates.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/gobuffalo/velvet"
	"github.com/spf13/cobra"
)

var outputBinName string
var zipBin bool

type builder struct {
	cleanup       []string
	original_main []byte
	workDir       string
}

func (b *builder) clean(name ...string) string {
	path := filepath.Join(name...)
	b.cleanup = append(b.cleanup, path)
	return path
}

func (b *builder) exec(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	fmt.Printf("--> running %s\n", strings.Join(cmd.Args, " "))
	cmd.Stdin = os.Stdin
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	return cmd.Run()
}

func (b *builder) execQuiet(name string, args ...string) error {
	cmd := exec.Command(name, args...)
	return cmd.Run()
}

func (b *builder) buildWebpack() error {
	_, err := os.Stat("webpack.config.js")
	if err == nil {
		// build webpack
		return b.exec("webpack")
	}
	return nil
}

func (b *builder) buildAPack() error {
	err := os.MkdirAll(b.clean("a"), 0766)
	if err != nil {
		return err
	}
	err = b.buildAInit()
	if err != nil {
		return err
	}
	err = b.buildDatabase()
	if err != nil {
		return err
	}
	return nil
}

func (b *builder) buildAInit() error {
	a, err := os.Create(b.clean("a", "a.go"))
	if err != nil {
		return err
	}
	a.WriteString(aGo)
	return nil
}

func (b *builder) buildDatabase() error {
	bb := &bytes.Buffer{}
	dgo, err := os.Create(b.clean("a", "database.go"))
	if err != nil {
		return err
	}
	_, err = os.Stat("database.yml")
	if err == nil {
		// copy the database.yml file to the migrations folder so it's available through rice
		d, err := os.Open("database.yml")
		if err != nil {
			return err
		}
		_, err = io.Copy(bb, d)
		if err != nil {
			return err
		}
	}
	dgo.WriteString("package a\n")
	dgo.WriteString(fmt.Sprintf("var DB_CONFIG = `%s`", bb.String()))
	return nil
}

func (b *builder) buildRiceZip() error {
	defer os.Chdir(b.workDir)
	_, err := exec.LookPath("rice")
	if err == nil {
		paths := map[string]bool{}
		// if rice exists, try and build some cleanup:
		err = filepath.Walk(b.workDir, func(path string, info os.FileInfo, err error) error {
			if info.IsDir() {
				base := filepath.Base(path)
				if base == "node_modules" || base == ".git" || base == "bin" {
					return filepath.SkipDir
				}
			} else {
				err = os.Chdir(filepath.Dir(path))
				if err != nil {
					return err
				}

				s, err := ioutil.ReadFile(path)
				if err != nil {
					return err
				}
				rx := regexp.MustCompile("(rice.FindBox|rice.MustFindBox)")
				if rx.Match(s) {
					gopath := filepath.Join(os.Getenv("GOPATH"), "src")
					pkg := filepath.Dir(strings.Replace(path, gopath+"/", "", -1))
					paths[pkg] = true
				}
			}
			return nil
		})
		if err != nil {
			return err
		}
		if len(paths) != 0 {
			args := []string{"append", "--exec", filepath.Join(b.workDir, outputBinName)}
			for k, _ := range paths {
				args = append(args, "-i", k)
			}
			return b.exec("rice", args...)
		}
		// rice append --exec example
	}
	return nil
}
func (b *builder) buildRiceEmbedded() error {
	defer os.Chdir(b.workDir)
	_, err := exec.LookPath("rice")
	if err == nil {
		// if rice exists, try and build some cleanup:
		err = filepath.Walk(b.workDir, func(path string, info os.FileInfo, err error) error {
			if info.IsDir() {
				base := filepath.Base(path)
				if base == "node_modules" || base == ".git" {
					return filepath.SkipDir
				}
				err = os.Chdir(path)
				if err != nil {
					return err
				}
				err = b.execQuiet("rice", "embed-go")
				if err == nil {
					bp := filepath.Join(path, "rice-box.go")
					_, err := os.Stat(bp)
					if err == nil {
						fmt.Printf("--> built rice box %s\n", bp)
						b.clean(bp)
					}
				}
			}
			return nil
		})
		if err != nil {
			return err
		}
		// rice append --exec example
	}
	return nil
}

func (b *builder) buildMain() error {
	new_main := strings.Replace(string(b.original_main), "func main()", "func original_main()", 1)
	maingo, err := os.Create("main.go")
	if err != nil {
		return err
	}
	_, err = maingo.WriteString(new_main)
	if err != nil {
		return err
	}

	root, err := rootPath("")
	if err != nil {
		return err
	}
	ctx := velvet.NewContext()
	ctx.Set("root", root)
	ctx.Set("modelsPack", filepath.Join(packagePath(root), "models"))
	ctx.Set("aPack", filepath.Join(packagePath(root), "a"))
	ctx.Set("name", filepath.Base(root))
	s, err := velvet.Render(buildMainTmpl, ctx)
	if err != nil {
		return err
	}
	f, err := os.Create(b.clean("buffalo_build_main.go"))
	if err != nil {
		return err
	}
	f.WriteString(s)

	return nil
}

func (b *builder) cleanupBuild() {
	fmt.Println("--> cleaning up build")
	for _, b := range b.cleanup {
		fmt.Printf("--> cleaning up %s\n", b)
		os.RemoveAll(b)
	}
	maingo, _ := os.Create("main.go")
	maingo.Write(b.original_main)
}

func (b *builder) run() error {
	err := b.buildMain()
	if err != nil {
		return err
	}

	err = b.buildWebpack()
	if err != nil {
		return err
	}

	err = b.buildAPack()
	if err != nil {
		return err
	}

	err = b.buildMain()
	if err != nil {
		return err
	}

	if zipBin {
		err = b.buildBin()
		if err != nil {
			return err
		}
		return b.buildRiceZip()
	}

	err = b.buildRiceEmbedded()
	if err != nil {
		return err
	}
	return b.buildBin()
}

func (b *builder) buildBin() error {
	buildArgs := []string{"build", "-v", "-o", outputBinName}
	_, err := exec.LookPath("git")
	buildTime := fmt.Sprintf("\"%s\"", time.Now().Format(time.RFC3339))
	version := buildTime
	if err == nil {
		cmd := exec.Command("git", "rev-parse", "--short", "HEAD")
		out := &bytes.Buffer{}
		cmd.Stdout = out
		err = cmd.Run()
		if err == nil && out.String() != "" {
			version = strings.TrimSpace(out.String())
		}
	}
	buildArgs = append(buildArgs, "-ldflags", fmt.Sprintf("-X main.version=%s -X main.buildTime=%s", version, buildTime))

	return b.exec("go", buildArgs...)
}

// buildCmd represents the build command
var buildCmd = &cobra.Command{
	Use:     "build",
	Aliases: []string{"b", "bill"},
	Short:   "Builds a Buffalo binary, including bundling of assets (go.rice & webpack)",
	RunE: func(cc *cobra.Command, args []string) error {
		original_main := &bytes.Buffer{}
		maingo, err := os.Open("main.go")
		_, err = original_main.ReadFrom(maingo)
		if err != nil {
			return err
		}
		maingo.Close()
		pwd, _ := os.Getwd()
		b := builder{
			cleanup:       []string{},
			original_main: original_main.Bytes(),
			workDir:       pwd,
		}
		defer b.cleanupBuild()

		return b.run()
	},
}

func init() {
	RootCmd.AddCommand(buildCmd)
	pwd, _ := os.Getwd()
	buildCmd.Flags().StringVarP(&outputBinName, "output", "o", filepath.Join("bin", filepath.Base(pwd)), "set the name of the binary")
	buildCmd.Flags().BoolVarP(&zipBin, "zip", "z", false, "zips the assets to the binary, this requires zip installed")
}

var buildMainTmpl = `package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	rice "github.com/GeertJohan/go.rice"
	_ "{{aPack}}"
	"{{modelsPack}}"
)

var version = "unknown"
var buildTime = "unknown"
var migrationBox *rice.Box

func main() {
	args := os.Args
	if len(args) == 1 {
		original_main()
	}
	c := args[1]
	switch c {
	case "migrate":
		migrate()
	case "start", "run", "serve":
		printVersion()
		original_main()
	case "version":
		printVersion()
	default:
		log.Fatalf("Could not find a command named: %s", c)
	}
}

func printVersion() {
	fmt.Printf("{{name}} version %s (%s)\n\n", version, buildTime)
}

func migrate() {
	var err error
	migrationBox, err = rice.FindBox("./migrations")
	if err != nil {
		fmt.Println("--> No migrations found.")
		return
	}
	fmt.Println("--> Running migrations")
	path, err := unpackMigrations()
	if err != nil {
		log.Fatalf("Failed to unpack migrations: %s", err)
	}
	defer os.RemoveAll(path)

	models.DB.MigrateUp(path)
}

func unpackMigrations() (string, error) {
	dir, err := ioutil.TempDir("", "{{name}}-migrations")
	if err != nil {
		log.Fatalf("Unable to create temp directory: %s", err)
	}

	migrationBox.Walk("", func(path string, fi os.FileInfo, e error) error {
		if !fi.IsDir() {
			content := migrationBox.MustBytes(path)
			file := filepath.Join(dir, path)
			if err := ioutil.WriteFile(file, content, 0666); err != nil {
				log.Fatalf("Failed to write migration to disk: %s", err)
			}
		}
		return e
	})

	return dir, nil
}`

var aGo = `package a

import (
	"log"
	"os"
)

func init() {
	dropDatabaseYml()
}

func dropDatabaseYml() {
	if DB_CONFIG != "" {

		_, err := os.Stat("database.yml")
		if err == nil {
			// yaml already exists, don't do anything
			return
		}
		f, err := os.Create("database.yml")
		if err != nil {
			log.Fatal(err)
		}
		_, err = f.WriteString(DB_CONFIG)
		if err != nil {
			log.Fatal(err)
		}
	}
}`
<|endoftext|>


<commit_before>package bolt

import (
	"bytes"
)

// Bucket represents a collection of key/value pairs inside the database.
type Bucket struct {
	*bucket
	name string
	tx   *Tx
}

// bucket represents the on-file representation of a bucket.
type bucket struct {
	root     pgid
	sequence uint64
}

// Name returns the name of the bucket.
func (b *Bucket) Name() string {
	return b.name
}

// Writable returns whether the bucket is writable.
func (b *Bucket) Writable() bool {
	return b.tx.writable
}

// Cursor creates a cursor associated with the bucket.
// The cursor is only valid as long as the transaction is open.
// Do not use a cursor after the transaction is closed.
func (b *Bucket) Cursor() *Cursor {
	return &Cursor{
		tx:    b.tx,
		root:  b.root,
		stack: make([]elemRef, 0),
	}
}

// Get retrieves the value for a key in the bucket.
// Returns a nil value if the key does not exist.
func (b *Bucket) Get(key []byte) []byte {
	c := b.Cursor()
	k, v := c.Seek(key)

	// If our target node isn't the same key as what's passed in then return nil.
	if !bytes.Equal(key, k) {
		return nil
	}
	return v
}

// Put sets the value for a key in the bucket.
// If the key exist then its previous value will be overwritten.
// Returns an error if the bucket was created from a read-only transaction, if the key is blank, if the key is too large, or if the value is too large.
func (b *Bucket) Put(key []byte, value []byte) error {
	if !b.Writable() {
		return ErrBucketNotWritable
	}

	// Validate the key and data size.
	if len(key) == 0 {
		return ErrKeyRequired
	} else if len(key) > MaxKeySize {
		return ErrKeyTooLarge
	} else if len(value) > MaxValueSize {
		return ErrValueTooLarge
	}

	// Move cursor to correct position.
	c := b.Cursor()
	c.Seek(key)

	// Insert the key/value.
	c.node(b.tx).put(key, key, value, 0)

	return nil
}

// Delete removes a key from the bucket.
// If the key does not exist then nothing is done and a nil error is returned.
// Returns an error if the bucket was created from a read-only transaction.
func (b *Bucket) Delete(key []byte) error {
	if !b.Writable() {
		return ErrBucketNotWritable
	}

	// Move cursor to correct position.
	c := b.Cursor()
	c.Seek(key)

	// Delete the node if we have a matching key.
	c.node(b.tx).del(key)

	return nil
}

// NextSequence returns an autoincrementing integer for the bucket.
func (b *Bucket) NextSequence() (int, error) {
	if !b.Writable() {
		return 0, ErrBucketNotWritable
	}

	// Make sure next sequence number will not be larger than the maximum
	// integer size of the system.
	if b.bucket.sequence == uint64(maxInt) {
		return 0, ErrSequenceOverflow
	}

	// Increment and return the sequence.
	b.bucket.sequence++

	return int(b.bucket.sequence), nil
}

// ForEach executes a function for each key/value pair in a bucket.
// An error is returned if the bucket cannot be found.
func (b *Bucket) ForEach(fn func(k, v []byte) error) error {
	c := b.Cursor()
	for k, v := c.First(); k != nil; k, v = c.Next() {
		if err := fn(k, v); err != nil {
			return err
		}
	}
	return nil
}

// Stat returns stats on a bucket.
func (b *Bucket) Stat() *BucketStat {
	s := &BucketStat{}
	b.tx.forEachPage(b.root, 0, func(p *page, depth int) {
		if (p.flags & leafPageFlag) != 0 {
			s.LeafPageCount++
			s.KeyCount += int(p.count)
		} else if (p.flags & branchPageFlag) != 0 {
			s.BranchPageCount++
		}

		s.OverflowPageCount += int(p.overflow)

		if depth+1 > s.MaxDepth {
			s.MaxDepth = (depth + 1)
		}
	})
	return s
}

// BucketStat represents stats on a bucket such as branch pages and leaf pages.
type BucketStat struct {
	BranchPageCount   int
	LeafPageCount     int
	OverflowPageCount int
	KeyCount          int
	MaxDepth          int
}
<commit_msg>Fix Bucket.ForEach() comment.<commit_after>package bolt

import (
	"bytes"
)

// Bucket represents a collection of key/value pairs inside the database.
type Bucket struct {
	*bucket
	name string
	tx   *Tx
}

// bucket represents the on-file representation of a bucket.
type bucket struct {
	root     pgid
	sequence uint64
}

// Name returns the name of the bucket.
func (b *Bucket) Name() string {
	return b.name
}

// Writable returns whether the bucket is writable.
func (b *Bucket) Writable() bool {
	return b.tx.writable
}

// Cursor creates a cursor associated with the bucket.
// The cursor is only valid as long as the transaction is open.
// Do not use a cursor after the transaction is closed.
func (b *Bucket) Cursor() *Cursor {
	return &Cursor{
		tx:    b.tx,
		root:  b.root,
		stack: make([]elemRef, 0),
	}
}

// Get retrieves the value for a key in the bucket.
// Returns a nil value if the key does not exist.
func (b *Bucket) Get(key []byte) []byte {
	c := b.Cursor()
	k, v := c.Seek(key)

	// If our target node isn't the same key as what's passed in then return nil.
	if !bytes.Equal(key, k) {
		return nil
	}
	return v
}

// Put sets the value for a key in the bucket.
// If the key exist then its previous value will be overwritten.
// Returns an error if the bucket was created from a read-only transaction, if the key is blank, if the key is too large, or if the value is too large.
func (b *Bucket) Put(key []byte, value []byte) error {
	if !b.Writable() {
		return ErrBucketNotWritable
	}

	// Validate the key and data size.
	if len(key) == 0 {
		return ErrKeyRequired
	} else if len(key) > MaxKeySize {
		return ErrKeyTooLarge
	} else if len(value) > MaxValueSize {
		return ErrValueTooLarge
	}

	// Move cursor to correct position.
	c := b.Cursor()
	c.Seek(key)

	// Insert the key/value.
	c.node(b.tx).put(key, key, value, 0)

	return nil
}

// Delete removes a key from the bucket.
// If the key does not exist then nothing is done and a nil error is returned.
// Returns an error if the bucket was created from a read-only transaction.
func (b *Bucket) Delete(key []byte) error {
	if !b.Writable() {
		return ErrBucketNotWritable
	}

	// Move cursor to correct position.
	c := b.Cursor()
	c.Seek(key)

	// Delete the node if we have a matching key.
	c.node(b.tx).del(key)

	return nil
}

// NextSequence returns an autoincrementing integer for the bucket.
func (b *Bucket) NextSequence() (int, error) {
	if !b.Writable() {
		return 0, ErrBucketNotWritable
	}

	// Make sure next sequence number will not be larger than the maximum
	// integer size of the system.
	if b.bucket.sequence == uint64(maxInt) {
		return 0, ErrSequenceOverflow
	}

	// Increment and return the sequence.
	b.bucket.sequence++

	return int(b.bucket.sequence), nil
}

// ForEach executes a function for each key/value pair in a bucket.
// If the provided function returns an error then the iteration is stopped and
// the error is returned to the caller.
func (b *Bucket) ForEach(fn func(k, v []byte) error) error {
	c := b.Cursor()
	for k, v := c.First(); k != nil; k, v = c.Next() {
		if err := fn(k, v); err != nil {
			return err
		}
	}
	return nil
}

// Stat returns stats on a bucket.
func (b *Bucket) Stat() *BucketStat {
	s := &BucketStat{}
	b.tx.forEachPage(b.root, 0, func(p *page, depth int) {
		if (p.flags & leafPageFlag) != 0 {
			s.LeafPageCount++
			s.KeyCount += int(p.count)
		} else if (p.flags & branchPageFlag) != 0 {
			s.BranchPageCount++
		}

		s.OverflowPageCount += int(p.overflow)

		if depth+1 > s.MaxDepth {
			s.MaxDepth = (depth + 1)
		}
	})
	return s
}

// BucketStat represents stats on a bucket such as branch pages and leaf pages.
type BucketStat struct {
	BranchPageCount   int
	LeafPageCount     int
	OverflowPageCount int
	KeyCount          int
	MaxDepth          int
}
<|endoftext|>


<commit_before>package git

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var logger = log.New(os.Stderr, "", 0)

type Repository struct {
	Origin    string
	LocalPath string
}

func (repo *Repository) cacheDir() string {
	return os.Getenv("HOME") + "/.dgtk/cache/git_repositories"
}

func (repo *Repository) Fetch() error {
	logger.Println("fetching origin")
	_, e := repo.executeGitCommand("fetch")
	return e
}

const githubPrefix = "git@github.com:"

func (repo *Repository) cachePath() string {
	name := repo.Name()
	if strings.HasPrefix(repo.Origin, githubPrefix) {
		name = strings.TrimPrefix(repo.Origin, githubPrefix)
		return repo.cacheDir() + "/github.com/" + name
	}
	return repo.cacheDir() + "/" + repo.Name()
}

func (repo *Repository) clone() error {
	logger.Printf("cloning %s into %s", repo.Origin, repo.cachePath())
	cmd := exec.Command("git", "clone", "--bare", repo.Origin, repo.cachePath())
	if b, e := cmd.CombinedOutput(); e != nil {
		logger.Printf("ERROR: %s", strings.TrimSpace(string(b)))
		return e
	}
	cmd = exec.Command("git", "--git-dir="+repo.cachePath(), "config", "remote.origin.fetch", "refs/heads/*:refs/heads/*")
	if b, e := cmd.CombinedOutput(); e != nil {
		logger.Printf("ERROR: %s", strings.TrimSpace(string(b)))
		return e
	}
	return nil
}

func (repo *Repository) Init() error {
	e := os.MkdirAll(repo.cacheDir(), 0755)
	if e != nil {
		return e
	}

	if !fileExists(repo.cachePath()) {
		if e := repo.clone(); e != nil {
			return e
		}
	} else {
		logger.Printf("already cloned %s to %s", repo.Origin, repo.cachePath())
	}
	return nil
}

func (repo *Repository) createGitCommand(gitCommand ...string) *exec.Cmd {
	return exec.Command("git", append([]string{"--git-dir=" + repo.cachePath()}, gitCommand...)...)
}

func (repo *Repository) executeGitCommand(gitCommand ...string) (b []byte, e error) {
	cmd := repo.createGitCommand(gitCommand...)
	b, e = cmd.CombinedOutput()
	if e != nil {
		logger.Printf("ERROR: %s (%v)", strings.TrimSpace(string(b)), cmd)
		return b, e
	}
	return b, nil
}

func (repo *Repository) MostRecentCommitFor(pattern string) (commit string, e error) {
	commits, e := repo.Commits(&CommitOptions{Limit: 1, Pattern: pattern})
	if e != nil {
		return "", e
	}
	if len(commits) == 0 {
		return "", e
	}
	return commits[0].Checksum, nil
}

var validTar = regexp.MustCompile("^([0-9a-f]{40})$")

// Writes tgz archive to the given tar writer.
func (repo *Repository) WriteArchiveToTar(revision string, w *tar.Writer) (e error) {
	if !validTar.MatchString(revision) {
		return fmt.Errorf("revision %q not valid (must be 40 digit git sha)", revision)
	}

	mtime, e := repo.DateOf(revision, ".")
	if e != nil {
		return e
	}

	e = repo.addArchiveToTar(revision, mtime, w)
	if e != nil {
		return e
	}

	return addFileToArchive("REVISION", []byte(revision), mtime, w)
}

func (repo *Repository) WriteFilesToTar(revision string, w *tar.Writer, files ...string) (e error) {
	if len(files) == 0 {
		return fmt.Errorf("empty file list given")
	}

	if !validTar.MatchString(revision) {
		return fmt.Errorf("revision %q not valid (must be 40 digit git sha)", revision)
	}

	for _, file := range files {
		mtime, e := repo.DateOf(revision, file)
		if e != nil {
			return e
		}

		buf, e := repo.getFileAtRevision(revision, file)
		if e != nil {
			return e
		}

		if e = addFileToArchive(file, buf, mtime, w); e != nil {
			return e
		}
	}

	return nil
}

func (repo *Repository) getFileAtRevision(revision, file string) (content []byte, e error) {
	buf := bytes.NewBuffer(nil)

	cmd := repo.createGitCommand("show", revision+":"+file)
	cmd.Stdout = buf

	if e = cmd.Run(); e != nil {
		return nil, e
	}

	return buf.Bytes(), nil
}

func addFileToArchive(name string, content []byte, modTime time.Time, w *tar.Writer) error {
	e := w.WriteHeader(&tar.Header{Name: name, Size: int64(len(content)), ModTime: modTime, Mode: 0644})
	if e != nil {
		return e
	}
	_, e = w.Write(content)
	return e
}

func (repo *Repository) addArchiveToTar(revision string, mtime time.Time, w *tar.Writer) (e error) {
	filename := repo.Name() + ".tar.gz"

	buf := bytes.NewBuffer(nil)

	cmd := repo.createGitCommand("archive", "--format=tar.gz", revision)
	cmd.Stdout = buf

	if e = cmd.Run(); e != nil {
		return e
	}

	e = w.WriteHeader(&tar.Header{Name: filename, Size: int64(buf.Len()), ModTime: mtime, Mode: 0644})
	if e != nil {
		return e
	}

	_, e = io.Copy(w, buf)
	return e
}

func (repo *Repository) Name() string {
	return strings.TrimSuffix(filepath.Base(repo.Origin), ".git")
}

func (repo *Repository) DateOf(revision, file string) (time.Time, error) {
	b, e := repo.executeGitCommand("log", "-1", "--format='%ct'", revision, "--", file)
	if e != nil {
		return time.Now(), e
	}
	d, e := strconv.Atoi(strings.Trim(string(b), "'\n"))
	if e != nil {
		return time.Now(), e
	}

	return time.Unix(int64(d), 0), nil
}

func (repo *Repository) Commits(options *CommitOptions) (commits []*Commit, e error) {
	if options == nil {
		options = &CommitOptions{Limit: 10}
	}
	path := repo.LocalPath
	if path == "" {
		path = repo.cachePath()
	}
	b, e := repo.executeGitCommand("log", "-n", strconv.Itoa(options.Limit), "--pretty=format:'%%H\t%%at\t%%s'", options.Pattern)
	if e != nil {
		return nil, e
	}
	lines := strings.Split(string(b), "\n")
	commits = make([]*Commit, 0, len(lines))

	for _, line := range lines {
		parts := strings.SplitN(line, "\t", 3)
		if len(parts) == 3 {
			if t, e := strconv.ParseInt(parts[1], 10, 64); e == nil {
				commits = append(commits, &Commit{Checksum: parts[0], AuthorDate: time.Unix(t, 0), Message: parts[2]})
			} else {
				logger.Printf("ERROR: %s", e.Error())
			}
		}
	}
	return commits, nil
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil || !os.IsNotExist(err)
}
<commit_msg>only add patterns to commits git command when present<commit_after>package git

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var logger = log.New(os.Stderr, "", 0)

type Repository struct {
	Origin    string
	LocalPath string
}

func (repo *Repository) cacheDir() string {
	return os.Getenv("HOME") + "/.dgtk/cache/git_repositories"
}

func (repo *Repository) Fetch() error {
	logger.Println("fetching origin")
	_, e := repo.executeGitCommand("fetch")
	return e
}

const githubPrefix = "git@github.com:"

func (repo *Repository) cachePath() string {
	name := repo.Name()
	if strings.HasPrefix(repo.Origin, githubPrefix) {
		name = strings.TrimPrefix(repo.Origin, githubPrefix)
		return repo.cacheDir() + "/github.com/" + name
	}
	return repo.cacheDir() + "/" + repo.Name()
}

func (repo *Repository) clone() error {
	logger.Printf("cloning %s into %s", repo.Origin, repo.cachePath())
	cmd := exec.Command("git", "clone", "--bare", repo.Origin, repo.cachePath())
	if b, e := cmd.CombinedOutput(); e != nil {
		logger.Printf("ERROR: %s", strings.TrimSpace(string(b)))
		return e
	}
	cmd = exec.Command("git", "--git-dir="+repo.cachePath(), "config", "remote.origin.fetch", "refs/heads/*:refs/heads/*")
	if b, e := cmd.CombinedOutput(); e != nil {
		logger.Printf("ERROR: %s", strings.TrimSpace(string(b)))
		return e
	}
	return nil
}

func (repo *Repository) Init() error {
	e := os.MkdirAll(repo.cacheDir(), 0755)
	if e != nil {
		return e
	}

	if !fileExists(repo.cachePath()) {
		if e := repo.clone(); e != nil {
			return e
		}
	} else {
		logger.Printf("already cloned %s to %s", repo.Origin, repo.cachePath())
	}
	return nil
}

func (repo *Repository) createGitCommand(gitCommand ...string) *exec.Cmd {
	return exec.Command("git", append([]string{"--git-dir=" + repo.cachePath()}, gitCommand...)...)
}

func (repo *Repository) executeGitCommand(gitCommand ...string) (b []byte, e error) {
	cmd := repo.createGitCommand(gitCommand...)
	b, e = cmd.CombinedOutput()
	if e != nil {
		logger.Printf("ERROR: %s (%v)", strings.TrimSpace(string(b)), cmd)
		return b, e
	}
	return b, nil
}

func (repo *Repository) MostRecentCommitFor(pattern string) (commit string, e error) {
	commits, e := repo.Commits(&CommitOptions{Limit: 1, Pattern: pattern})
	if e != nil {
		return "", e
	}
	if len(commits) == 0 {
		return "", e
	}
	return commits[0].Checksum, nil
}

var validTar = regexp.MustCompile("^([0-9a-f]{40})$")

// Writes tgz archive to the given tar writer.
func (repo *Repository) WriteArchiveToTar(revision string, w *tar.Writer) (e error) {
	if !validTar.MatchString(revision) {
		return fmt.Errorf("revision %q not valid (must be 40 digit git sha)", revision)
	}

	mtime, e := repo.DateOf(revision, ".")
	if e != nil {
		return e
	}

	e = repo.addArchiveToTar(revision, mtime, w)
	if e != nil {
		return e
	}

	return addFileToArchive("REVISION", []byte(revision), mtime, w)
}

func (repo *Repository) WriteFilesToTar(revision string, w *tar.Writer, files ...string) (e error) {
	if len(files) == 0 {
		return fmt.Errorf("empty file list given")
	}

	if !validTar.MatchString(revision) {
		return fmt.Errorf("revision %q not valid (must be 40 digit git sha)", revision)
	}

	for _, file := range files {
		mtime, e := repo.DateOf(revision, file)
		if e != nil {
			return e
		}

		buf, e := repo.getFileAtRevision(revision, file)
		if e != nil {
			return e
		}

		if e = addFileToArchive(file, buf, mtime, w); e != nil {
			return e
		}
	}

	return nil
}

func (repo *Repository) getFileAtRevision(revision, file string) (content []byte, e error) {
	buf := bytes.NewBuffer(nil)

	cmd := repo.createGitCommand("show", revision+":"+file)
	cmd.Stdout = buf

	if e = cmd.Run(); e != nil {
		return nil, e
	}

	return buf.Bytes(), nil
}

func addFileToArchive(name string, content []byte, modTime time.Time, w *tar.Writer) error {
	e := w.WriteHeader(&tar.Header{Name: name, Size: int64(len(content)), ModTime: modTime, Mode: 0644})
	if e != nil {
		return e
	}
	_, e = w.Write(content)
	return e
}

func (repo *Repository) addArchiveToTar(revision string, mtime time.Time, w *tar.Writer) (e error) {
	filename := repo.Name() + ".tar.gz"

	buf := bytes.NewBuffer(nil)

	cmd := repo.createGitCommand("archive", "--format=tar.gz", revision)
	cmd.Stdout = buf

	if e = cmd.Run(); e != nil {
		return e
	}

	e = w.WriteHeader(&tar.Header{Name: filename, Size: int64(buf.Len()), ModTime: mtime, Mode: 0644})
	if e != nil {
		return e
	}

	_, e = io.Copy(w, buf)
	return e
}

func (repo *Repository) Name() string {
	return strings.TrimSuffix(filepath.Base(repo.Origin), ".git")
}

func (repo *Repository) DateOf(revision, file string) (time.Time, error) {
	b, e := repo.executeGitCommand("log", "-1", "--format='%ct'", revision, "--", file)
	if e != nil {
		return time.Now(), e
	}
	d, e := strconv.Atoi(strings.Trim(string(b), "'\n"))
	if e != nil {
		return time.Now(), e
	}

	return time.Unix(int64(d), 0), nil
}

func (repo *Repository) Commits(options *CommitOptions) (commits []*Commit, e error) {
	if options == nil {
		options = &CommitOptions{Limit: 10}
	}
	path := repo.LocalPath
	if path == "" {
		path = repo.cachePath()
	}
	parts := []string{"log", "-n", strconv.Itoa(options.Limit), "--pretty=format:%H\t%at\t%s"}
	if options.Pattern != "" {
		parts = append(parts, options.Pattern)
	}
	b, e := repo.executeGitCommand(parts...)
	if e != nil {
		return nil, e
	}
	lines := strings.Split(string(b), "\n")
	commits = make([]*Commit, 0, len(lines))

	for _, line := range lines {
		parts := strings.SplitN(line, "\t", 3)
		if len(parts) == 3 {
			if t, e := strconv.ParseInt(parts[1], 10, 64); e == nil {
				commits = append(commits, &Commit{Checksum: parts[0], AuthorDate: time.Unix(t, 0), Message: parts[2]})
			} else {
				logger.Printf("ERROR: %s", e.Error())
			}
		}
	}
	return commits, nil
}

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil || !os.IsNotExist(err)
}
<|endoftext|>


<commit_before>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 Julien Schmidt. All rights reserved.
// http://www.julienschmidt.com
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import "io"

const defaultBufSize = 4096

// A read buffer similar to bufio.Reader but zero-copy-ish
// Also highly optimized for this particular use case.
type buffer struct {
	buf    []byte
	rd     io.Reader
	idx    int
	length int
}

func newBuffer(rd io.Reader) *buffer {
	return &buffer{
		buf: make([]byte, defaultBufSize),
		rd:  rd,
	}
}

// fill reads into the buffer until at least _need_ bytes are in it
func (b *buffer) fill(need int) (err error) {
	// move existing data to the beginning
	if b.length > 0 && b.idx > 0 {
		copy(b.buf[0:b.length], b.buf[b.idx:])
	}

	// grow buffer if necessary
	if need > len(b.buf) {
		for {
			b.buf = append(b.buf, 0)
			b.buf = b.buf[:cap(b.buf)]

			if cap(b.buf) < need {
				continue
			}
			break
		}
	}

	b.idx = 0

	var n int
	for {
		n, err = b.rd.Read(b.buf[b.length:])
		b.length += n

		if b.length < need && err == nil {
			continue
		}
		return // err
	}
}

// returns next N bytes from buffer.
// The returned slice is only guaranteed to be valid until the next read
func (b *buffer) readNext(need int) (p []byte, err error) {
	if b.length < need {
		// refill
		err = b.fill(need) // err deferred
	}

	p = b.buf[b.idx : b.idx+need]
	b.idx += need
	b.length -= need
	return
}
<commit_msg>fix build with Go 1.0 + branch optimization<commit_after>// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
//
// Copyright 2013 Julien Schmidt. All rights reserved.
// http://www.julienschmidt.com
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package mysql

import "io"

const defaultBufSize = 4096

// A read buffer similar to bufio.Reader but zero-copy-ish
// Also highly optimized for this particular use case.
type buffer struct {
	buf    []byte
	rd     io.Reader
	idx    int
	length int
}

func newBuffer(rd io.Reader) *buffer {
	return &buffer{
		buf: make([]byte, defaultBufSize),
		rd:  rd,
	}
}

// fill reads into the buffer until at least _need_ bytes are in it
func (b *buffer) fill(need int) (err error) {
	// move existing data to the beginning
	if b.length > 0 && b.idx > 0 {
		copy(b.buf[0:b.length], b.buf[b.idx:])
	}

	// grow buffer if necessary
	if need > len(b.buf) {
		for {
			b.buf = append(b.buf, 0)
			b.buf = b.buf[:cap(b.buf)]

			if cap(b.buf) >= need {
				break
			}
		}
	}

	b.idx = 0

	var n int
	for {
		n, err = b.rd.Read(b.buf[b.length:])
		b.length += n

		if b.length < need && err == nil {
			continue
		}
		return // err
	}
	return
}

// returns next N bytes from buffer.
// The returned slice is only guaranteed to be valid until the next read
func (b *buffer) readNext(need int) (p []byte, err error) {
	if b.length < need {
		// refill
		err = b.fill(need) // err deferred
	}

	p = b.buf[b.idx : b.idx+need]
	b.idx += need
	b.length -= need
	return
}
<|endoftext|>


<commit_before>package broker

import (
	"context"
	"fmt"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/messaging_pb"
	"github.com/chrislusf/seaweedfs/weed/security"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (broker *MessageBroker) appendToFile(targetFile string, topicConfig *messaging_pb.TopicConfiguration, data []byte) error {

	assignResult, uploadResult, err2 := broker.assignAndUpload(topicConfig, data)
	if err2 != nil {
		return err2
	}

	dir, name := util.FullPath(targetFile).DirAndName()

	chunk := &filer_pb.FileChunk{
		FileId:    assignResult.Fid,
		Offset:    0, // needs to be fixed during appending
		Size:      uint64(uploadResult.Size),
		Mtime:     time.Now().UnixNano(),
		ETag:      uploadResult.ETag,
		IsGzipped: uploadResult.Gzip > 0,
	}

	// append the chunk
	if err := broker.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		request := &filer_pb.AppendToEntryRequest{
			Directory: dir,
			EntryName: name,
			Chunks:    []*filer_pb.FileChunk{chunk},
		}

		_, err := client.AppendToEntry(context.Background(), request)
		if err != nil {
			glog.V(0).Infof("append to file %v: %v", request, err)
			return err
		}

		return nil
	}); err != nil {
		return fmt.Errorf("append to file %v: %v", targetFile, err)
	}

	return nil
}

func (broker *MessageBroker) assignAndUpload(topicConfig *messaging_pb.TopicConfiguration, data []byte) (*operation.AssignResult, *operation.UploadResult, error) {

	var assignResult = &operation.AssignResult{}

	// assign a volume location
	if err := broker.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		request := &filer_pb.AssignVolumeRequest{
			Count:       1,
			Replication: topicConfig.Replication,
			Collection:  topicConfig.Collection,
		}

		resp, err := client.AssignVolume(context.Background(), request)
		if err != nil {
			glog.V(0).Infof("assign volume failure %v: %v", request, err)
			return err
		}
		if resp.Error != "" {
			return fmt.Errorf("assign volume failure %v: %v", request, resp.Error)
		}

		assignResult.Auth = security.EncodedJwt(resp.Auth)
		assignResult.Fid = resp.FileId
		assignResult.Url = resp.Url
		assignResult.PublicUrl = resp.PublicUrl
		assignResult.Count = uint64(resp.Count)

		return nil
	}); err != nil {
		return nil, nil, err
	}

	// upload data
	targetUrl := fmt.Sprintf("http://%s/%s", assignResult.Url, assignResult.Fid)
	uploadResult, err := operation.UploadData(targetUrl, "", broker.option.Cipher, data, false, "", nil, assignResult.Auth)
	if err != nil {
		return nil, nil, fmt.Errorf("upload data %s: %v", targetUrl, err)
	}
	// println("uploaded to", targetUrl)
	return assignResult, uploadResult, nil
}

var _ = filer_pb.FilerClient(&MessageBroker{})

func (broker *MessageBroker) WithFilerClient(fn func(filer_pb.SeaweedFilerClient) error) (err error) {

	for _, filer := range broker.option.Filers {
		if err = pb.WithFilerClient(filer, broker.grpcDialOption, fn); err != nil {
			glog.V(0).Infof("fail to connect to %s: %v", filer, err)
		} else {
			break
		}
	}

	return

}

func (broker *MessageBroker) AdjustedUrl(hostAndPort string) string {
	return hostAndPort
}
<commit_msg>add cipher key<commit_after>package broker

import (
	"context"
	"fmt"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/messaging_pb"
	"github.com/chrislusf/seaweedfs/weed/security"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (broker *MessageBroker) appendToFile(targetFile string, topicConfig *messaging_pb.TopicConfiguration, data []byte) error {

	assignResult, uploadResult, err2 := broker.assignAndUpload(topicConfig, data)
	if err2 != nil {
		return err2
	}

	dir, name := util.FullPath(targetFile).DirAndName()

	chunk := &filer_pb.FileChunk{
		FileId:    assignResult.Fid,
		Offset:    0, // needs to be fixed during appending
		Size:      uint64(uploadResult.Size),
		Mtime:     time.Now().UnixNano(),
		ETag:      uploadResult.ETag,
		CipherKey: uploadResult.CipherKey,
		IsGzipped: uploadResult.Gzip > 0,
	}

	// append the chunk
	if err := broker.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		request := &filer_pb.AppendToEntryRequest{
			Directory: dir,
			EntryName: name,
			Chunks:    []*filer_pb.FileChunk{chunk},
		}

		_, err := client.AppendToEntry(context.Background(), request)
		if err != nil {
			glog.V(0).Infof("append to file %v: %v", request, err)
			return err
		}

		return nil
	}); err != nil {
		return fmt.Errorf("append to file %v: %v", targetFile, err)
	}

	return nil
}

func (broker *MessageBroker) assignAndUpload(topicConfig *messaging_pb.TopicConfiguration, data []byte) (*operation.AssignResult, *operation.UploadResult, error) {

	var assignResult = &operation.AssignResult{}

	// assign a volume location
	if err := broker.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		request := &filer_pb.AssignVolumeRequest{
			Count:       1,
			Replication: topicConfig.Replication,
			Collection:  topicConfig.Collection,
		}

		resp, err := client.AssignVolume(context.Background(), request)
		if err != nil {
			glog.V(0).Infof("assign volume failure %v: %v", request, err)
			return err
		}
		if resp.Error != "" {
			return fmt.Errorf("assign volume failure %v: %v", request, resp.Error)
		}

		assignResult.Auth = security.EncodedJwt(resp.Auth)
		assignResult.Fid = resp.FileId
		assignResult.Url = resp.Url
		assignResult.PublicUrl = resp.PublicUrl
		assignResult.Count = uint64(resp.Count)

		return nil
	}); err != nil {
		return nil, nil, err
	}

	// upload data
	targetUrl := fmt.Sprintf("http://%s/%s", assignResult.Url, assignResult.Fid)
	uploadResult, err := operation.UploadData(targetUrl, "", broker.option.Cipher, data, false, "", nil, assignResult.Auth)
	if err != nil {
		return nil, nil, fmt.Errorf("upload data %s: %v", targetUrl, err)
	}
	// println("uploaded to", targetUrl)
	return assignResult, uploadResult, nil
}

var _ = filer_pb.FilerClient(&MessageBroker{})

func (broker *MessageBroker) WithFilerClient(fn func(filer_pb.SeaweedFilerClient) error) (err error) {

	for _, filer := range broker.option.Filers {
		if err = pb.WithFilerClient(filer, broker.grpcDialOption, fn); err != nil {
			glog.V(0).Infof("fail to connect to %s: %v", filer, err)
		} else {
			break
		}
	}

	return

}

func (broker *MessageBroker) AdjustedUrl(hostAndPort string) string {
	return hostAndPort
}
<|endoftext|>