<commit_before>package gomigrate

import (
	"path/filepath"
	"regexp"
	"strconv"
)

var (
	upMigrationFile   = regexp.MustCompile(`(\d+)_(\w+)_up\.sql`)
	downMigrationFile = regexp.MustCompile(`(\d+)_(\w+)_down\.sql`)
	subMigrationSplit = regexp.MustCompile(`;\s*\\n`)
	allWhitespace     = regexp.MustCompile(`^\s*$`)
)

// Returns the migration number, type and base name, so 1, "up", "migration" from "01_migration_up.sql"
func parseMigrationPath(path string) (uint64, migrationType, string, error) {
	filebase := filepath.Base(path)

	matches := upMigrationFile.FindAllSubmatch([]byte(filebase), -1)
	if matches != nil {
		return parseMatches(matches, upMigration)
	}
	matches = downMigrationFile.FindAllSubmatch([]byte(filebase), -1)
	if matches != nil {
		return parseMatches(matches, downMigration)
	}

	return 0, "", "", InvalidMigrationFile
}

// Parses matches given by a migration file regex.
func parseMatches(matches [][][]byte, mType migrationType) (uint64, migrationType, string, error) {
	num := matches[0][1]
	name := matches[0][2]
	parsedNum, err := strconv.ParseUint(string(num), 10, 64)
	if err != nil {
		return 0, "", "", err
	}
	return parsedNum, mType, string(name), nil
}

// Splits migration sql into different strings separated by a semi-colon.
func splitMigrationString(sql string) []string {
	return subMigrationSplit.Split(sql, -1)
}

// This type is used to sort migration ids.
type uint64slice []uint64

func (u uint64slice) Len() int {
	return len(u)
}

func (u uint64slice) Less(a, b int) bool {
	return u[a] < u[b]
}

func (u uint64slice) Swap(a, b int) {
	tempA := u[a]
	u[a] = u[b]
	u[b] = tempA
}
<commit_msg>Be less strict<commit_after>package gomigrate

import (
	"path/filepath"
	"regexp"
	"strconv"
)

var (
	upMigrationFile   = regexp.MustCompile(`(\d+)_(\w+)_up\.sql`)
	downMigrationFile = regexp.MustCompile(`(\d+)_(\w+)_down\.sql`)
	subMigrationSplit = regexp.MustCompile(`;\s*`)
	allWhitespace     = regexp.MustCompile(`^\s*$`)
)

// Returns the migration number, type and base name, so 1, "up", "migration" from "01_migration_up.sql"
func parseMigrationPath(path string) (uint64, migrationType, string, error) {
	filebase := filepath.Base(path)

	matches := upMigrationFile.FindAllSubmatch([]byte(filebase), -1)
	if matches != nil {
		return parseMatches(matches, upMigration)
	}
	matches = downMigrationFile.FindAllSubmatch([]byte(filebase), -1)
	if matches != nil {
		return parseMatches(matches, downMigration)
	}

	return 0, "", "", InvalidMigrationFile
}

// Parses matches given by a migration file regex.
func parseMatches(matches [][][]byte, mType migrationType) (uint64, migrationType, string, error) {
	num := matches[0][1]
	name := matches[0][2]
	parsedNum, err := strconv.ParseUint(string(num), 10, 64)
	if err != nil {
		return 0, "", "", err
	}
	return parsedNum, mType, string(name), nil
}

// Splits migration sql into different strings separated by a semi-colon.
func splitMigrationString(sql string) []string {
	return subMigrationSplit.Split(sql, -1)
}

// This type is used to sort migration ids.
type uint64slice []uint64

func (u uint64slice) Len() int {
	return len(u)
}

func (u uint64slice) Less(a, b int) bool {
	return u[a] < u[b]
}

func (u uint64slice) Swap(a, b int) {
	tempA := u[a]
	u[a] = u[b]
	u[b] = tempA
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"runtime"
	"time"
)

const MEGABYTE = 1024 * 1024

func startRuntimeProfiler() {
	m := &runtime.MemStats{}

	for {
		runtime.ReadMemStats(m)

		fmt.Println("-----------------------")
		fmt.Println("Goroutines:", runtime.NumGoroutine())
		fmt.Println("Memory acquired:", m.Sys, "bytes,", m.Sys/MEGABYTE, "mb")
		fmt.Println("Memory used:", m.Alloc, "bytes,", m.Alloc/MEGABYTE, "mb")

		time.Sleep(time.Minute)
	}
}
<commit_msg>Set profiler sleep time to 30 seconds<commit_after>package main

import (
	"fmt"
	"runtime"
	"time"
)

const MEGABYTE = 1024 * 1024

func startRuntimeProfiler() {
	m := &runtime.MemStats{}

	for {
		runtime.ReadMemStats(m)

		fmt.Println("-----------------------")
		fmt.Println("Goroutines:", runtime.NumGoroutine())
		fmt.Println("Memory acquired:", m.Sys, "bytes,", m.Sys/MEGABYTE, "mb")
		fmt.Println("Memory used:", m.Alloc, "bytes,", m.Alloc/MEGABYTE, "mb")

		time.Sleep(time.Second * 30)
	}
}
<|endoftext|>


<commit_before>package mpfluentd

import (
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.fluentd")

// FluentdMetrics plugin for fluentd
type FluentdMetrics struct {
	Target          string
	Tempfile        string
	pluginType      string
	pluginIDPattern *regexp.Regexp
	extendedMetrics []string

	plugins []FluentdPluginMetrics
}

// FluentdPluginMetrics metrics
type FluentdPluginMetrics struct {
	RetryCount            uint64 `json:"retry_count"`
	BufferQueueLength     uint64 `json:"buffer_queue_length"`
	BufferTotalQueuedSize uint64 `json:"buffer_total_queued_size"`
	OutputPlugin          bool   `json:"output_plugin"`
	Type                  string `json:"type"`
	PluginCategory        string `json:"plugin_category"`
	PluginID              string `json:"plugin_id"`
	normalizedPluginID    string

	// extended metrics fluentd >= 1.6
	// https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released
	EmitRecords                      uint64  `json:"emit_records"`
	EmitCount                        uint64  `json:"emit_count"`
	WriteCount                       uint64  `json:"write_count"`
	RollbackCount                    uint64  `json:"rollback_count"`
	SlowFlushCount                   uint64  `json:"slow_flush_count"`
	FlushTimeCount                   uint64  `json:"flush_time_count"`
	BufferStageLength                uint64  `json:"buffer_stage_length"`
	BufferStageByteSize              uint64  `json:"buffer_stage_byte_size"`
	BufferQueueByteSize              uint64  `json:"buffer_queue_byte_size"`
	BufferAvailableBufferSpaceRatios float64 `json:"buffer_available_buffer_space_ratios"`
}

func (f FluentdPluginMetrics) getExtended(name string) float64 {
	switch name {
	case "emit_records":
		return float64(f.EmitRecords)
	case "emit_count":
		return float64(f.EmitCount)
	case "write_count":
		return float64(f.WriteCount)
	case "rollback_count":
		return float64(f.RollbackCount)
	case "slow_flush_count":
		return float64(f.SlowFlushCount)
	case "flush_time_count":
		return float64(f.FlushTimeCount)
	case "buffer_stage_length":
		return float64(f.BufferStageLength)
	case "buffer_stage_byte_size":
		return float64(f.BufferStageByteSize)
	case "buffer_queue_byte_size":
		return float64(f.BufferQueueByteSize)
	case "buffer_available_buffer_space_ratios":
		return f.BufferAvailableBufferSpaceRatios
	default:
		logger.Warningf("extended-metrics %s not defined", name)
	}
	return 0
}

// FluentMonitorJSON monitor json
type FluentMonitorJSON struct {
	Plugins []FluentdPluginMetrics `json:"plugins"`
}

var normalizePluginIDRe = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizePluginID(in string) string {
	return normalizePluginIDRe.ReplaceAllString(in, "_")
}

func (fpm FluentdPluginMetrics) getNormalizedPluginID() string {
	if fpm.normalizedPluginID == "" {
		fpm.normalizedPluginID = normalizePluginID(fpm.PluginID)
	}
	return fpm.normalizedPluginID
}

func (f *FluentdMetrics) parseStats(body []byte) (map[string]interface{}, error) {
	var j FluentMonitorJSON
	err := json.Unmarshal(body, &j)
	f.plugins = j.Plugins

	metrics := make(map[string]interface{})
	for _, p := range f.plugins {
		if f.nonTargetPlugin(p) {
			continue
		}
		pid := p.getNormalizedPluginID()
		metrics["fluentd.retry_count."+pid] = float64(p.RetryCount)
		metrics["fluentd.buffer_queue_length."+pid] = float64(p.BufferQueueLength)
		metrics["fluentd.buffer_total_queued_size."+pid] = float64(p.BufferTotalQueuedSize)
		for _, name := range f.extendedMetrics {
			key := strings.Join([]string{"fluentd", name, pid}, ".")
			metrics[key] = p.getExtended(name)
		}
	}
	return metrics, err
}

func (f *FluentdMetrics) nonTargetPlugin(plugin FluentdPluginMetrics) bool {
	if plugin.PluginCategory != "output" {
		return true
	}
	if f.pluginType != "" && f.pluginType != plugin.Type {
		return true
	}
	if f.pluginIDPattern != nil && !f.pluginIDPattern.MatchString(plugin.PluginID) {
		return true
	}
	return false
}

// FetchMetrics interface for mackerelplugin
func (f FluentdMetrics) FetchMetrics() (map[string]interface{}, error) {
	req, err := http.NewRequest(http.MethodGet, f.Target, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "mackerel-plugin-fluentd")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return f.parseStats(body)
}

var defaultGraphs = map[string]mp.Graphs{
	"fluentd.retry_count": {
		Label: "Fluentd retry count",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_queue_length": {
		Label: "Fluentd queue length",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_total_queued_size": {
		Label: "Fluentd buffer total queued size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

var extendedGraphs = map[string]mp.Graphs{
	"fluentd.emit_records": {
		Label: "Fluentd emitted records",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.emit_count": {
		Label: "Fluentd emit calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.write_count": {
		Label: "Fluentd write/try_write calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.rollback_count": {
		Label: "Fluentd rollbacks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.slow_flush_count": {
		Label: "Fluentd slow flushes",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.flush_time_count": {
		Label: "Fluentd buffer flush time in msec",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.buffer_stage_length": {
		Label: "Fluentd length of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_stage_byte_size": {
		Label: "Fluentd bytesize of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_queue_byte_size": {
		Label: "Fluentd bytesize of queued buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_available_buffer_space_ratios": {
		Label: "Fluentd available space for buffer",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

// GraphDefinition interface for mackerelplugin
func (f FluentdMetrics) GraphDefinition() map[string]mp.Graphs {
	graphs := make(map[string]mp.Graphs, len(defaultGraphs))
	for key, g := range defaultGraphs {
		g := g
		graphs[key] = g
	}
	for _, name := range f.extendedMetrics {
		if g, ok := extendedGraphs["fluentd."+name]; ok {
			graphs["fluentd."+name] = g
		}
	}
	return graphs
}

// Do the plugin
func Do() {
	host := flag.String("host", "localhost", "fluentd monitor_agent host")
	port := flag.String("port", "24220", "fluentd monitor_agent port")
	pluginType := flag.String("plugin-type", "", "Gets the metric that matches this plugin type")
	pluginIDPatternString := flag.String("plugin-id-pattern", "", "Gets the metric that matches this plugin id pattern")
	tempFile := flag.String("tempfile", "", "Temp file name")
	extendedMetricNames := flag.String("extended_metrics", "", "extended metric names joind with ',' or 'all' (fluentd >= v1.6.0)")
	flag.Parse()

	var pluginIDPattern *regexp.Regexp
	var err error
	if *pluginIDPatternString != "" {
		pluginIDPattern, err = regexp.Compile(*pluginIDPatternString)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to exec mackerel-plugin-fluentd: invalid plugin-id-pattern: %s\n", err)
			os.Exit(1)
		}
	}

	var extendedMetrics []string
	switch *extendedMetricNames {
	case "all":
		for key := range extendedGraphs {
			extendedMetrics = append(extendedMetrics, strings.TrimPrefix(key, "fluentd."))
		}
	case "":
	default:
		extendedMetrics = strings.Split(*extendedMetricNames, ",")
	}
	f := FluentdMetrics{
		Target:          fmt.Sprintf("http://%s:%s/api/plugins.json", *host, *port),
		Tempfile:        *tempFile,
		pluginType:      *pluginType,
		pluginIDPattern: pluginIDPattern,
		extendedMetrics: extendedMetrics,
	}

	helper := mp.NewMackerelPlugin(f)

	helper.Tempfile = *tempFile
	if *tempFile == "" {
		tempFileSuffix := []string{*host, *port}
		if *pluginType != "" {
			tempFileSuffix = append(tempFileSuffix, *pluginType)
		}
		if *pluginIDPatternString != "" {
			tempFileSuffix = append(tempFileSuffix, fmt.Sprintf("%x", md5.Sum([]byte(*pluginIDPatternString))))
		}
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-fluentd-%s", strings.Join(tempFileSuffix, "-")))
	}

	helper.Run()
}
<commit_msg>fix receiver name<commit_after>package mpfluentd

import (
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.fluentd")

// FluentdMetrics plugin for fluentd
type FluentdMetrics struct {
	Target          string
	Tempfile        string
	pluginType      string
	pluginIDPattern *regexp.Regexp
	extendedMetrics []string

	plugins []FluentdPluginMetrics
}

// FluentdPluginMetrics metrics
type FluentdPluginMetrics struct {
	RetryCount            uint64 `json:"retry_count"`
	BufferQueueLength     uint64 `json:"buffer_queue_length"`
	BufferTotalQueuedSize uint64 `json:"buffer_total_queued_size"`
	OutputPlugin          bool   `json:"output_plugin"`
	Type                  string `json:"type"`
	PluginCategory        string `json:"plugin_category"`
	PluginID              string `json:"plugin_id"`
	normalizedPluginID    string

	// extended metrics fluentd >= 1.6
	// https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released
	EmitRecords                      uint64  `json:"emit_records"`
	EmitCount                        uint64  `json:"emit_count"`
	WriteCount                       uint64  `json:"write_count"`
	RollbackCount                    uint64  `json:"rollback_count"`
	SlowFlushCount                   uint64  `json:"slow_flush_count"`
	FlushTimeCount                   uint64  `json:"flush_time_count"`
	BufferStageLength                uint64  `json:"buffer_stage_length"`
	BufferStageByteSize              uint64  `json:"buffer_stage_byte_size"`
	BufferQueueByteSize              uint64  `json:"buffer_queue_byte_size"`
	BufferAvailableBufferSpaceRatios float64 `json:"buffer_available_buffer_space_ratios"`
}

func (fpm FluentdPluginMetrics) getExtended(name string) float64 {
	switch name {
	case "emit_records":
		return float64(f.EmitRecords)
	case "emit_count":
		return float64(f.EmitCount)
	case "write_count":
		return float64(f.WriteCount)
	case "rollback_count":
		return float64(f.RollbackCount)
	case "slow_flush_count":
		return float64(f.SlowFlushCount)
	case "flush_time_count":
		return float64(f.FlushTimeCount)
	case "buffer_stage_length":
		return float64(f.BufferStageLength)
	case "buffer_stage_byte_size":
		return float64(f.BufferStageByteSize)
	case "buffer_queue_byte_size":
		return float64(f.BufferQueueByteSize)
	case "buffer_available_buffer_space_ratios":
		return f.BufferAvailableBufferSpaceRatios
	default:
		logger.Warningf("extended-metrics %s not defined", name)
	}
	return 0
}

// FluentMonitorJSON monitor json
type FluentMonitorJSON struct {
	Plugins []FluentdPluginMetrics `json:"plugins"`
}

var normalizePluginIDRe = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizePluginID(in string) string {
	return normalizePluginIDRe.ReplaceAllString(in, "_")
}

func (fpm FluentdPluginMetrics) getNormalizedPluginID() string {
	if fpm.normalizedPluginID == "" {
		fpm.normalizedPluginID = normalizePluginID(fpm.PluginID)
	}
	return fpm.normalizedPluginID
}

func (f *FluentdMetrics) parseStats(body []byte) (map[string]interface{}, error) {
	var j FluentMonitorJSON
	err := json.Unmarshal(body, &j)
	f.plugins = j.Plugins

	metrics := make(map[string]interface{})
	for _, p := range f.plugins {
		if f.nonTargetPlugin(p) {
			continue
		}
		pid := p.getNormalizedPluginID()
		metrics["fluentd.retry_count."+pid] = float64(p.RetryCount)
		metrics["fluentd.buffer_queue_length."+pid] = float64(p.BufferQueueLength)
		metrics["fluentd.buffer_total_queued_size."+pid] = float64(p.BufferTotalQueuedSize)
		for _, name := range f.extendedMetrics {
			key := strings.Join([]string{"fluentd", name, pid}, ".")
			metrics[key] = p.getExtended(name)
		}
	}
	return metrics, err
}

func (f *FluentdMetrics) nonTargetPlugin(plugin FluentdPluginMetrics) bool {
	if plugin.PluginCategory != "output" {
		return true
	}
	if f.pluginType != "" && f.pluginType != plugin.Type {
		return true
	}
	if f.pluginIDPattern != nil && !f.pluginIDPattern.MatchString(plugin.PluginID) {
		return true
	}
	return false
}

// FetchMetrics interface for mackerelplugin
func (f FluentdMetrics) FetchMetrics() (map[string]interface{}, error) {
	req, err := http.NewRequest(http.MethodGet, f.Target, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "mackerel-plugin-fluentd")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return f.parseStats(body)
}

var defaultGraphs = map[string]mp.Graphs{
	"fluentd.retry_count": {
		Label: "Fluentd retry count",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_queue_length": {
		Label: "Fluentd queue length",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_total_queued_size": {
		Label: "Fluentd buffer total queued size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

var extendedGraphs = map[string]mp.Graphs{
	"fluentd.emit_records": {
		Label: "Fluentd emitted records",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.emit_count": {
		Label: "Fluentd emit calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.write_count": {
		Label: "Fluentd write/try_write calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.rollback_count": {
		Label: "Fluentd rollbacks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.slow_flush_count": {
		Label: "Fluentd slow flushes",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.flush_time_count": {
		Label: "Fluentd buffer flush time in msec",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.buffer_stage_length": {
		Label: "Fluentd length of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_stage_byte_size": {
		Label: "Fluentd bytesize of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_queue_byte_size": {
		Label: "Fluentd bytesize of queued buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_available_buffer_space_ratios": {
		Label: "Fluentd available space for buffer",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

// GraphDefinition interface for mackerelplugin
func (f FluentdMetrics) GraphDefinition() map[string]mp.Graphs {
	graphs := make(map[string]mp.Graphs, len(defaultGraphs))
	for key, g := range defaultGraphs {
		g := g
		graphs[key] = g
	}
	for _, name := range f.extendedMetrics {
		if g, ok := extendedGraphs["fluentd."+name]; ok {
			graphs["fluentd."+name] = g
		}
	}
	return graphs
}

// Do the plugin
func Do() {
	host := flag.String("host", "localhost", "fluentd monitor_agent host")
	port := flag.String("port", "24220", "fluentd monitor_agent port")
	pluginType := flag.String("plugin-type", "", "Gets the metric that matches this plugin type")
	pluginIDPatternString := flag.String("plugin-id-pattern", "", "Gets the metric that matches this plugin id pattern")
	tempFile := flag.String("tempfile", "", "Temp file name")
	extendedMetricNames := flag.String("extended_metrics", "", "extended metric names joind with ',' or 'all' (fluentd >= v1.6.0)")
	flag.Parse()

	var pluginIDPattern *regexp.Regexp
	var err error
	if *pluginIDPatternString != "" {
		pluginIDPattern, err = regexp.Compile(*pluginIDPatternString)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to exec mackerel-plugin-fluentd: invalid plugin-id-pattern: %s\n", err)
			os.Exit(1)
		}
	}

	var extendedMetrics []string
	switch *extendedMetricNames {
	case "all":
		for key := range extendedGraphs {
			extendedMetrics = append(extendedMetrics, strings.TrimPrefix(key, "fluentd."))
		}
	case "":
	default:
		extendedMetrics = strings.Split(*extendedMetricNames, ",")
	}
	f := FluentdMetrics{
		Target:          fmt.Sprintf("http://%s:%s/api/plugins.json", *host, *port),
		Tempfile:        *tempFile,
		pluginType:      *pluginType,
		pluginIDPattern: pluginIDPattern,
		extendedMetrics: extendedMetrics,
	}

	helper := mp.NewMackerelPlugin(f)

	helper.Tempfile = *tempFile
	if *tempFile == "" {
		tempFileSuffix := []string{*host, *port}
		if *pluginType != "" {
			tempFileSuffix = append(tempFileSuffix, *pluginType)
		}
		if *pluginIDPatternString != "" {
			tempFileSuffix = append(tempFileSuffix, fmt.Sprintf("%x", md5.Sum([]byte(*pluginIDPatternString))))
		}
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-fluentd-%s", strings.Join(tempFileSuffix, "-")))
	}

	helper.Run()
}
<|endoftext|>


<commit_before>// http://codingdojo.org/cgi-bin/wiki.pl?back=KataBankOCR
package main

import (
	"fmt"
	"log"
	"strings"
)

func assert(t bool, note string) {
	if !t {
		log.Fatalf("%s FAILED !", note)
	}
}

type bigdigit [3]string

/* GLOBAL VARIABLES */

var zero = bigdigit{
	" _ ",
	"| |",
	"|_|",
}

var one = bigdigit{
	"   ",
	"  |",
	"  |",
}

var two = bigdigit{
	" _ ",
	" _|",
	"|_ ",
}

var three = bigdigit{
	" _ ",
	" _|",
	" _|",
}

var four = bigdigit{
	"   ",
	"|_|",
	"  |",
}

var five = bigdigit{
	" _ ",
	"|_ ",
	" _|",
}

var six = bigdigit{
	" _ ",
	"|_ ",
	"|_|",
}

var seven = bigdigit{
	" _ ",
	"  |",
	"  |",
}

var eight = bigdigit{
	" _ ",
	"|_|",
	"|_|",
}

var nine = bigdigit{
	" _ ",
	"|_|",
	" _|",
}

var digitlist = []bigdigit{
	zero, one, two, three, four, five, six, seven, eight, nine,
}

func init() {

}

func equal(x, y bigdigit) bool {
	if x[0] == y[0] && x[1] == y[1] && x[2] == y[2] {
		return true
	}
	return false
}

func scan(d bigdigit) int {
	for i := 0; i < len(digitlist); i++ {
		if equal(d, digitlist[i]) {
			return i
		}
	}

	return -1
}

// user story 1 : take the account number picture from fax and convert it to account number
func OCR(d []bigdigit) string {
	result := ""
	for _, c := range d {
		if v := scan(c); v >= 0 {
			result += fmt.Sprintf("%d", v)
		} else {
			result += "?"
		}
	}

	return result
}

func stringToDigit(s []string) []bigdigit {
	var digits []bigdigit

	for c := 0; c < 27; c += 3 {
		var d bigdigit
		d[0], d[1], d[2] = s[0][c:c+3], s[1][c:c+3], s[2][c:c+3]
		digits = append(digits, d)
	}

	return digits
}

func isValid(s string) bool {
	// calculate checksum
	d9 := int(s[0]) - '0'
	d8 := int(s[1]) - '0'
	d7 := int(s[2]) - '0'
	d6 := int(s[3]) - '0'
	d5 := int(s[4]) - '0'
	d4 := int(s[5]) - '0'
	d3 := int(s[6]) - '0'
	d2 := int(s[7]) - '0'
	d1 := int(s[8]) - '0'
	chksum := (d1 + 2*d2 + 3*d3 + 4*d4 + 5*d5 + 6*d6 + 7*d7 + 8*d8 + 9*d9) % 11
	return chksum == 0
}

/* TESTING */

func TestStory1() {
	log.Printf("TEST USER STORY 1 STARTED ...")
	acct0 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"| || || || || || || || || |",
		"|_||_||_||_||_||_||_||_||_|",
	}

	acct1 := []string{
		"                           ",
		"  |  |  |  |  |  |  |  |  |",
		"  |  |  |  |  |  |  |  |  |",
	}

	acct2 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		" _| _| _| _| _| _| _| _| _|",
		"|_ |_ |_ |_ |_ |_ |_ |_ |_ ",
	}

	acct3 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		" _| _| _| _| _| _| _| _| _|",
		" _| _| _| _| _| _| _| _| _|",
	}

	acct4 := []string{
		"                           ",
		"|_||_||_||_||_||_||_||_||_|",
		"  |  |  |  |  |  |  |  |  |",
	}

	acct5 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_ |_ |_ |_ |_ |_ |_ |_ |_ ",
		" _| _| _| _| _| _| _| _| _|",
	}

	acct6 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_ |_ |_ |_ |_ |_ |_ |_ |_ ",
		"|_||_||_||_||_||_||_||_||_|",
	}

	acct7 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"  |  |  |  |  |  |  |  |  |",
		"  |  |  |  |  |  |  |  |  |",
	}

	acct8 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_||_||_||_||_||_||_||_||_|",
		"|_||_||_||_||_||_||_||_||_|",
	}

	acct9 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_||_||_||_||_||_||_||_||_|",
		" _| _| _| _| _| _| _| _| _|",
	}

	acct10 := []string{
		"    _  _     _  _  _  _  _ ",
		"  | _| _||_||_ |_   ||_||_|",
		"  ||_  _|  | _||_|  ||_| _|",
	}

	assert("000000000" == OCR(stringToDigit(acct0)), "test 0")
	assert("111111111" == OCR(stringToDigit(acct1)), "test 1")
	assert("222222222" == OCR(stringToDigit(acct2)), "test 2")
	assert("333333333" == OCR(stringToDigit(acct3)), "test 3")
	assert("444444444" == OCR(stringToDigit(acct4)), "test 4")
	assert("555555555" == OCR(stringToDigit(acct5)), "test 5")
	assert("666666666" == OCR(stringToDigit(acct6)), "test 6")
	assert("777777777" == OCR(stringToDigit(acct7)), "test 7")
	assert("888888888" == OCR(stringToDigit(acct8)), "test 8")
	assert("999999999" == OCR(stringToDigit(acct9)), "test 9")
	assert("123456789" == OCR(stringToDigit(acct10)), "test 10")
}

func validate(s string) string {
	if strings.Contains(s, "?") {
		return s + " ILL"
	}

	if isValid(s) {
		return s
	}

	return s + " ERR"
}

func TestStory3() {
	log.Printf("TEST USER STORY 3 STARTED ...")

	acct0 := []string{
		" _  _  _  _  _  _  _  _    ",
		"| || || || || || || ||_   |",
		"|_||_||_||_||_||_||_| _|  |",
	}

	acct1 := []string{
		"    _  _  _  _  _  _     _ ",
		"|_||_|| || ||_   |  |  | _ ",
		"  | _||_||_||_|  |  |  | _|",
	}

	acct2 := []string{
		"    _  _     _  _  _  _  _ ",
		"  | _| _||_| _ |_   ||_||_|",
		"  ||_  _|  | _||_|  ||_| _ ",
	}

	assert("000000051" == validate(OCR(stringToDigit(acct0))), "test 11")
	assert("49006771? ILL" == validate(OCR(stringToDigit(acct1))), "test 12")
	assert("1234?678? ILL" == validate(OCR(stringToDigit(acct2))), "test 13")
	assert("664371495 ERR" == validate("664371495"), "test 14")
	assert("86110??36 ILL" == validate("86110??36"), "test 15")
}

func main() {
	TestStory1()
	TestStory3()
}
<commit_msg>add user story 4<commit_after>// http://codingdojo.org/cgi-bin/wiki.pl?KataBankOCR
package main

import (
	"fmt"
	"log"
	"strings"
)

func assert(t bool, note string) {
	if !t {
		log.Fatalf("%s FAILED !", note)
	}
}

const (
	VALID = 0
	ILL   = 1
	ERR   = 2
)

type bigdigit [3]string

type element struct {
	graphic       *bigdigit // graphic representation
	possibilities []int     // possible character
}

/* GLOBAL VARIABLES */

var zero = bigdigit{
	" _ ",
	"| |",
	"|_|",
}

var one = bigdigit{
	"   ",
	"  |",
	"  |",
}

var two = bigdigit{
	" _ ",
	" _|",
	"|_ ",
}

var three = bigdigit{
	" _ ",
	" _|",
	" _|",
}

var four = bigdigit{
	"   ",
	"|_|",
	"  |",
}

var five = bigdigit{
	" _ ",
	"|_ ",
	" _|",
}

var six = bigdigit{
	" _ ",
	"|_ ",
	"|_|",
}

var seven = bigdigit{
	" _ ",
	"  |",
	"  |",
}

var eight = bigdigit{
	" _ ",
	"|_|",
	"|_|",
}

var nine = bigdigit{
	" _ ",
	"|_|",
	" _|",
}

var digitlist = []bigdigit{
	zero, one, two, three, four, five, six, seven, eight, nine,
}

var alldigit []element

func isReplaceable(a, b *bigdigit) bool {
	count := 0
	for i := 0; i < 3; i++ {
		for j := 0; j < 3; j++ {
			//			if a[i][j] == ' ' && (b[i][j] == '|' || b[i][j] == '_') {
			// if a[i][j] != b[i][j] && (a[i][j] == ' ' || b[i][j] == ' ') {
			if a[i][j] != b[i][j] {
				count++
			} else if a[i][j] != b[i][j] {
				return false
			}
		}
	}

	if count == 1 {
		return true
	}

	return false
}

func init() {
	alldigit = make([]element, 10)
	alldigit[0] = element{&zero, []int{}}
	alldigit[1] = element{&one, []int{}}
	alldigit[2] = element{&two, []int{}}
	alldigit[3] = element{&three, []int{}}
	alldigit[4] = element{&four, []int{}}
	alldigit[5] = element{&five, []int{}}
	alldigit[6] = element{&six, []int{}}
	alldigit[7] = element{&seven, []int{}}
	alldigit[8] = element{&eight, []int{}}
	alldigit[9] = element{&nine, []int{}}

	// initialize the possibilities
	for i := 0; i < 10; i++ {
		for j := 0; j < 10; j++ {
			if i == j {
				continue
			}
			if k := isReplaceable(alldigit[i].graphic, alldigit[j].graphic); k == true {
				alldigit[i].possibilities = append(alldigit[i].possibilities, j)
			}
		}
	}

	/*
		for i, v := range alldigit {
			fmt.Printf("%d %v\n", i, v.possibilities)
		}
	*/
}

func equal(x, y bigdigit) bool {
	if x[0] == y[0] && x[1] == y[1] && x[2] == y[2] {
		return true
	}
	return false
}

func scan(d bigdigit) int {
	for i := 0; i < len(digitlist); i++ {
		if equal(d, digitlist[i]) {
			return i
		}
	}

	return -1
}

// user story 1 : take the account number picture from fax and convert it to account number
func OCR(d []bigdigit) string {
	result := ""
	for _, c := range d {
		if v := scan(c); v >= 0 {
			result += fmt.Sprintf("%d", v)
		} else {
			result += "?"
		}
	}

	return result
}

func stringToDigit(s []string) []bigdigit {
	var digits []bigdigit

	for c := 0; c < 27; c += 3 {
		var d bigdigit
		d[0], d[1], d[2] = s[0][c:c+3], s[1][c:c+3], s[2][c:c+3]
		digits = append(digits, d)
	}

	return digits
}

func isValid(s string) bool {
	// calculate checksum
	d9 := int(s[0]) - '0'
	d8 := int(s[1]) - '0'
	d7 := int(s[2]) - '0'
	d6 := int(s[3]) - '0'
	d5 := int(s[4]) - '0'
	d4 := int(s[5]) - '0'
	d3 := int(s[6]) - '0'
	d2 := int(s[7]) - '0'
	d1 := int(s[8]) - '0'
	chksum := (d1 + 2*d2 + 3*d3 + 4*d4 + 5*d5 + 6*d6 + 7*d7 + 8*d8 + 9*d9) % 11
	return chksum == 0
}

/* TESTING */

func TestStory1() {
	log.Printf("TEST USER STORY 1 STARTED ...")
	acct0 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"| || || || || || || || || |",
		"|_||_||_||_||_||_||_||_||_|",
	}

	acct1 := []string{
		"                           ",
		"  |  |  |  |  |  |  |  |  |",
		"  |  |  |  |  |  |  |  |  |",
	}

	acct2 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		" _| _| _| _| _| _| _| _| _|",
		"|_ |_ |_ |_ |_ |_ |_ |_ |_ ",
	}

	acct3 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		" _| _| _| _| _| _| _| _| _|",
		" _| _| _| _| _| _| _| _| _|",
	}

	acct4 := []string{
		"                           ",
		"|_||_||_||_||_||_||_||_||_|",
		"  |  |  |  |  |  |  |  |  |",
	}

	acct5 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_ |_ |_ |_ |_ |_ |_ |_ |_ ",
		" _| _| _| _| _| _| _| _| _|",
	}

	acct6 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_ |_ |_ |_ |_ |_ |_ |_ |_ ",
		"|_||_||_||_||_||_||_||_||_|",
	}

	acct7 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"  |  |  |  |  |  |  |  |  |",
		"  |  |  |  |  |  |  |  |  |",
	}

	acct8 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_||_||_||_||_||_||_||_||_|",
		"|_||_||_||_||_||_||_||_||_|",
	}

	acct9 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_||_||_||_||_||_||_||_||_|",
		" _| _| _| _| _| _| _| _| _|",
	}

	acct10 := []string{
		"    _  _     _  _  _  _  _ ",
		"  | _| _||_||_ |_   ||_||_|",
		"  ||_  _|  | _||_|  ||_| _|",
	}

	assert("000000000" == OCR(stringToDigit(acct0)), "test 0")
	assert("111111111" == OCR(stringToDigit(acct1)), "test 1")
	assert("222222222" == OCR(stringToDigit(acct2)), "test 2")
	assert("333333333" == OCR(stringToDigit(acct3)), "test 3")
	assert("444444444" == OCR(stringToDigit(acct4)), "test 4")
	assert("555555555" == OCR(stringToDigit(acct5)), "test 5")
	assert("666666666" == OCR(stringToDigit(acct6)), "test 6")
	assert("777777777" == OCR(stringToDigit(acct7)), "test 7")
	assert("888888888" == OCR(stringToDigit(acct8)), "test 8")
	assert("999999999" == OCR(stringToDigit(acct9)), "test 9")
	assert("123456789" == OCR(stringToDigit(acct10)), "test 10")
}

func validate(s string) int {
	if strings.Contains(s, "?") {
		return ILL
	}

	if isValid(s) {
		return VALID
	}

	return ERR
}

func validatestr(s string) string {
	if v := validate(s); v == ILL {
		return s + " ILL"
	} else if v == ERR {
		return s + " ERR"
	}
	return s
}

func TestStory3() {
	log.Printf("TEST USER STORY 3 STARTED ...")

	acct0 := []string{
		" _  _  _  _  _  _  _  _    ",
		"| || || || || || || ||_   |",
		"|_||_||_||_||_||_||_| _|  |",
	}

	acct1 := []string{
		"    _  _  _  _  _  _     _ ",
		"|_||_|| || ||_   |  |  | _ ",
		"  | _||_||_||_|  |  |  | _|",
	}

	acct2 := []string{
		"    _  _     _  _  _  _  _ ",
		"  | _| _||_| _ |_   ||_||_|",
		"  ||_  _|  | _||_|  ||_| _ ",
	}

	assert("000000051" == validatestr(OCR(stringToDigit(acct0))), "test 11")
	assert("49006771? ILL" == validatestr(OCR(stringToDigit(acct1))), "test 12")
	assert("1234?678? ILL" == validatestr(OCR(stringToDigit(acct2))), "test 13")
	assert("664371495 ERR" == validatestr("664371495"), "test 14")
	assert("86110??36 ILL" == validatestr("86110??36"), "test 15")
}

func fixAccount(s []string) string {
	acctNo := OCR(stringToDigit(s))
	result := []string{acctNo}

	if status := validate(acctNo); status == ERR {
		for i := 0; i < len(acctNo); i++ {
			keep := int(acctNo[i]) - '0'
			safe := acctNo
			for j := 0; j < len(alldigit[keep].possibilities); j++ {
				t := []byte(acctNo)
				t[i] = byte(alldigit[keep].possibilities[j]) + 48
				acctNo = string(t)

				if validate(acctNo) == VALID {
					result = append(result, acctNo)
				}
			}
			acctNo = safe
		}
	} else if status == ILL {
		//fmt.Println(acctNo)
		v := strings.Index(acctNo, "?")
		// invariant: when s == ILL, v >= 0

		d := stringToDigit(s)[v]
		//fmt.Println(d)
		for i := 0; i < 10; i++ {
			if isReplaceable(&d, alldigit[i].graphic) {
				t := []byte(acctNo)
				//	fmt.Println(t)
				t[v] = byte(i) + 48
				acctNo = string(t)
				//	fmt.Println(acctNo)

				if validate(acctNo) == VALID {
					result = append(result, acctNo)
					break
				}
			}
		}

	} else {
		return acctNo
	}

	return strings.Join(result[1:], ",")
}

func TestStory4() {
	log.Printf("TEST USER STORY 4 STARTED ...")

	tc0 := []string{
		"                           ",
		"  |  |  |  |  |  |  |  |  |",
		"  |  |  |  |  |  |  |  |  |",
	}
	tc1 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"  |  |  |  |  |  |  |  |  |",
		"  |  |  |  |  |  |  |  |  |",
	}
	tc2 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		" _|| || || || || || || || |",
		"|_ |_||_||_||_||_||_||_||_|",
	}
	tc3 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		" _| _| _| _| _| _| _| _| _|",
		" _| _| _| _| _| _| _| _| _|",
	}
	tc4 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_||_||_||_||_||_||_||_||_|",
		"|_||_||_||_||_||_||_||_||_|",
	}
	tc5 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_ |_ |_ |_ |_ |_ |_ |_ |_ ",
		" _| _| _| _| _| _| _| _| _|",
	}
	//=> 555555555 AMB ['555655555', '559555555']
	tc6 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_ |_ |_ |_ |_ |_ |_ |_ |_ ",
		"|_||_||_||_||_||_||_||_||_|",
	}
	//=> 666666666 AMB ['666566666', '686666666']
	tc7 := []string{
		" _  _  _  _  _  _  _  _  _ ",
		"|_||_||_||_||_||_||_||_||_|",
		" _| _| _| _| _| _| _| _| _|",
	}
	//=> 999999999 AMB ['899999999', '993999999', '999959999']
	tc8 := []string{
		"    _  _  _  _  _  _     _ ",
		"|_||_|| || ||_   |  |  ||_ ",
		"  | _||_||_||_|  |  |  | _|",
	}
	//=> 490067715 AMB ['490067115', '490067719', '490867715']

	tc9 := []string{
		"    _  _     _  _  _  _  _ ",
		" _| _| _||_||_ |_   ||_||_|",
		"  ||_  _|  | _||_|  ||_| _|",
	}
	//=> 123456789

	tc10 := []string{
		" _     _  _  _  _  _  _    ",
		"| || || || || || || ||_   |",
		"|_||_||_||_||_||_||_| _|  |",
	}
	//=> 000000051
	tc11 := []string{
		"    _  _  _  _  _  _     _ ",
		"|_||_|| ||_||_   |  |  | _ ",
		"  | _||_||_||_|  |  |  | _|",
	}
	//=> 490867715
	//fmt.Println(fixAccount(tc0))
	assert("711111111" == fixAccount(tc0), "test 16")
	assert("777777177" == fixAccount(tc1), "test 17")
	assert("200800000" == fixAccount(tc2), "test 18")
	assert("333393333" == fixAccount(tc3), "test 19")
	assert("888886888,888888988,888888880" == fixAccount(tc4), "test 20")
	assert("559555555,555655555" == fixAccount(tc5), "test 21")
	assert("686666666,666566666" == fixAccount(tc6), "test 22")
	assert("899999999,993999999,999959999" == fixAccount(tc7), "test 23")
	assert("490867715,490067115,490067719" == fixAccount(tc8), "test 24")
	assert("123456789" == fixAccount(tc9), "test 25")
	assert("000000051" == fixAccount(tc10), "test 26")
	assert("490867715" == fixAccount(tc11), "test 27")

}

func main() {
	TestStory1()
	TestStory3()
	TestStory4()
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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 priorities

import (
	"context"
	"fmt"
	"math"
	"sync/atomic"

	v1 "k8s.io/api/core/v1"
	"k8s.io/client-go/util/workqueue"
	"k8s.io/kubernetes/pkg/scheduler/algorithm/predicates"
	framework "k8s.io/kubernetes/pkg/scheduler/framework/v1alpha1"
	schedulerlisters "k8s.io/kubernetes/pkg/scheduler/listers"
	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"
	schedutil "k8s.io/kubernetes/pkg/scheduler/util"

	"k8s.io/klog"
)

type topologyPair struct {
	key   string
	value string
}

type podTopologySpreadMap struct {
	// nodeNameSet is a string set holding all node names which have all constraints[*].topologyKey present.
	nodeNameSet map[string]struct{}
	// topologyPairToPodCounts is keyed with topologyPair, and valued with the number of matching pods.
	topologyPairToPodCounts map[topologyPair]*int64
}

func newTopologySpreadConstraintsMap() *podTopologySpreadMap {
	return &podTopologySpreadMap{
		nodeNameSet:             make(map[string]struct{}),
		topologyPairToPodCounts: make(map[topologyPair]*int64),
	}
}

// buildPodTopologySpreadMap prepares necessary data (podTopologySpreadMap) for incoming pod on the filteredNodes.
// Later Priority function will use 'podTopologySpreadMap' to perform the Scoring calculations.
func buildPodTopologySpreadMap(pod *v1.Pod, filteredNodes []*v1.Node, allNodes []*schedulernodeinfo.NodeInfo) *podTopologySpreadMap {
	// return if incoming pod doesn't have soft topology spread constraints.
	constraints := getSoftTopologySpreadConstraints(pod)
	if len(constraints) == 0 || len(filteredNodes) == 0 || len(allNodes) == 0 {
		return nil
	}

	// initialize podTopologySpreadMap which will be used in Score plugin.
	m := newTopologySpreadConstraintsMap()
	m.initialize(pod, filteredNodes)

	errCh := schedutil.NewErrorChannel()
	ctx, cancel := context.WithCancel(context.Background())
	processAllNode := func(i int) {
		nodeInfo := allNodes[i]
		node := nodeInfo.Node()
		if node == nil {
			return
		}
		// (1) `node` should satisfy incoming pod's NodeSelector/NodeAffinity
		// (2) All topologyKeys need to be present in `node`
		if !predicates.PodMatchesNodeSelectorAndAffinityTerms(pod, node) ||
			!predicates.NodeLabelsMatchSpreadConstraints(node.Labels, constraints) {
			return
		}

		for _, constraint := range constraints {
			pair := topologyPair{key: constraint.TopologyKey, value: node.Labels[constraint.TopologyKey]}
			// If current topology pair is not associated with any candidate node,
			// continue to avoid unnecessary calculation.
			if m.topologyPairToPodCounts[pair] == nil {
				continue
			}

			// <matchSum> indicates how many pods (on current node) match the <constraint>.
			matchSum := int64(0)
			for _, existingPod := range nodeInfo.Pods() {
				match, err := predicates.PodMatchesSpreadConstraint(existingPod.Labels, constraint)
				if err != nil {
					errCh.SendErrorWithCancel(err, cancel)
					return
				}
				if match {
					matchSum++
				}
			}
			atomic.AddInt64(m.topologyPairToPodCounts[pair], matchSum)
		}
	}
	workqueue.ParallelizeUntil(ctx, 16, len(allNodes), processAllNode)
	if err := errCh.ReceiveError(); err != nil {
		klog.Error(err)
		return nil
	}

	return m
}

// initialize iterates "filteredNodes" to filter out the nodes which don't have required topologyKey(s),
// and initialize two maps:
// 1) m.topologyPairToPodCounts: keyed with both eligible topology pair and node names.
// 2) m.nodeNameSet: keyed with node name, and valued with a *int64 pointer for eligible node only.
func (m *podTopologySpreadMap) initialize(pod *v1.Pod, filteredNodes []*v1.Node) {
	constraints := getSoftTopologySpreadConstraints(pod)
	for _, node := range filteredNodes {
		if !predicates.NodeLabelsMatchSpreadConstraints(node.Labels, constraints) {
			continue
		}
		for _, constraint := range constraints {
			pair := topologyPair{key: constraint.TopologyKey, value: node.Labels[constraint.TopologyKey]}
			if m.topologyPairToPodCounts[pair] == nil {
				m.topologyPairToPodCounts[pair] = new(int64)
			}
		}
		m.nodeNameSet[node.Name] = struct{}{}
		// For those nodes which don't have all required topologyKeys present, it's intentional to leave
		// their entries absent in nodeNameSet, so that we're able to score them to 0 afterwards.
	}
}

// CalculateEvenPodsSpreadPriorityMap calculate the number of matching pods on the passed-in "node",
// and return the number as Score.
func CalculateEvenPodsSpreadPriorityMap(pod *v1.Pod, meta interface{}, nodeInfo *schedulernodeinfo.NodeInfo) (framework.NodeScore, error) {
	node := nodeInfo.Node()
	if node == nil {
		return framework.NodeScore{}, fmt.Errorf("node not found")
	}

	var m *podTopologySpreadMap
	if priorityMeta, ok := meta.(*priorityMetadata); ok {
		m = priorityMeta.podTopologySpreadMap
	}
	if m == nil {
		return framework.NodeScore{}, nil
	}

	// no need to continue if the node is not qualified.
	if _, ok := m.nodeNameSet[node.Name]; !ok {
		return framework.NodeScore{Name: node.Name, Score: 0}, nil
	}

	constraints := getSoftTopologySpreadConstraints(pod)
	// For each present <pair>, current node gets a credit of <matchSum>.
	// And we sum up <matchSum> and return it as this node's score.
	var score int64
	for _, constraint := range constraints {
		if tpVal, ok := node.Labels[constraint.TopologyKey]; ok {
			pair := topologyPair{key: constraint.TopologyKey, value: tpVal}
			matchSum := *m.topologyPairToPodCounts[pair]
			score += matchSum
		}
	}
	return framework.NodeScore{Name: node.Name, Score: score}, nil
}

// CalculateEvenPodsSpreadPriorityReduce normalizes the score for each filteredNode,
// The basic rule is: the bigger the score(matching number of pods) is, the smaller the
// final normalized score will be.
func CalculateEvenPodsSpreadPriorityReduce(pod *v1.Pod, meta interface{}, sharedLister schedulerlisters.SharedLister,
	result framework.NodeScoreList) error {
	var m *podTopologySpreadMap
	if priorityMeta, ok := meta.(*priorityMetadata); ok {
		m = priorityMeta.podTopologySpreadMap
	}
	if m == nil {
		return nil
	}

	// Calculate the summed <total> score and <minScore>.
	var minScore int64 = math.MaxInt64
	var total int64
	for _, score := range result {
		// it's mandatory to check if <score.Name> is present in m.nodeNameSet
		if _, ok := m.nodeNameSet[score.Name]; !ok {
			continue
		}
		total += score.Score
		if score.Score < minScore {
			minScore = score.Score
		}
	}

	maxMinDiff := total - minScore
	for i := range result {
		nodeInfo, err := sharedLister.NodeInfos().Get(result[i].Name)
		if err != nil {
			return err
		}
		node := nodeInfo.Node()
		// Debugging purpose: print the score for each node.
		// Score must be a pointer here, otherwise it's always 0.
		if klog.V(10) {
			defer func(score *int64, nodeName string) {
				klog.Infof("%v -> %v: PodTopologySpread NormalizeScore, Score: (%d)", pod.Name, nodeName, *score)
			}(&result[i].Score, node.Name)
		}

		if maxMinDiff == 0 {
			result[i].Score = framework.MaxNodeScore
			continue
		}

		if _, ok := m.nodeNameSet[node.Name]; !ok {
			result[i].Score = 0
			continue
		}

		flippedScore := total - result[i].Score
		fScore := float64(framework.MaxNodeScore) * (float64(flippedScore) / float64(maxMinDiff))
		result[i].Score = int64(fScore)
	}
	return nil
}

// TODO(Huang-Wei): combine this with getHardTopologySpreadConstraints() in predicates package
func getSoftTopologySpreadConstraints(pod *v1.Pod) (constraints []v1.TopologySpreadConstraint) {
	if pod != nil {
		for _, constraint := range pod.Spec.TopologySpreadConstraints {
			if constraint.WhenUnsatisfiable == v1.ScheduleAnyway {
				constraints = append(constraints, constraint)
			}
		}
	}
	return
}
<commit_msg>fix: return NodeScore in even pods spread priority<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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 priorities

import (
	"context"
	"fmt"
	"math"
	"sync/atomic"

	v1 "k8s.io/api/core/v1"
	"k8s.io/client-go/util/workqueue"
	"k8s.io/kubernetes/pkg/scheduler/algorithm/predicates"
	framework "k8s.io/kubernetes/pkg/scheduler/framework/v1alpha1"
	schedulerlisters "k8s.io/kubernetes/pkg/scheduler/listers"
	schedulernodeinfo "k8s.io/kubernetes/pkg/scheduler/nodeinfo"
	schedutil "k8s.io/kubernetes/pkg/scheduler/util"

	"k8s.io/klog"
)

type topologyPair struct {
	key   string
	value string
}

type podTopologySpreadMap struct {
	// nodeNameSet is a string set holding all node names which have all constraints[*].topologyKey present.
	nodeNameSet map[string]struct{}
	// topologyPairToPodCounts is keyed with topologyPair, and valued with the number of matching pods.
	topologyPairToPodCounts map[topologyPair]*int64
}

func newTopologySpreadConstraintsMap() *podTopologySpreadMap {
	return &podTopologySpreadMap{
		nodeNameSet:             make(map[string]struct{}),
		topologyPairToPodCounts: make(map[topologyPair]*int64),
	}
}

// buildPodTopologySpreadMap prepares necessary data (podTopologySpreadMap) for incoming pod on the filteredNodes.
// Later Priority function will use 'podTopologySpreadMap' to perform the Scoring calculations.
func buildPodTopologySpreadMap(pod *v1.Pod, filteredNodes []*v1.Node, allNodes []*schedulernodeinfo.NodeInfo) *podTopologySpreadMap {
	// return if incoming pod doesn't have soft topology spread constraints.
	constraints := getSoftTopologySpreadConstraints(pod)
	if len(constraints) == 0 || len(filteredNodes) == 0 || len(allNodes) == 0 {
		return nil
	}

	// initialize podTopologySpreadMap which will be used in Score plugin.
	m := newTopologySpreadConstraintsMap()
	m.initialize(pod, filteredNodes)

	errCh := schedutil.NewErrorChannel()
	ctx, cancel := context.WithCancel(context.Background())
	processAllNode := func(i int) {
		nodeInfo := allNodes[i]
		node := nodeInfo.Node()
		if node == nil {
			return
		}
		// (1) `node` should satisfy incoming pod's NodeSelector/NodeAffinity
		// (2) All topologyKeys need to be present in `node`
		if !predicates.PodMatchesNodeSelectorAndAffinityTerms(pod, node) ||
			!predicates.NodeLabelsMatchSpreadConstraints(node.Labels, constraints) {
			return
		}

		for _, constraint := range constraints {
			pair := topologyPair{key: constraint.TopologyKey, value: node.Labels[constraint.TopologyKey]}
			// If current topology pair is not associated with any candidate node,
			// continue to avoid unnecessary calculation.
			if m.topologyPairToPodCounts[pair] == nil {
				continue
			}

			// <matchSum> indicates how many pods (on current node) match the <constraint>.
			matchSum := int64(0)
			for _, existingPod := range nodeInfo.Pods() {
				match, err := predicates.PodMatchesSpreadConstraint(existingPod.Labels, constraint)
				if err != nil {
					errCh.SendErrorWithCancel(err, cancel)
					return
				}
				if match {
					matchSum++
				}
			}
			atomic.AddInt64(m.topologyPairToPodCounts[pair], matchSum)
		}
	}
	workqueue.ParallelizeUntil(ctx, 16, len(allNodes), processAllNode)
	if err := errCh.ReceiveError(); err != nil {
		klog.Error(err)
		return nil
	}

	return m
}

// initialize iterates "filteredNodes" to filter out the nodes which don't have required topologyKey(s),
// and initialize two maps:
// 1) m.topologyPairToPodCounts: keyed with both eligible topology pair and node names.
// 2) m.nodeNameSet: keyed with node name, and valued with a *int64 pointer for eligible node only.
func (m *podTopologySpreadMap) initialize(pod *v1.Pod, filteredNodes []*v1.Node) {
	constraints := getSoftTopologySpreadConstraints(pod)
	for _, node := range filteredNodes {
		if !predicates.NodeLabelsMatchSpreadConstraints(node.Labels, constraints) {
			continue
		}
		for _, constraint := range constraints {
			pair := topologyPair{key: constraint.TopologyKey, value: node.Labels[constraint.TopologyKey]}
			if m.topologyPairToPodCounts[pair] == nil {
				m.topologyPairToPodCounts[pair] = new(int64)
			}
		}
		m.nodeNameSet[node.Name] = struct{}{}
		// For those nodes which don't have all required topologyKeys present, it's intentional to leave
		// their entries absent in nodeNameSet, so that we're able to score them to 0 afterwards.
	}
}

// CalculateEvenPodsSpreadPriorityMap calculate the number of matching pods on the passed-in "node",
// and return the number as Score.
func CalculateEvenPodsSpreadPriorityMap(pod *v1.Pod, meta interface{}, nodeInfo *schedulernodeinfo.NodeInfo) (framework.NodeScore, error) {
	node := nodeInfo.Node()
	if node == nil {
		return framework.NodeScore{}, fmt.Errorf("node not found")
	}

	var m *podTopologySpreadMap
	if priorityMeta, ok := meta.(*priorityMetadata); ok {
		m = priorityMeta.podTopologySpreadMap
	}
	if m == nil {
		return framework.NodeScore{Name: node.Name, Score: 0}, nil
	}

	// no need to continue if the node is not qualified.
	if _, ok := m.nodeNameSet[node.Name]; !ok {
		return framework.NodeScore{Name: node.Name, Score: 0}, nil
	}

	constraints := getSoftTopologySpreadConstraints(pod)
	// For each present <pair>, current node gets a credit of <matchSum>.
	// And we sum up <matchSum> and return it as this node's score.
	var score int64
	for _, constraint := range constraints {
		if tpVal, ok := node.Labels[constraint.TopologyKey]; ok {
			pair := topologyPair{key: constraint.TopologyKey, value: tpVal}
			matchSum := *m.topologyPairToPodCounts[pair]
			score += matchSum
		}
	}
	return framework.NodeScore{Name: node.Name, Score: score}, nil
}

// CalculateEvenPodsSpreadPriorityReduce normalizes the score for each filteredNode,
// The basic rule is: the bigger the score(matching number of pods) is, the smaller the
// final normalized score will be.
func CalculateEvenPodsSpreadPriorityReduce(pod *v1.Pod, meta interface{}, sharedLister schedulerlisters.SharedLister,
	result framework.NodeScoreList) error {
	var m *podTopologySpreadMap
	if priorityMeta, ok := meta.(*priorityMetadata); ok {
		m = priorityMeta.podTopologySpreadMap
	}
	if m == nil {
		return nil
	}

	// Calculate the summed <total> score and <minScore>.
	var minScore int64 = math.MaxInt64
	var total int64
	for _, score := range result {
		// it's mandatory to check if <score.Name> is present in m.nodeNameSet
		if _, ok := m.nodeNameSet[score.Name]; !ok {
			continue
		}
		total += score.Score
		if score.Score < minScore {
			minScore = score.Score
		}
	}

	maxMinDiff := total - minScore
	for i := range result {
		nodeInfo, err := sharedLister.NodeInfos().Get(result[i].Name)
		if err != nil {
			return err
		}
		node := nodeInfo.Node()
		// Debugging purpose: print the score for each node.
		// Score must be a pointer here, otherwise it's always 0.
		if klog.V(10) {
			defer func(score *int64, nodeName string) {
				klog.Infof("%v -> %v: PodTopologySpread NormalizeScore, Score: (%d)", pod.Name, nodeName, *score)
			}(&result[i].Score, node.Name)
		}

		if maxMinDiff == 0 {
			result[i].Score = framework.MaxNodeScore
			continue
		}

		if _, ok := m.nodeNameSet[node.Name]; !ok {
			result[i].Score = 0
			continue
		}

		flippedScore := total - result[i].Score
		fScore := float64(framework.MaxNodeScore) * (float64(flippedScore) / float64(maxMinDiff))
		result[i].Score = int64(fScore)
	}
	return nil
}

// TODO(Huang-Wei): combine this with getHardTopologySpreadConstraints() in predicates package
func getSoftTopologySpreadConstraints(pod *v1.Pod) (constraints []v1.TopologySpreadConstraint) {
	if pod != nil {
		for _, constraint := range pod.Spec.TopologySpreadConstraints {
			if constraint.WhenUnsatisfiable == v1.ScheduleAnyway {
				constraints = append(constraints, constraint)
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup

import (
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/aws/sdb"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/crypto"
	"math/rand"
	"regexp"
	"strconv"
	"time"
	"unicode/utf8"
)

const (
	// The item name we use for "domain is already in use" markers.
	markerItemName = "comeback_marker"

	// The attribute name we use for storing crypto-compatibility data.
	cryptoMarkerAttributeName = "encrypted_data"

	// A time format that works properly with range queries.
	iso8601TimeFormat = "2006-01-02T15:04:05Z"
)

// A regexp for selected item names.
var itemNameRegexp = regexp.MustCompile(`^backup_([0-9a-f]{16})$`)

// A record in the backup registry describing a successful backup job.
type CompletedJob struct {
	// A unique ID for this completed run.
	Id uint64

	// The name of the backup job.
	Name string

	// The time at which the backup was started.
	StartTime time.Time

	// The score representing the contents of the backup.
	Score blob.Score
}

type Registry interface {
	// Record that the named backup job has completed.
	RecordBackup(j CompletedJob) (err error)

	// Return a list of the most recent completed backups.
	ListRecentBackups() (jobs []CompletedJob, err error)

	// Find a particular completed job by ID.
	FindBackup(jobId uint64) (job CompletedJob, err error)
}

func verifyCompatible(
	markerAttr sdb.Attribute,
	crypter crypto.Crypter) (err error) {
	// The ciphertext should be base64-encoded.
	encoded := markerAttr.Value
	ciphertext, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		err = fmt.Errorf("base64.DecodeString(%s): %v", encoded, err)
		return
	}

	// Attempt to decrypt the ciphertext.
	if _, err = crypter.Decrypt(ciphertext); err != nil {
		// Special case: Did the crypter signal that the key was wrong?
		if _, ok := err.(*crypto.NotAuthenticError); ok {
			err = &IncompatibleCrypterError{
				"The supplied crypter is not compatible with the data in the domain.",
			}
			return
		}

		// Generic error.
		err = fmt.Errorf("Decrypt: %v", err)
		return
	}

	return
}

// Create a registry that stores data in the supplied SimpleDB domain.
//
// Before doing so, check to see whether this domain has been used as a
// registry before. If not, write an encrypted marker with the supplied
// crypter. If it has been used before, make sure that it was used with a
// crypter compatible with the supplied one. This prevents accidentally writing
// data with the wrong key, as if the user entered the wrong password.
//
// The crypter must be set up such that it is guaranteed to return an error if
// it is used to decrypt ciphertext encrypted with a different key. In that
// case, this function will return an *IncompatibleCrypterError.
func NewRegistry(
	domain sdb.Domain,
	crypter crypto.Crypter,
	randSrc *rand.Rand,
) (r Registry, err error) {
	// Set up a tentative result.
	r = &registry{crypter, domain.Db(), domain}

	// Ask for the data that will tell us whether the crypter is compatible with
	// the previous one used in this domain, if any.
	attrs, err := domain.GetAttributes(
		markerItemName,
		false, // No need to ask for a consistent read
		[]string{cryptoMarkerAttributeName},
	)

	if err != nil {
		err = fmt.Errorf("GetAttributes: %v", err)
		return
	}

	// If we got back an attribute, we must verify that it is compatible.
	if len(attrs) > 0 {
		err = verifyCompatible(attrs[0], crypter)
		return
	}

	// Otherwise, we want to claim this domain. Encrypt some random data, base64
	// encode it, then write it out. Make sure to use a precondition to defeat
	// the race condition where another machine is doing the same simultaneously.
	plaintext := get8RandBytes(randSrc)
	ciphertext, err := crypter.Encrypt(plaintext)
	if err != nil {
		err = fmt.Errorf("Encrypt: %v", err)
		return
	}

	encoded := base64.StdEncoding.EncodeToString(ciphertext)

	err = domain.PutAttributes(
		markerItemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: cryptoMarkerAttributeName, Value: encoded},
		},
		&sdb.Precondition{Name: cryptoMarkerAttributeName, Value: nil},
	)

	if err != nil {
		err = fmt.Errorf("PutAttributes: %v", err)
		return
	}

	// All is good.
	return
}

func get8RandBytes(src *rand.Rand) []byte {
	a := src.Uint32()
	b := src.Uint32()

	return []byte{
		byte(a),
		byte(a >> 8),
		byte(a >> 16),
		byte(a >> 24),
		byte(b),
		byte(b >> 8),
		byte(b >> 16),
		byte(b >> 24),
	}
}

type registry struct {
	crypter crypto.Crypter
	db      sdb.SimpleDB
	domain  sdb.Domain
}

func (r *registry) RecordBackup(job CompletedJob) (err error) {
	// Job names must be valid SimpleDB attribute values.
	if len(job.Name) == 0 || len(job.Name) > 1024 || !utf8.ValidString(job.Name) {
		err = fmt.Errorf(
			"Job names must be non-empty UTF-8 no more than 1024 bytes long.")
		return
	}

	// Format the item name.
	itemName := sdb.ItemName(fmt.Sprintf("backup_%016x", job.Id))

	// Use a time format that works correctly with range queries. Make sure to
	// standardize on UTC.
	formattedTime := job.StartTime.UTC().Format(iso8601TimeFormat)

	// Call the domain.
	updates := []sdb.PutUpdate{
		sdb.PutUpdate{Name: "job_name", Value: job.Name},
		sdb.PutUpdate{Name: "start_time", Value: formattedTime},
		sdb.PutUpdate{Name: "score", Value: job.Score.Hex()},
	}

	precond := sdb.Precondition{Name: "score", Value: nil}

	if err = r.domain.PutAttributes(itemName, updates, &precond); err != nil {
		err = fmt.Errorf("PutAttributes: %v", err)
		return
	}

	return
}

func convertSelectedItem(item sdb.SelectedItem) (j CompletedJob, err error) {
	// Convert the item name.
	subMatches := itemNameRegexp.FindStringSubmatch(string(item.Name))
	if subMatches == nil {
		err = fmt.Errorf("Invalid item name: %s", item.Name)
		return
	}

	if j.Id, err = strconv.ParseUint(subMatches[1], 16, 64); err != nil {
		panic(fmt.Sprintf("Unexpected result for name: %s", item.Name))
	}

	// Convert each attribute.
	for _, attr := range item.Attributes {
		switch attr.Name {
		case "job_name":
			j.Name = attr.Value

		case "start_time":
			if j.StartTime, err = time.Parse(iso8601TimeFormat, attr.Value); err != nil {
				err = fmt.Errorf("Invalid start_time value: %v", err)
				return
			}

		case "score":
			var decoded []byte
			decoded, err = hex.DecodeString(attr.Value)
			if err != nil || len(decoded) != 20 {
				err = fmt.Errorf("Invalid score: %s", attr.Value)
				return
			}

			j.Score = blob.Score(decoded)
		}
	}

	// Everything must have been returned.
	if j.Name == "" {
		err = fmt.Errorf("Missing name attribute.")
		return
	}

	if j.StartTime.IsZero() {
		err = fmt.Errorf("Missing start_time attribute.")
		return
	}

	if j.Score == nil {
		err = fmt.Errorf("Missing score attribute.")
		return
	}

	return
}

func (r *registry) ListRecentBackups() (jobs []CompletedJob, err error) {
	// Call the database.
	query := fmt.Sprintf(
		"select job_name, start_time, score from `%s` where "+
			"start_time is not null order by start_time desc",
		r.domain.Name(),
	)

	results, _, err := r.db.Select(
		query,
		false, // No need for consistent reads.
		nil,
	)

	if err != nil {
		err = fmt.Errorf("Select: %v", err)
		return
	}

	// Convert each result.
	for _, item := range results {
		var job CompletedJob
		job, err = convertSelectedItem(item)
		if err != nil {
			err = fmt.Errorf("Item %s is invalid: %v", item.Name, err)
			return
		}

		jobs = append(jobs, job)
	}

	return
}

func (r *registry) FindBackup(jobId uint64) (job CompletedJob, err error) {
	err = fmt.Errorf("TODO")
	return
}

type IncompatibleCrypterError struct {
	s string
}

func (e *IncompatibleCrypterError) Error() string {
	return e.s
}
<commit_msg>Refactored.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup

import (
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"github.com/jacobsa/aws/sdb"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/crypto"
	"math/rand"
	"regexp"
	"strconv"
	"time"
	"unicode/utf8"
)

const (
	// The item name we use for "domain is already in use" markers.
	markerItemName = "comeback_marker"

	// The attribute name we use for storing crypto-compatibility data.
	cryptoMarkerAttributeName = "encrypted_data"

	// A time format that works properly with range queries.
	iso8601TimeFormat = "2006-01-02T15:04:05Z"
)

// A regexp for selected item names.
var itemNameRegexp = regexp.MustCompile(`^backup_([0-9a-f]{16})$`)

// A record in the backup registry describing a successful backup job.
type CompletedJob struct {
	// A unique ID for this completed run.
	Id uint64

	// The name of the backup job.
	Name string

	// The time at which the backup was started.
	StartTime time.Time

	// The score representing the contents of the backup.
	Score blob.Score
}

type Registry interface {
	// Record that the named backup job has completed.
	RecordBackup(j CompletedJob) (err error)

	// Return a list of the most recent completed backups.
	ListRecentBackups() (jobs []CompletedJob, err error)

	// Find a particular completed job by ID.
	FindBackup(jobId uint64) (job CompletedJob, err error)
}

func verifyCompatible(
	markerAttr sdb.Attribute,
	crypter crypto.Crypter) (err error) {
	// The ciphertext should be base64-encoded.
	encoded := markerAttr.Value
	ciphertext, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		err = fmt.Errorf("base64.DecodeString(%s): %v", encoded, err)
		return
	}

	// Attempt to decrypt the ciphertext.
	if _, err = crypter.Decrypt(ciphertext); err != nil {
		// Special case: Did the crypter signal that the key was wrong?
		if _, ok := err.(*crypto.NotAuthenticError); ok {
			err = &IncompatibleCrypterError{
				"The supplied crypter is not compatible with the data in the domain.",
			}
			return
		}

		// Generic error.
		err = fmt.Errorf("Decrypt: %v", err)
		return
	}

	return
}

// Create a registry that stores data in the supplied SimpleDB domain.
//
// Before doing so, check to see whether this domain has been used as a
// registry before. If not, write an encrypted marker with the supplied
// crypter. If it has been used before, make sure that it was used with a
// crypter compatible with the supplied one. This prevents accidentally writing
// data with the wrong key, as if the user entered the wrong password.
//
// The crypter must be set up such that it is guaranteed to return an error if
// it is used to decrypt ciphertext encrypted with a different key. In that
// case, this function will return an *IncompatibleCrypterError.
func NewRegistry(
	domain sdb.Domain,
	crypter crypto.Crypter,
	randSrc *rand.Rand,
) (r Registry, err error) {
	// Set up a tentative result.
	r = &registry{crypter, domain.Db(), domain}

	// Ask for the data that will tell us whether the crypter is compatible with
	// the previous one used in this domain, if any.
	attrs, err := domain.GetAttributes(
		markerItemName,
		false, // No need to ask for a consistent read
		[]string{cryptoMarkerAttributeName},
	)

	if err != nil {
		err = fmt.Errorf("GetAttributes: %v", err)
		return
	}

	// If we got back an attribute, we must verify that it is compatible.
	if len(attrs) > 0 {
		err = verifyCompatible(attrs[0], crypter)
		return
	}

	// Otherwise, we want to claim this domain. Encrypt some random data, base64
	// encode it, then write it out. Make sure to use a precondition to defeat
	// the race condition where another machine is doing the same simultaneously.
	plaintext := get8RandBytes(randSrc)
	ciphertext, err := crypter.Encrypt(plaintext)
	if err != nil {
		err = fmt.Errorf("Encrypt: %v", err)
		return
	}

	encoded := base64.StdEncoding.EncodeToString(ciphertext)

	err = domain.PutAttributes(
		markerItemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: cryptoMarkerAttributeName, Value: encoded},
		},
		&sdb.Precondition{Name: cryptoMarkerAttributeName, Value: nil},
	)

	if err != nil {
		err = fmt.Errorf("PutAttributes: %v", err)
		return
	}

	// All is good.
	return
}

func get8RandBytes(src *rand.Rand) []byte {
	a := src.Uint32()
	b := src.Uint32()

	return []byte{
		byte(a),
		byte(a >> 8),
		byte(a >> 16),
		byte(a >> 24),
		byte(b),
		byte(b >> 8),
		byte(b >> 16),
		byte(b >> 24),
	}
}

type registry struct {
	crypter crypto.Crypter
	db      sdb.SimpleDB
	domain  sdb.Domain
}

func (r *registry) RecordBackup(job CompletedJob) (err error) {
	// Job names must be valid SimpleDB attribute values.
	if len(job.Name) == 0 || len(job.Name) > 1024 || !utf8.ValidString(job.Name) {
		err = fmt.Errorf(
			"Job names must be non-empty UTF-8 no more than 1024 bytes long.")
		return
	}

	// Format the item name.
	itemName := sdb.ItemName(fmt.Sprintf("backup_%016x", job.Id))

	// Use a time format that works correctly with range queries. Make sure to
	// standardize on UTC.
	formattedTime := job.StartTime.UTC().Format(iso8601TimeFormat)

	// Call the domain.
	updates := []sdb.PutUpdate{
		sdb.PutUpdate{Name: "job_name", Value: job.Name},
		sdb.PutUpdate{Name: "start_time", Value: formattedTime},
		sdb.PutUpdate{Name: "score", Value: job.Score.Hex()},
	}

	precond := sdb.Precondition{Name: "score", Value: nil}

	if err = r.domain.PutAttributes(itemName, updates, &precond); err != nil {
		err = fmt.Errorf("PutAttributes: %v", err)
		return
	}

	return
}

func convertAttributes(attrs []sdb.Attribute) (j CompletedJob, err error) {
	for _, attr := range attrs {
		switch attr.Name {
		case "job_name":
			j.Name = attr.Value

		case "start_time":
			if j.StartTime, err = time.Parse(iso8601TimeFormat, attr.Value); err != nil {
				err = fmt.Errorf("Invalid start_time value: %v", err)
				return
			}

		case "score":
			var decoded []byte
			decoded, err = hex.DecodeString(attr.Value)
			if err != nil || len(decoded) != 20 {
				err = fmt.Errorf("Invalid score: %s", attr.Value)
				return
			}

			j.Score = blob.Score(decoded)
		}
	}

	// Everything must have been returned.
	if j.Name == "" {
		err = fmt.Errorf("Missing name attribute.")
		return
	}

	if j.StartTime.IsZero() {
		err = fmt.Errorf("Missing start_time attribute.")
		return
	}

	if j.Score == nil {
		err = fmt.Errorf("Missing score attribute.")
		return
	}

	return
}

func convertSelectedItem(item sdb.SelectedItem) (j CompletedJob, err error) {
	// Convert the item's attributes.
	if j, err = convertAttributes(item.Attributes); err != nil {
		return
	}

	// Convert the item name.
	subMatches := itemNameRegexp.FindStringSubmatch(string(item.Name))
	if subMatches == nil {
		err = fmt.Errorf("Invalid item name: %s", item.Name)
		return
	}

	if j.Id, err = strconv.ParseUint(subMatches[1], 16, 64); err != nil {
		panic(fmt.Sprintf("Unexpected result for name: %s", item.Name))
	}

	return
}

func (r *registry) ListRecentBackups() (jobs []CompletedJob, err error) {
	// Call the database.
	query := fmt.Sprintf(
		"select job_name, start_time, score from `%s` where "+
			"start_time is not null order by start_time desc",
		r.domain.Name(),
	)

	results, _, err := r.db.Select(
		query,
		false, // No need for consistent reads.
		nil,
	)

	if err != nil {
		err = fmt.Errorf("Select: %v", err)
		return
	}

	// Convert each result.
	for _, item := range results {
		var job CompletedJob
		job, err = convertSelectedItem(item)
		if err != nil {
			err = fmt.Errorf("Item %s is invalid: %v", item.Name, err)
			return
		}

		jobs = append(jobs, job)
	}

	return
}

func (r *registry) FindBackup(jobId uint64) (job CompletedJob, err error) {
	err = fmt.Errorf("TODO")
	return
}

type IncompatibleCrypterError struct {
	s string
}

func (e *IncompatibleCrypterError) Error() string {
	return e.s
}
<|endoftext|>


<commit_before>// Copyright 2014 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package assert

import (
	"errors"
	"testing"
)

func TestGetCallerInfo(t *testing.T) {
	str := getCallerInfo()
	// NOTE:注意这里涉及到调用函数的行号信息
	if str != "TestGetCallerInfo(assert_test.go:13)" {
		t.Errorf("getCallerInfo返回的信息不正确，其返回值为：%v", str)
	}

	// 嵌套调用，第二个参数为当前的行号
	testGetCallerInfo(t, "20")
	testGetCallerInfo(t, "21")

	// 闭合函数，line为调用所在的行号。
	f := func(line string) {
		str := getCallerInfo()
		if str != "TestGetCallerInfo(assert_test.go:"+line+")" ||
			str != "TestGetCallerInfo.fun1(assert_test.go:"+line+")" { // Go1.5之后变这样了
			t.Errorf("getCallerInfo返回的信息不正确，其返回值为：%v", str)
		}
	}
	f("31") // 参数为当前等号
	f("32")
}

// 参数line，为调用此函数所在的行号。
func testGetCallerInfo(t *testing.T, line string) {
	str := getCallerInfo()
	if str != "TestGetCallerInfo(assert_test.go:"+line+")" {
		t.Error("getCallerInfo返回的信息不正确，其返回值为：%v", str)
	}

}

func TestFormatMsg(t *testing.T) {
	msg1 := []interface{}{}
	msg2 := []interface{}{[]rune("msg:%v"), 2}
	msg3 := []interface{}{"msg:%v", 3}
	msg4 := []interface{}{123, 456}

	str := formatMessage(msg1, msg2)
	if str != "msg:2" {
		t.Errorf("formatMessage(msg1,msg2)返回信息错误:[%v]", str)
	}

	str = formatMessage(nil, msg2)
	if str != "msg:2" {
		t.Errorf("formatMessage(msg1,msg2)返回信息错误:[%v]", str)
	}

	str = formatMessage(msg2, msg3)
	if str != "msg:2" {
		t.Errorf("formatMessage(msg2,msg3)返回信息错误:[%v]", str)
	}

	str = formatMessage(nil, nil)
	if str != "<未提供任何错误信息>" {
		t.Errorf("formatMessage(nil,nil)返回信息错误:[%v]", str)
	}

	str = formatMessage(nil, msg4)
	if str != "<无法正确转换错误提示信息>" {
		t.Errorf("formatMessage(nil,nil)返回信息错误:[%v]", str)
	}
}

func TestTrue(t *testing.T) {
	True(t, true)
	True(t, 1 == 1, "True(1==1) falid")
}

func TestFalse(t *testing.T) {
	False(t, false, "False falid")
	False(t, 1 == 2, "False(1==2) falid")
}

func TestNil(t *testing.T) {
	Nil(t, nil, "Nil falid")

	var v interface{}
	Nil(t, v, "Nil(v) falid")
}

func TestNotNil(t *testing.T) {
	NotNil(t, 5, "NotNil falid")

	var v interface{} = 5
	NotNil(t, v, "NotNil falid")
}

func TestEqual(t *testing.T) {
	Equal(t, 5, 5, "Equal(5,5) falid")

	var v1, v2 interface{}
	v1 = 5
	v2 = 5

	Equal(t, 5, v1)
	Equal(t, v1, v2, "Equal(v1,v2) falid")
	Equal(t, int8(126), 126)
	Equal(t, int64(126), int8(126))
	Equal(t, uint(7), int(7))
}

func TestNotEqual(t *testing.T) {
	NotEqual(t, 5, 6, "NotEqual(5,6) falid")

	var v1, v2 interface{} = 5, 6

	NotEqual(t, 5, v2, "NotEqual(5,v2) falid")
	NotEqual(t, v1, v2, "NotEqual(v1,v2) falid")
	NotEqual(t, 128, int8(127))
}

func TestEmpty(t *testing.T) {
	Empty(t, 0, "Empty(0) falid")
	Empty(t, "", "Empty(``) falid")
	Empty(t, false, "Empty(false) falid")
	Empty(t, []string{}, "Empty(slice{}) falid")
	Empty(t, []int{}, "Empty(slice{}) falid")
}

func TestNotEmpty(t *testing.T) {
	NotEmpty(t, 1, "NotEmpty(1) falid")
	NotEmpty(t, true, "NotEmpty(true) falid")
	NotEmpty(t, []string{"ab"}, "NotEmpty(slice(abc)) falid")
}

type ErrorImpl struct {
	msg string
}

func (err *ErrorImpl) Error() string {
	return err.msg
}

func TestError(t *testing.T) {
	err1 := errors.New("test")
	Error(t, err1, "Error(err) falid")

	err2 := &ErrorImpl{msg: "msg"}
	Error(t, err2, "Error(ErrorImpl) falid")
}

func TestNotError(t *testing.T) {
	NotError(t, "123", "NotError(123) falid")

	var err1 error = nil
	NotError(t, err1, "var err1 error falid")

	err2 := &ErrorImpl{msg: "msg"}
	Error(t, err2, "Error(ErrorImpl) falid")
}

func TestFileExists(t *testing.T) {
	FileExists(t, "./assert.go", "FileExists() falid")
	FileExists(t, "./", "FileExists() falid")
}

func TestFileNotExists(t *testing.T) {
	FileNotExists(t, "c:/win", "FileNotExists() falid")
	FileNotExists(t, "./abcefg/", "FileNotExists() falid")
}

func TestPanic(t *testing.T) {
	f1 := func() {
		panic("panic")
	}

	Panic(t, f1)
}

func TestNotPanic(t *testing.T) {
	f1 := func() {
	}

	NotPanic(t, f1)
}

func TestContains(t *testing.T) {
	Contains(t, []int{1, 2, 3}, []int8{1, 2})
}

func TestNotContains(t *testing.T) {
	NotContains(t, []int{1, 2, 3}, []int8{1, 3})
}
<commit_msg>修复测试用例判断错误<commit_after>// Copyright 2014 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package assert

import (
	"errors"
	"testing"
)

func TestGetCallerInfo(t *testing.T) {
	str := getCallerInfo()
	// NOTE:注意这里涉及到调用函数的行号信息
	if str != "TestGetCallerInfo(assert_test.go:13)" {
		t.Errorf("getCallerInfo返回的信息不正确，其返回值为：%v", str)
	}

	// 嵌套调用，第二个参数为当前的行号
	testGetCallerInfo(t, "20")
	testGetCallerInfo(t, "21")

	// 闭合函数，line为调用所在的行号。
	f := func(line string) {
		str := getCallerInfo()
		if str != "TestGetCallerInfo(assert_test.go:"+line+")" &&
			str != "TestGetCallerInfo.fun1(assert_test.go:"+line+")" { // Go1.5之后变这样了
			t.Errorf("getCallerInfo返回的信息不正确，其返回值为：%v", str)
		}
	}
	f("31") // 参数为当前等号
	f("32")
}

// 参数line，为调用此函数所在的行号。
func testGetCallerInfo(t *testing.T, line string) {
	str := getCallerInfo()
	if str != "TestGetCallerInfo(assert_test.go:"+line+")" {
		t.Error("getCallerInfo返回的信息不正确，其返回值为：%v", str)
	}

}

func TestFormatMsg(t *testing.T) {
	msg1 := []interface{}{}
	msg2 := []interface{}{[]rune("msg:%v"), 2}
	msg3 := []interface{}{"msg:%v", 3}
	msg4 := []interface{}{123, 456}

	str := formatMessage(msg1, msg2)
	if str != "msg:2" {
		t.Errorf("formatMessage(msg1,msg2)返回信息错误:[%v]", str)
	}

	str = formatMessage(nil, msg2)
	if str != "msg:2" {
		t.Errorf("formatMessage(msg1,msg2)返回信息错误:[%v]", str)
	}

	str = formatMessage(msg2, msg3)
	if str != "msg:2" {
		t.Errorf("formatMessage(msg2,msg3)返回信息错误:[%v]", str)
	}

	str = formatMessage(nil, nil)
	if str != "<未提供任何错误信息>" {
		t.Errorf("formatMessage(nil,nil)返回信息错误:[%v]", str)
	}

	str = formatMessage(nil, msg4)
	if str != "<无法正确转换错误提示信息>" {
		t.Errorf("formatMessage(nil,nil)返回信息错误:[%v]", str)
	}
}

func TestTrue(t *testing.T) {
	True(t, true)
	True(t, 1 == 1, "True(1==1) falid")
}

func TestFalse(t *testing.T) {
	False(t, false, "False falid")
	False(t, 1 == 2, "False(1==2) falid")
}

func TestNil(t *testing.T) {
	Nil(t, nil, "Nil falid")

	var v interface{}
	Nil(t, v, "Nil(v) falid")
}

func TestNotNil(t *testing.T) {
	NotNil(t, 5, "NotNil falid")

	var v interface{} = 5
	NotNil(t, v, "NotNil falid")
}

func TestEqual(t *testing.T) {
	Equal(t, 5, 5, "Equal(5,5) falid")

	var v1, v2 interface{}
	v1 = 5
	v2 = 5

	Equal(t, 5, v1)
	Equal(t, v1, v2, "Equal(v1,v2) falid")
	Equal(t, int8(126), 126)
	Equal(t, int64(126), int8(126))
	Equal(t, uint(7), int(7))
}

func TestNotEqual(t *testing.T) {
	NotEqual(t, 5, 6, "NotEqual(5,6) falid")

	var v1, v2 interface{} = 5, 6

	NotEqual(t, 5, v2, "NotEqual(5,v2) falid")
	NotEqual(t, v1, v2, "NotEqual(v1,v2) falid")
	NotEqual(t, 128, int8(127))
}

func TestEmpty(t *testing.T) {
	Empty(t, 0, "Empty(0) falid")
	Empty(t, "", "Empty(``) falid")
	Empty(t, false, "Empty(false) falid")
	Empty(t, []string{}, "Empty(slice{}) falid")
	Empty(t, []int{}, "Empty(slice{}) falid")
}

func TestNotEmpty(t *testing.T) {
	NotEmpty(t, 1, "NotEmpty(1) falid")
	NotEmpty(t, true, "NotEmpty(true) falid")
	NotEmpty(t, []string{"ab"}, "NotEmpty(slice(abc)) falid")
}

type ErrorImpl struct {
	msg string
}

func (err *ErrorImpl) Error() string {
	return err.msg
}

func TestError(t *testing.T) {
	err1 := errors.New("test")
	Error(t, err1, "Error(err) falid")

	err2 := &ErrorImpl{msg: "msg"}
	Error(t, err2, "Error(ErrorImpl) falid")
}

func TestNotError(t *testing.T) {
	NotError(t, "123", "NotError(123) falid")

	var err1 error = nil
	NotError(t, err1, "var err1 error falid")

	err2 := &ErrorImpl{msg: "msg"}
	Error(t, err2, "Error(ErrorImpl) falid")
}

func TestFileExists(t *testing.T) {
	FileExists(t, "./assert.go", "FileExists() falid")
	FileExists(t, "./", "FileExists() falid")
}

func TestFileNotExists(t *testing.T) {
	FileNotExists(t, "c:/win", "FileNotExists() falid")
	FileNotExists(t, "./abcefg/", "FileNotExists() falid")
}

func TestPanic(t *testing.T) {
	f1 := func() {
		panic("panic")
	}

	Panic(t, f1)
}

func TestNotPanic(t *testing.T) {
	f1 := func() {
	}

	NotPanic(t, f1)
}

func TestContains(t *testing.T) {
	Contains(t, []int{1, 2, 3}, []int8{1, 2})
}

func TestNotContains(t *testing.T) {
	NotContains(t, []int{1, 2, 3}, []int8{1, 3})
}
<|endoftext|>


<commit_before>// Package apexlogs implements a handler for Apex Logs https://apex.sh/logs/.
package apexlogs

import (
	"context"
	stdlog "log"
	"net/http"
	"os"
	"sync"

	"github.com/tj/go-buffer"

	"github.com/apex/log"
	"github.com/apex/logs/go/logs"
)

// logger instance.
var logger = stdlog.New(os.Stderr, "buffer ", stdlog.LstdFlags)

// levelMap is a mapping of severity levels.
var levelMap = map[log.Level]string{
	log.DebugLevel: "debug",
	log.InfoLevel:  "info",
	log.WarnLevel:  "warning",
	log.ErrorLevel: "error",
	log.FatalLevel: "emergency",
}

// Handler implementation.
type Handler struct {
	url           string
	projectID     string
	authToken     string
	httpClient    *http.Client
	bufferOptions []buffer.Option

	once sync.Once
	b    *buffer.Buffer
	c    logs.Client
}

// Option function.
type Option func(*Handler)

// New Apex Logs handler with the url, projectID, authToken and options.
func New(url, projectID, authToken string, options ...Option) *Handler {
	var v Handler
	v.url = url
	v.projectID = projectID
	v.authToken = authToken
	for _, o := range options {
		o(&v)
	}
	return &v
}

// WithHTTPClient sets the HTTP client used for requests.
func WithHTTPClient(client *http.Client) Option {
	return func(v *Handler) {
		v.httpClient = client
	}
}

// WithBufferOptions sets options for the underlying buffer used to batch logs.
func WithBufferOptions(options ...buffer.Option) Option {
	return func(v *Handler) {
		v.bufferOptions = options
	}
}

// init the client and buffer.
func (h *Handler) init() {
	h.c = logs.Client{
		URL:        h.url,
		HTTPClient: h.httpClient,
		AuthToken:  h.authToken,
	}

	var options []buffer.Option
	options = append(options, buffer.WithFlushHandler(h.handleFlush))
	options = append(options, buffer.WithErrorHandler(h.handleError))
	options = append(options, h.bufferOptions...)
	h.b = buffer.New(options...)
}

// HandleLog implements log.Handler.
func (h *Handler) HandleLog(e *log.Entry) error {
	h.once.Do(h.init)

	h.b.Push(logs.Event{
		Level:     levelMap[e.Level],
		Message:   e.Message,
		Fields:    map[string]interface{}(e.Fields),
		Timestamp: e.Timestamp,
	})

	return nil
}

// Flush any pending logs.
func (h *Handler) Flush() {
	h.b.Flush()
}

// Close flushes any pending logs, and waits for flushing to complete. This
// method should be called before exiting your program to ensure entries have
// flushed properly.
func (h *Handler) Close() {
	h.b.Close()
}

// handleFlush implementation.
func (h *Handler) handleFlush(ctx context.Context, values []interface{}) error {
	var events []logs.Event

	for _, v := range values {
		events = append(events, v.(logs.Event))
	}

	return h.c.AddEvents(logs.AddEventsInput{
		ProjectID: h.projectID,
		Events:    events,
	})
}

// handleError implementation.
func (h *Handler) handleError(err error) {
	logger.Printf("error flushing logs: %v", err)
}
<commit_msg>docs: mention that Flush() is non-blocking now, use Close()<commit_after>// Package apexlogs implements a handler for Apex Logs https://apex.sh/logs/.
package apexlogs

import (
	"context"
	stdlog "log"
	"net/http"
	"os"
	"sync"

	"github.com/tj/go-buffer"

	"github.com/apex/log"
	"github.com/apex/logs/go/logs"
)

// logger instance.
var logger = stdlog.New(os.Stderr, "buffer ", stdlog.LstdFlags)

// levelMap is a mapping of severity levels.
var levelMap = map[log.Level]string{
	log.DebugLevel: "debug",
	log.InfoLevel:  "info",
	log.WarnLevel:  "warning",
	log.ErrorLevel: "error",
	log.FatalLevel: "emergency",
}

// Handler implementation.
type Handler struct {
	url           string
	projectID     string
	authToken     string
	httpClient    *http.Client
	bufferOptions []buffer.Option

	once sync.Once
	b    *buffer.Buffer
	c    logs.Client
}

// Option function.
type Option func(*Handler)

// New Apex Logs handler with the url, projectID, authToken and options.
func New(url, projectID, authToken string, options ...Option) *Handler {
	var v Handler
	v.url = url
	v.projectID = projectID
	v.authToken = authToken
	for _, o := range options {
		o(&v)
	}
	return &v
}

// WithHTTPClient sets the HTTP client used for requests.
func WithHTTPClient(client *http.Client) Option {
	return func(v *Handler) {
		v.httpClient = client
	}
}

// WithBufferOptions sets options for the underlying buffer used to batch logs.
func WithBufferOptions(options ...buffer.Option) Option {
	return func(v *Handler) {
		v.bufferOptions = options
	}
}

// init the client and buffer.
func (h *Handler) init() {
	h.c = logs.Client{
		URL:        h.url,
		HTTPClient: h.httpClient,
		AuthToken:  h.authToken,
	}

	var options []buffer.Option
	options = append(options, buffer.WithFlushHandler(h.handleFlush))
	options = append(options, buffer.WithErrorHandler(h.handleError))
	options = append(options, h.bufferOptions...)
	h.b = buffer.New(options...)
}

// HandleLog implements log.Handler.
func (h *Handler) HandleLog(e *log.Entry) error {
	h.once.Do(h.init)

	h.b.Push(logs.Event{
		Level:     levelMap[e.Level],
		Message:   e.Message,
		Fields:    map[string]interface{}(e.Fields),
		Timestamp: e.Timestamp,
	})

	return nil
}

// Flush any pending logs. This method is non-blocking.
func (h *Handler) Flush() {
	h.b.Flush()
}

// Close flushes any pending logs, and waits for flushing to complete. This
// method should be called before exiting your program to ensure entries have
// flushed properly.
func (h *Handler) Close() {
	h.b.Close()
}

// handleFlush implementation.
func (h *Handler) handleFlush(ctx context.Context, values []interface{}) error {
	var events []logs.Event

	for _, v := range values {
		events = append(events, v.(logs.Event))
	}

	return h.c.AddEvents(logs.AddEventsInput{
		ProjectID: h.projectID,
		Events:    events,
	})
}

// handleError implementation.
func (h *Handler) handleError(err error) {
	logger.Printf("error flushing logs: %v", err)
}
<|endoftext|>