<commit_before>package log

import (
	"bytes"
	"fmt"
	"regexp"
	"strings"
	"text/template"
	"text/template/parse"

	"github.com/Masterminds/sprig/v3"

	"github.com/grafana/loki/pkg/logqlmodel"
)

var (
	_ Stage = &LineFormatter{}
	_ Stage = &LabelsFormatter{}

	// Available map of functions for the text template engine.
	functionMap = template.FuncMap{
		// olds functions deprecated.
		"ToLower":    strings.ToLower,
		"ToUpper":    strings.ToUpper,
		"Replace":    strings.Replace,
		"Trim":       strings.Trim,
		"TrimLeft":   strings.TrimLeft,
		"TrimRight":  strings.TrimRight,
		"TrimPrefix": strings.TrimPrefix,
		"TrimSuffix": strings.TrimSuffix,
		"TrimSpace":  strings.TrimSpace,
		"regexReplaceAll": func(regex string, s string, repl string) string {
			r := regexp.MustCompile(regex)
			return r.ReplaceAllString(s, repl)
		},
		"regexReplaceAllLiteral": func(regex string, s string, repl string) string {
			r := regexp.MustCompile(regex)
			return r.ReplaceAllLiteralString(s, repl)
		},
	}

	// sprig template functions
	templateFunctions = []string{
		"lower",
		"upper",
		"title",
		"trunc",
		"substr",
		"contains",
		"hasPrefix",
		"hasSuffix",
		"indent",
		"nindent",
		"replace",
		"repeat",
		"trim",
		"trimAll",
		"trimSuffix",
		"trimPrefix",
		"int",
		"float64",
		"add",
		"sub",
		"mul",
		"div",
		"mod",
		"addf",
		"subf",
		"mulf",
		"divf",
		"max",
		"min",
		"maxf",
		"minf",
		"ceil",
		"floor",
		"round",
		"fromJson",
		"date",
		"toDate",
		"now",
		"unixEpoch",
	}
)

func init() {
	sprigFuncMap := sprig.GenericFuncMap()
	for _, v := range templateFunctions {
		if function, ok := sprigFuncMap[v]; ok {
			functionMap[v] = function
		}
	}
}

type LineFormatter struct {
	*template.Template
	buf *bytes.Buffer
}

// NewFormatter creates a new log line formatter from a given text template.
func NewFormatter(tmpl string) (*LineFormatter, error) {
	t, err := template.New("line").Option("missingkey=zero").Funcs(functionMap).Parse(tmpl)
	if err != nil {
		return nil, fmt.Errorf("invalid line template: %s", err)
	}
	return &LineFormatter{
		Template: t,
		buf:      bytes.NewBuffer(make([]byte, 4096)),
	}, nil
}

func (lf *LineFormatter) Process(line []byte, lbs *LabelsBuilder) ([]byte, bool) {
	lf.buf.Reset()
	if err := lf.Template.Execute(lf.buf, lbs.Labels().Map()); err != nil {
		lbs.SetErr(errTemplateFormat)
		return line, true
	}
	// todo(cyriltovena): we might want to reuse the input line or a bytes buffer.
	res := make([]byte, len(lf.buf.Bytes()))
	copy(res, lf.buf.Bytes())
	return res, true
}

func (lf *LineFormatter) RequiredLabelNames() []string {
	return uniqueString(listNodeFields(lf.Root))
}

func listNodeFields(node parse.Node) []string {
	var res []string
	if node.Type() == parse.NodeAction {
		res = append(res, listNodeFieldsFromPipe(node.(*parse.ActionNode).Pipe)...)
	}
	res = append(res, listNodeFieldsFromBranch(node)...)
	if ln, ok := node.(*parse.ListNode); ok {
		for _, n := range ln.Nodes {
			res = append(res, listNodeFields(n)...)
		}
	}
	return res
}

func listNodeFieldsFromBranch(node parse.Node) []string {
	var res []string
	var b parse.BranchNode
	switch node.Type() {
	case parse.NodeIf:
		b = node.(*parse.IfNode).BranchNode
	case parse.NodeWith:
		b = node.(*parse.WithNode).BranchNode
	case parse.NodeRange:
		b = node.(*parse.RangeNode).BranchNode
	default:
		return res
	}
	if b.Pipe != nil {
		res = append(res, listNodeFieldsFromPipe(b.Pipe)...)
	}
	if b.List != nil {
		res = append(res, listNodeFields(b.List)...)
	}
	if b.ElseList != nil {
		res = append(res, listNodeFields(b.ElseList)...)
	}
	return res
}

func listNodeFieldsFromPipe(p *parse.PipeNode) []string {
	var res []string
	for _, c := range p.Cmds {
		for _, a := range c.Args {
			if f, ok := a.(*parse.FieldNode); ok {
				res = append(res, f.Ident...)
			}
		}
	}
	return res
}

// LabelFmt is a configuration struct for formatting a label.
type LabelFmt struct {
	Name  string
	Value string

	Rename bool
}

// NewRenameLabelFmt creates a configuration to rename a label.
func NewRenameLabelFmt(dst, target string) LabelFmt {
	return LabelFmt{
		Name:   dst,
		Rename: true,
		Value:  target,
	}
}

// NewTemplateLabelFmt creates a configuration to format a label using text template.
func NewTemplateLabelFmt(dst, template string) LabelFmt {
	return LabelFmt{
		Name:   dst,
		Rename: false,
		Value:  template,
	}
}

type labelFormatter struct {
	tmpl *template.Template
	LabelFmt
}

type LabelsFormatter struct {
	formats []labelFormatter
	buf     *bytes.Buffer
}

// NewLabelsFormatter creates a new formatter that can format multiple labels at once.
// Either by renaming or using text template.
// It is not allowed to reformat the same label twice within the same formatter.
func NewLabelsFormatter(fmts []LabelFmt) (*LabelsFormatter, error) {
	if err := validate(fmts); err != nil {
		return nil, err
	}
	formats := make([]labelFormatter, 0, len(fmts))
	for _, fm := range fmts {
		toAdd := labelFormatter{LabelFmt: fm}
		if !fm.Rename {
			t, err := template.New("label").Option("missingkey=zero").Funcs(functionMap).Parse(fm.Value)
			if err != nil {
				return nil, fmt.Errorf("invalid template for label '%s': %s", fm.Name, err)
			}
			toAdd.tmpl = t
		}
		formats = append(formats, toAdd)
	}
	return &LabelsFormatter{
		formats: formats,
		buf:     bytes.NewBuffer(make([]byte, 1024)),
	}, nil
}

func validate(fmts []LabelFmt) error {
	// it would be too confusing to rename and change the same label value.
	// To avoid confusion we allow to have a label name only once per stage.
	uniqueLabelName := map[string]struct{}{}
	for _, f := range fmts {
		if f.Name == logqlmodel.ErrorLabel {
			return fmt.Errorf("%s cannot be formatted", f.Name)
		}
		if _, ok := uniqueLabelName[f.Name]; ok {
			return fmt.Errorf("multiple label name '%s' not allowed in a single format operation", f.Name)
		}
		uniqueLabelName[f.Name] = struct{}{}
	}
	return nil
}

func (lf *LabelsFormatter) Process(l []byte, lbs *LabelsBuilder) ([]byte, bool) {
	var data interface{}
	for _, f := range lf.formats {
		if f.Rename {
			v, ok := lbs.Get(f.Value)
			if ok {
				lbs.Set(f.Name, v)
				lbs.Del(f.Value)
			}
			continue
		}
		lf.buf.Reset()
		if data == nil {
			data = lbs.Labels().Map()
		}
		if err := f.tmpl.Execute(lf.buf, data); err != nil {
			lbs.SetErr(errTemplateFormat)
			continue
		}
		lbs.Set(f.Name, lf.buf.String())
	}
	return l, true
}

func (lf *LabelsFormatter) RequiredLabelNames() []string {
	var names []string
	for _, fm := range lf.formats {
		if fm.Rename {
			names = append(names, fm.Value)
			continue
		}
		names = append(names, listNodeFields(fm.tmpl.Root)...)
	}
	return uniqueString(names)
}

func trunc(c int, s string) string {
	runes := []rune(s)
	l := len(runes)
	if c < 0 && l+c > 0 {
		return string(runes[l+c:])
	}
	if c >= 0 && l > c {
		return string(runes[:c])
	}
	return s
}

// substring creates a substring of the given string.
//
// If start is < 0, this calls string[:end].
//
// If start is >= 0 and end < 0 or end bigger than s length, this calls string[start:]
//
// Otherwise, this calls string[start, end].
func substring(start, end int, s string) string {
	runes := []rune(s)
	l := len(runes)
	if end > l {
		end = l
	}
	if start > l {
		start = l
	}
	if start < 0 {
		if end < 0 {
			return ""
		}
		return string(runes[:end])
	}
	if end < 0 {
		return string(runes[start:])
	}
	if start > end {
		return ""
	}
	return string(runes[start:end])
}
<commit_msg>Improve LogQL format stages requireLabel (#4769)<commit_after>package log

import (
	"bytes"
	"fmt"
	"regexp"
	"strings"
	"text/template"
	"text/template/parse"

	"github.com/Masterminds/sprig/v3"

	"github.com/grafana/loki/pkg/logqlmodel"
)

var (
	_ Stage = &LineFormatter{}
	_ Stage = &LabelsFormatter{}

	// Available map of functions for the text template engine.
	functionMap = template.FuncMap{
		// olds functions deprecated.
		"ToLower":    strings.ToLower,
		"ToUpper":    strings.ToUpper,
		"Replace":    strings.Replace,
		"Trim":       strings.Trim,
		"TrimLeft":   strings.TrimLeft,
		"TrimRight":  strings.TrimRight,
		"TrimPrefix": strings.TrimPrefix,
		"TrimSuffix": strings.TrimSuffix,
		"TrimSpace":  strings.TrimSpace,
		"regexReplaceAll": func(regex string, s string, repl string) string {
			r := regexp.MustCompile(regex)
			return r.ReplaceAllString(s, repl)
		},
		"regexReplaceAllLiteral": func(regex string, s string, repl string) string {
			r := regexp.MustCompile(regex)
			return r.ReplaceAllLiteralString(s, repl)
		},
	}

	// sprig template functions
	templateFunctions = []string{
		"lower",
		"upper",
		"title",
		"trunc",
		"substr",
		"contains",
		"hasPrefix",
		"hasSuffix",
		"indent",
		"nindent",
		"replace",
		"repeat",
		"trim",
		"trimAll",
		"trimSuffix",
		"trimPrefix",
		"int",
		"float64",
		"add",
		"sub",
		"mul",
		"div",
		"mod",
		"addf",
		"subf",
		"mulf",
		"divf",
		"max",
		"min",
		"maxf",
		"minf",
		"ceil",
		"floor",
		"round",
		"fromJson",
		"date",
		"toDate",
		"now",
		"unixEpoch",
	}
)

func init() {
	sprigFuncMap := sprig.GenericFuncMap()
	for _, v := range templateFunctions {
		if function, ok := sprigFuncMap[v]; ok {
			functionMap[v] = function
		}
	}
}

type LineFormatter struct {
	*template.Template
	buf *bytes.Buffer
}

// NewFormatter creates a new log line formatter from a given text template.
func NewFormatter(tmpl string) (*LineFormatter, error) {
	t, err := template.New("line").Option("missingkey=zero").Funcs(functionMap).Parse(tmpl)
	if err != nil {
		return nil, fmt.Errorf("invalid line template: %w", err)
	}
	return &LineFormatter{
		Template: t,
		buf:      bytes.NewBuffer(make([]byte, 4096)),
	}, nil
}

func (lf *LineFormatter) Process(line []byte, lbs *LabelsBuilder) ([]byte, bool) {
	lf.buf.Reset()

	if err := lf.Template.Execute(lf.buf, lbs.Labels().Map()); err != nil {
		lbs.SetErr(errTemplateFormat)
		return line, true
	}
	// todo(cyriltovena): we might want to reuse the input line or a bytes buffer.
	res := make([]byte, len(lf.buf.Bytes()))
	copy(res, lf.buf.Bytes())
	return res, true
}

func (lf *LineFormatter) RequiredLabelNames() []string {
	return uniqueString(listNodeFields([]parse.Node{lf.Root}))
}

func listNodeFields(nodes []parse.Node) []string {
	var res []string
	for _, node := range nodes {
		switch node.Type() {
		case parse.NodePipe:
			res = append(res, listNodeFieldsFromPipe(node.(*parse.PipeNode))...)
		case parse.NodeAction:
			res = append(res, listNodeFieldsFromPipe(node.(*parse.ActionNode).Pipe)...)
		case parse.NodeList:
			res = append(res, listNodeFields(node.(*parse.ListNode).Nodes)...)
		case parse.NodeCommand:
			res = append(res, listNodeFields(node.(*parse.CommandNode).Args)...)
		case parse.NodeIf, parse.NodeWith, parse.NodeRange:
			res = append(res, listNodeFieldsFromBranch(node)...)
		case parse.NodeField:
			res = append(res, node.(*parse.FieldNode).Ident...)
		}
	}
	return res
}

func listNodeFieldsFromBranch(node parse.Node) []string {
	var res []string
	var b parse.BranchNode
	switch node.Type() {
	case parse.NodeIf:
		b = node.(*parse.IfNode).BranchNode
	case parse.NodeWith:
		b = node.(*parse.WithNode).BranchNode
	case parse.NodeRange:
		b = node.(*parse.RangeNode).BranchNode
	default:
		return res
	}
	if b.Pipe != nil {
		res = append(res, listNodeFieldsFromPipe(b.Pipe)...)
	}
	if b.List != nil {
		res = append(res, listNodeFields(b.List.Nodes)...)
	}
	if b.ElseList != nil {
		res = append(res, listNodeFields(b.ElseList.Nodes)...)
	}
	return res
}

func listNodeFieldsFromPipe(p *parse.PipeNode) []string {
	var res []string
	for _, c := range p.Cmds {
		res = append(res, listNodeFields(c.Args)...)
	}
	return res
}

// LabelFmt is a configuration struct for formatting a label.
type LabelFmt struct {
	Name  string
	Value string

	Rename bool
}

// NewRenameLabelFmt creates a configuration to rename a label.
func NewRenameLabelFmt(dst, target string) LabelFmt {
	return LabelFmt{
		Name:   dst,
		Rename: true,
		Value:  target,
	}
}

// NewTemplateLabelFmt creates a configuration to format a label using text template.
func NewTemplateLabelFmt(dst, template string) LabelFmt {
	return LabelFmt{
		Name:   dst,
		Rename: false,
		Value:  template,
	}
}

type labelFormatter struct {
	tmpl *template.Template
	LabelFmt
}

type LabelsFormatter struct {
	formats []labelFormatter
	buf     *bytes.Buffer
}

// NewLabelsFormatter creates a new formatter that can format multiple labels at once.
// Either by renaming or using text template.
// It is not allowed to reformat the same label twice within the same formatter.
func NewLabelsFormatter(fmts []LabelFmt) (*LabelsFormatter, error) {
	if err := validate(fmts); err != nil {
		return nil, err
	}
	formats := make([]labelFormatter, 0, len(fmts))

	for _, fm := range fmts {
		toAdd := labelFormatter{LabelFmt: fm}
		if !fm.Rename {
			t, err := template.New("label").Option("missingkey=zero").Funcs(functionMap).Parse(fm.Value)
			if err != nil {
				return nil, fmt.Errorf("invalid template for label '%s': %s", fm.Name, err)
			}
			toAdd.tmpl = t
		}
		formats = append(formats, toAdd)
	}
	return &LabelsFormatter{
		formats: formats,
		buf:     bytes.NewBuffer(make([]byte, 1024)),
	}, nil
}

func validate(fmts []LabelFmt) error {
	// it would be too confusing to rename and change the same label value.
	// To avoid confusion we allow to have a label name only once per stage.
	uniqueLabelName := map[string]struct{}{}
	for _, f := range fmts {
		if f.Name == logqlmodel.ErrorLabel {
			return fmt.Errorf("%s cannot be formatted", f.Name)
		}
		if _, ok := uniqueLabelName[f.Name]; ok {
			return fmt.Errorf("multiple label name '%s' not allowed in a single format operation", f.Name)
		}
		uniqueLabelName[f.Name] = struct{}{}
	}
	return nil
}

func (lf *LabelsFormatter) Process(l []byte, lbs *LabelsBuilder) ([]byte, bool) {
	var data interface{}
	for _, f := range lf.formats {
		if f.Rename {
			v, ok := lbs.Get(f.Value)
			if ok {
				lbs.Set(f.Name, v)
				lbs.Del(f.Value)
			}
			continue
		}
		lf.buf.Reset()
		if data == nil {
			data = lbs.Labels().Map()
		}
		if err := f.tmpl.Execute(lf.buf, data); err != nil {
			lbs.SetErr(errTemplateFormat)
			continue
		}
		lbs.Set(f.Name, lf.buf.String())
	}
	return l, true
}

func (lf *LabelsFormatter) RequiredLabelNames() []string {
	var names []string
	for _, fm := range lf.formats {
		if fm.Rename {
			names = append(names, fm.Value)
			continue
		}
		names = append(names, listNodeFields([]parse.Node{fm.tmpl.Root})...)
	}
	return uniqueString(names)
}

func trunc(c int, s string) string {
	runes := []rune(s)
	l := len(runes)
	if c < 0 && l+c > 0 {
		return string(runes[l+c:])
	}
	if c >= 0 && l > c {
		return string(runes[:c])
	}
	return s
}

// substring creates a substring of the given string.
//
// If start is < 0, this calls string[:end].
//
// If start is >= 0 and end < 0 or end bigger than s length, this calls string[start:]
//
// Otherwise, this calls string[start, end].
func substring(start, end int, s string) string {
	runes := []rune(s)
	l := len(runes)
	if end > l {
		end = l
	}
	if start > l {
		start = l
	}
	if start < 0 {
		if end < 0 {
			return ""
		}
		return string(runes[:end])
	}
	if end < 0 {
		return string(runes[start:])
	}
	if start > end {
		return ""
	}
	return string(runes[start:end])
}
<|endoftext|>


<commit_before>// Copyright 2012-2020 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains support functions for msr access for Linux.
package msr

import (
	"encoding/binary"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
)

// CPUs is a slice of the various cpus to read or write the MSR to.
type CPUs []uint64

func AllCPUs() (CPUs, error) {
	c, errs := GlobCPUs("*")
	if errs != nil || len(c) == 0 {
		return nil, fmt.Errorf("error finding all cpus, maybe try modprobe msr? : %v", errs)
	}
	return c, nil
}

// GlobCPUs allow the user to specify CPUs using a glob as one would in /dev/cpu
func GlobCPUs(g string) (CPUs, []error) {
	var hadErr bool

	f, err := filepath.Glob(filepath.Join("/dev/cpu", g, "msr"))
	if err != nil {
		return nil, []error{err}
	}

	c := make([]uint64, len(f))
	errs := make([]error, len(f))
	for i, v := range f {
		c[i], errs[i] = strconv.ParseUint(filepath.Base(filepath.Dir(v)), 0, 64)
		if errs[i] != nil {
			hadErr = true
		}
	}
	if hadErr {
		return nil, errs
	}
	return c, nil
}

// MSR is the address of the MSR we want to target.
type MSR uint32

func (m MSR) String() string {
	return fmt.Sprintf("%#x", uint32(m))
}

func (c CPUs) paths() []string {
	var p = make([]string, len(c))

	for i, v := range c {
		p[i] = filepath.Join("/dev/cpu", strconv.Itoa(int(v)), "msr")
	}
	return p
}

func (m MSR) Read(c CPUs) ([]uint64, []error) {
	var hadErr bool
	var regs = make([]uint64, len(c))

	paths := c.paths()
	f, errs := openAll(paths, os.O_RDONLY)
	if errs != nil {
		return nil, errs
	}
	errs = make([]error, len(f))
	for i := range f {
		defer f[i].Close()
		errs[i] = doIO(f[i], m, func(port *os.File) error {
			return binary.Read(port, binary.LittleEndian, &regs[i])
		})
		if errs[i] != nil {
			hadErr = true
		}
	}
	if hadErr {
		return nil, errs
	}

	return regs, nil
}

// Write writes the corresponding data to their specified msrs
func (m MSR) Write(c CPUs, data ...uint64) []error {
	var hadErr bool

	if len(data) != len(c) && len(data) != 1 {
		return []error{fmt.Errorf("mismatched lengths: cpus %v, data %v", c, data)}
	}

	paths := c.paths()
	f, errs := openAll(paths, os.O_RDWR)

	if errs != nil {
		return errs
	}
	errs = make([]error, len(f))
	for i := range f {
		defer f[i].Close()
		errs[i] = doIO(f[i], m, func(port *os.File) error {
			return binary.Write(port, binary.LittleEndian, data[i])
		})
		if errs[i] != nil {
			hadErr = true
		}
	}
	if hadErr {
		return errs
	}
	return nil
}

// TestAndSet takes a mask of bits to clear and to set, and applies them to the specified MSR in
// each of the CPUs. Note that TestAndSet does not write if the mask does not change the MSR
func (m MSR) TestAndSet(c CPUs, clearMask uint64, setMask uint64) []error {
	paths := c.paths()
	f, errs := openAll(paths, os.O_RDWR)

	if errs != nil {
		return errs
	}
	errs = make([]error, len(f))
	for i := range f {
		defer f[i].Close()
		errs[i] = doIO(f[i], m, func(port *os.File) error {
			var v uint64
			err := binary.Read(port, binary.LittleEndian, &v)
			if err != nil {
				return err
			}
			n := v & ^clearMask
			n |= setMask
			// We write only if there is a change. This is to avoid
			// cases where we try to set a lock bit again, but the bit is
			// already set
			if n != v {
				return binary.Write(port, binary.LittleEndian, n)
			}
			return nil
		})
	}
	return errs
}

func openAll(m []string, o int) ([]*os.File, []error) {
	var (
		f      = make([]*os.File, len(m))
		errs   = make([]error, len(m))
		hadErr bool
	)
	for i := range m {
		f[i], errs[i] = os.OpenFile(m[i], o, 0)
		if errs[i] != nil {
			hadErr = true
			f[i] = nil // Not sure if I need to do this, it doesn't seem guaranteed.
		}
	}
	if hadErr {
		for i := range f {
			if f[i] != nil {
				f[i].Close()
			}
		}
		return nil, errs
	}
	return f, nil
}

func doIO(msr *os.File, addr MSR, f func(*os.File) error) error {
	if _, err := msr.Seek(int64(addr), 0); err != nil {
		return fmt.Errorf("bad address %v: %v", addr, err)
	}
	return f(msr)
}
<commit_msg>Fix msr's Write function to handle one data element<commit_after>// Copyright 2012-2020 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains support functions for msr access for Linux.
package msr

import (
	"encoding/binary"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
)

// CPUs is a slice of the various cpus to read or write the MSR to.
type CPUs []uint64

func AllCPUs() (CPUs, error) {
	c, errs := GlobCPUs("*")
	if errs != nil || len(c) == 0 {
		return nil, fmt.Errorf("error finding all cpus, maybe try modprobe msr? : %v", errs)
	}
	return c, nil
}

// GlobCPUs allow the user to specify CPUs using a glob as one would in /dev/cpu
func GlobCPUs(g string) (CPUs, []error) {
	var hadErr bool

	f, err := filepath.Glob(filepath.Join("/dev/cpu", g, "msr"))
	if err != nil {
		return nil, []error{err}
	}

	c := make([]uint64, len(f))
	errs := make([]error, len(f))
	for i, v := range f {
		c[i], errs[i] = strconv.ParseUint(filepath.Base(filepath.Dir(v)), 0, 64)
		if errs[i] != nil {
			hadErr = true
		}
	}
	if hadErr {
		return nil, errs
	}
	return c, nil
}

// MSR is the address of the MSR we want to target.
type MSR uint32

func (m MSR) String() string {
	return fmt.Sprintf("%#x", uint32(m))
}

func (c CPUs) paths() []string {
	var p = make([]string, len(c))

	for i, v := range c {
		p[i] = filepath.Join("/dev/cpu", strconv.Itoa(int(v)), "msr")
	}
	return p
}

func (m MSR) Read(c CPUs) ([]uint64, []error) {
	var hadErr bool
	var regs = make([]uint64, len(c))

	paths := c.paths()
	f, errs := openAll(paths, os.O_RDONLY)
	if errs != nil {
		return nil, errs
	}
	errs = make([]error, len(f))
	for i := range f {
		defer f[i].Close()
		errs[i] = doIO(f[i], m, func(port *os.File) error {
			return binary.Read(port, binary.LittleEndian, &regs[i])
		})
		if errs[i] != nil {
			hadErr = true
		}
	}
	if hadErr {
		return nil, errs
	}

	return regs, nil
}

// Write writes the corresponding data to their specified msrs
func (m MSR) Write(c CPUs, data ...uint64) []error {
	var hadErr bool

	if len(data) == 1 {
		// Expand value to all cpus
		for i := 1; i < len(c); i++ {
			data = append(data, data[0])
		}
	}
	if len(data) != len(c) {
		return []error{fmt.Errorf("mismatched lengths: cpus %v, data %v", c, data)}
	}

	paths := c.paths()
	f, errs := openAll(paths, os.O_RDWR)

	if errs != nil {
		return errs
	}
	errs = make([]error, len(f))
	for i := range f {
		defer f[i].Close()
		errs[i] = doIO(f[i], m, func(port *os.File) error {
			return binary.Write(port, binary.LittleEndian, data[i])
		})
		if errs[i] != nil {
			hadErr = true
		}
	}
	if hadErr {
		return errs
	}
	return nil
}

// TestAndSet takes a mask of bits to clear and to set, and applies them to the specified MSR in
// each of the CPUs. Note that TestAndSet does not write if the mask does not change the MSR
func (m MSR) TestAndSet(c CPUs, clearMask uint64, setMask uint64) []error {
	paths := c.paths()
	f, errs := openAll(paths, os.O_RDWR)

	if errs != nil {
		return errs
	}
	errs = make([]error, len(f))
	for i := range f {
		defer f[i].Close()
		errs[i] = doIO(f[i], m, func(port *os.File) error {
			var v uint64
			err := binary.Read(port, binary.LittleEndian, &v)
			if err != nil {
				return err
			}
			n := v & ^clearMask
			n |= setMask
			// We write only if there is a change. This is to avoid
			// cases where we try to set a lock bit again, but the bit is
			// already set
			if n != v {
				return binary.Write(port, binary.LittleEndian, n)
			}
			return nil
		})
	}
	return errs
}

func openAll(m []string, o int) ([]*os.File, []error) {
	var (
		f      = make([]*os.File, len(m))
		errs   = make([]error, len(m))
		hadErr bool
	)
	for i := range m {
		f[i], errs[i] = os.OpenFile(m[i], o, 0)
		if errs[i] != nil {
			hadErr = true
			f[i] = nil // Not sure if I need to do this, it doesn't seem guaranteed.
		}
	}
	if hadErr {
		for i := range f {
			if f[i] != nil {
				f[i].Close()
			}
		}
		return nil, errs
	}
	return f, nil
}

func doIO(msr *os.File, addr MSR, f func(*os.File) error) error {
	if _, err := msr.Seek(int64(addr), 0); err != nil {
		return fmt.Errorf("bad address %v: %v", addr, err)
	}
	return f(msr)
}
<|endoftext|>


<commit_before>package emailnotifier

import (
	"errors"
	"fmt"
	"socialapi/config"
	socialmodels "socialapi/models"
	"socialapi/workers/emailnotifier/models"
	"socialapi/workers/helper"
	notificationmodels "socialapi/workers/notification/models"
	"time"

	"github.com/koding/logging"
	"github.com/koding/rabbitmq"
	"github.com/streadway/amqp"
)

var emailConfig = map[string]string{
	notificationmodels.NotificationContent_TYPE_COMMENT: "comment",
	notificationmodels.NotificationContent_TYPE_LIKE:    "likeActivities",
	notificationmodels.NotificationContent_TYPE_MENTION: "mention",
}

const (
	DAY           = 24 * time.Hour
	TIMEFORMAT    = "20060102"
	CACHEPREFIX   = "dailymail"
	RECIPIENTSKEY = "recipients"
)

type Controller struct {
	log      logging.Logger
	rmqConn  *amqp.Connection
	settings *models.EmailSettings
}

func (n *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	n.log.Error("an error occured: %s", err)
	delivery.Ack(false)

	return false
}

func New(rmq *rabbitmq.RabbitMQ, log logging.Logger, es *models.EmailSettings) (*Controller, error) {
	rmqConn, err := rmq.Connect("NewEmailNotifierWorkerController")
	if err != nil {
		return nil, err
	}

	nwc := &Controller{
		log:      log,
		rmqConn:  rmqConn.Conn(),
		settings: es,
	}

	return nwc, nil
}

func (n *Controller) SendInstantEmail(notification *notificationmodels.Notification) error {
	channel, err := n.rmqConn.Channel()
	if err != nil {
		return errors.New("channel connection error")
	}
	defer channel.Close()

	// fetch latest activity for checking actor
	activity, nc, err := notification.FetchLastActivity()
	if err != nil {
		return err
	}

	if !n.validNotification(activity, notification) {
		return nil
	}

	uc, err := models.FetchUserContact(notification.AccountId)
	if err != nil {
		return fmt.Errorf("an error occurred while fetching user contact: %s", err)
	}

	if !n.checkMailSettings(uc, activity, nc) {
		return nil
	}

	mc := models.NewMailerContainer()
	mc.AccountId = notification.AccountId
	mc.Activity = activity
	mc.Content = nc

	if err := mc.PrepareContainer(); err != nil {
		return err
	}

	mc.CreatedAt = notification.ActivatedAt

	tp := models.NewTemplateParser()
	tp.UserContact = uc
	body, err := tp.RenderInstantTemplate(mc)
	if err != nil {
		return fmt.Errorf("an error occurred while preparing notification email: %s", err)
	}

	tg := &models.TokenGenerator{
		UserContact:      uc,
		NotificationType: emailConfig[nc.TypeConstant],
	}

	if err := tg.CreateToken(); err != nil {
		return err
	}

	mailer := models.NewMailer()
	mailer.EmailSettings = n.settings
	mailer.UserContact = uc
	mailer.Body = body
	mailer.Subject = prepareSubject(mc)

	if err := mailer.SendMail(); err != nil {
		return err
	}

	n.log.Info("%s notified by email", uc.Username)

	return nil
}

func prepareSubject(mc *models.MailerContainer) string {
	t, err := mc.Content.GetContentType()
	if err != nil {
		return ""
	}

	return t.GetDefinition()
}

func (c *Controller) validNotification(a *notificationmodels.NotificationActivity, n *notificationmodels.Notification) bool {
	// do not notify actor for her own action
	if a.ActorId == n.AccountId {
		return false
	}

	// do not notify actor for troll activity
	acc, err := socialmodels.FetchAccountFromCache(a.ActorId)
	if err != nil {
		c.log.Error("Invalid notification: %s", err)
		return false
	}

	if acc.IsTroll {
		return false
	}

	// do not notify user when notification is not yet activated
	return !n.ActivatedAt.IsZero()
}

func (n *Controller) checkMailSettings(uc *models.UserContact, a *notificationmodels.NotificationActivity,
	nc *notificationmodels.NotificationContent) bool {
	// notifications are disabled
	if val := uc.EmailSettings.Global; !val {
		return false
	}

	notificationEnabled := checkMailSettings(nc, uc)
	// daily notifications are enabled
	if val := uc.EmailSettings.Daily; val {
		if notificationEnabled {
			go n.saveDailyMail(uc.AccountId, a.Id)
		}

		return false
	}

	// get config
	return notificationEnabled
}

func (n *Controller) saveDailyMail(accountId, activityId int64) {
	if err := saveRecipient(accountId); err != nil {
		n.log.Error("daily mail error: %s", err)
	}

	if err := saveActivity(accountId, activityId); err != nil {
		n.log.Error("daily mail error: %s", err)
	}
}

func saveRecipient(accountId int64) error {
	redisConn := helper.MustGetRedisConn()
	key := prepareRecipientsCacheKey()
	if _, err := redisConn.AddSetMembers(key, accountId); err != nil {
		return err
	}

	if err := redisConn.Expire(key, DAY); err != nil {
		return fmt.Errorf("Could not set ttl of recipients: %s", err)
	}

	return nil
}

func saveActivity(accountId, activityId int64) error {
	redisConn := helper.MustGetRedisConn()
	key := prepareSetterCacheKey(accountId)
	if _, err := redisConn.AddSetMembers(key, activityId); err != nil {
		return err
	}

	if err := redisConn.Expire(key, DAY); err != nil {
		return fmt.Errorf("Could not set ttl of activity: %s", err)
	}

	return nil
}

func checkMailSettings(nc *notificationmodels.NotificationContent, uc *models.UserContact) bool {
	switch nc.TypeConstant {
	case notificationmodels.NotificationContent_TYPE_COMMENT:
		return uc.EmailSettings.Comment
	case notificationmodels.NotificationContent_TYPE_LIKE:
		return uc.EmailSettings.Like
	case notificationmodels.NotificationContent_TYPE_MENTION:
		return uc.EmailSettings.Mention
	}

	return false
}

func prepareRecipientsCacheKey() string {
	return fmt.Sprintf("%s:%s:%s:%s",
		config.MustGet().Environment,
		CACHEPREFIX,
		RECIPIENTSKEY,
		time.Now().Format(TIMEFORMAT))
}

func prepareSetterCacheKey(accountId int64) string {
	return fmt.Sprintf("%s:%s:%d:%s",
		config.MustGet().Environment,
		CACHEPREFIX,
		accountId,
		time.Now().Format(TIMEFORMAT))
}
<commit_msg>EmailNotifier: Use FetchAccountById instead of FetchAccountFromCache<commit_after>package emailnotifier

import (
	"errors"
	"fmt"
	"socialapi/config"
	socialmodels "socialapi/models"
	"socialapi/workers/emailnotifier/models"
	"socialapi/workers/helper"
	notificationmodels "socialapi/workers/notification/models"
	"time"

	"github.com/koding/logging"
	"github.com/koding/rabbitmq"
	"github.com/streadway/amqp"
)

var emailConfig = map[string]string{
	notificationmodels.NotificationContent_TYPE_COMMENT: "comment",
	notificationmodels.NotificationContent_TYPE_LIKE:    "likeActivities",
	notificationmodels.NotificationContent_TYPE_MENTION: "mention",
}

const (
	DAY           = 24 * time.Hour
	TIMEFORMAT    = "20060102"
	CACHEPREFIX   = "dailymail"
	RECIPIENTSKEY = "recipients"
)

type Controller struct {
	log      logging.Logger
	rmqConn  *amqp.Connection
	settings *models.EmailSettings
}

func (n *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	n.log.Error("an error occured: %s", err)
	delivery.Ack(false)

	return false
}

func New(rmq *rabbitmq.RabbitMQ, log logging.Logger, es *models.EmailSettings) (*Controller, error) {
	rmqConn, err := rmq.Connect("NewEmailNotifierWorkerController")
	if err != nil {
		return nil, err
	}

	nwc := &Controller{
		log:      log,
		rmqConn:  rmqConn.Conn(),
		settings: es,
	}

	return nwc, nil
}

func (n *Controller) SendInstantEmail(notification *notificationmodels.Notification) error {
	channel, err := n.rmqConn.Channel()
	if err != nil {
		return errors.New("channel connection error")
	}
	defer channel.Close()

	// fetch latest activity for checking actor
	activity, nc, err := notification.FetchLastActivity()
	if err != nil {
		return err
	}

	if !n.validNotification(activity, notification) {
		return nil
	}

	uc, err := models.FetchUserContact(notification.AccountId)
	if err != nil {
		return fmt.Errorf("an error occurred while fetching user contact: %s", err)
	}

	if !n.checkMailSettings(uc, activity, nc) {
		return nil
	}

	mc := models.NewMailerContainer()
	mc.AccountId = notification.AccountId
	mc.Activity = activity
	mc.Content = nc

	if err := mc.PrepareContainer(); err != nil {
		return err
	}

	mc.CreatedAt = notification.ActivatedAt

	tp := models.NewTemplateParser()
	tp.UserContact = uc
	body, err := tp.RenderInstantTemplate(mc)
	if err != nil {
		return fmt.Errorf("an error occurred while preparing notification email: %s", err)
	}

	tg := &models.TokenGenerator{
		UserContact:      uc,
		NotificationType: emailConfig[nc.TypeConstant],
	}

	if err := tg.CreateToken(); err != nil {
		return err
	}

	mailer := models.NewMailer()
	mailer.EmailSettings = n.settings
	mailer.UserContact = uc
	mailer.Body = body
	mailer.Subject = prepareSubject(mc)

	if err := mailer.SendMail(); err != nil {
		return err
	}

	n.log.Info("%s notified by email", uc.Username)

	return nil
}

func prepareSubject(mc *models.MailerContainer) string {
	t, err := mc.Content.GetContentType()
	if err != nil {
		return ""
	}

	return t.GetDefinition()
}

func (c *Controller) validNotification(a *notificationmodels.NotificationActivity, n *notificationmodels.Notification) bool {
	// do not notify actor for her own action
	if a.ActorId == n.AccountId {
		return false
	}

	// do not notify actor for troll activity
	acc, err := socialmodels.FetchAccountById(a.ActorId)
	if err != nil {
		c.log.Error("Invalid notification: %s", err)
		return false
	}

	if acc.IsTroll {
		return false
	}

	// do not notify user when notification is not yet activated
	return !n.ActivatedAt.IsZero()
}

func (n *Controller) checkMailSettings(uc *models.UserContact, a *notificationmodels.NotificationActivity,
	nc *notificationmodels.NotificationContent) bool {
	// notifications are disabled
	if val := uc.EmailSettings.Global; !val {
		return false
	}

	notificationEnabled := checkMailSettings(nc, uc)
	// daily notifications are enabled
	if val := uc.EmailSettings.Daily; val {
		if notificationEnabled {
			go n.saveDailyMail(uc.AccountId, a.Id)
		}

		return false
	}

	// get config
	return notificationEnabled
}

func (n *Controller) saveDailyMail(accountId, activityId int64) {
	if err := saveRecipient(accountId); err != nil {
		n.log.Error("daily mail error: %s", err)
	}

	if err := saveActivity(accountId, activityId); err != nil {
		n.log.Error("daily mail error: %s", err)
	}
}

func saveRecipient(accountId int64) error {
	redisConn := helper.MustGetRedisConn()
	key := prepareRecipientsCacheKey()
	if _, err := redisConn.AddSetMembers(key, accountId); err != nil {
		return err
	}

	if err := redisConn.Expire(key, DAY); err != nil {
		return fmt.Errorf("Could not set ttl of recipients: %s", err)
	}

	return nil
}

func saveActivity(accountId, activityId int64) error {
	redisConn := helper.MustGetRedisConn()
	key := prepareSetterCacheKey(accountId)
	if _, err := redisConn.AddSetMembers(key, activityId); err != nil {
		return err
	}

	if err := redisConn.Expire(key, DAY); err != nil {
		return fmt.Errorf("Could not set ttl of activity: %s", err)
	}

	return nil
}

func checkMailSettings(nc *notificationmodels.NotificationContent, uc *models.UserContact) bool {
	switch nc.TypeConstant {
	case notificationmodels.NotificationContent_TYPE_COMMENT:
		return uc.EmailSettings.Comment
	case notificationmodels.NotificationContent_TYPE_LIKE:
		return uc.EmailSettings.Like
	case notificationmodels.NotificationContent_TYPE_MENTION:
		return uc.EmailSettings.Mention
	}

	return false
}

func prepareRecipientsCacheKey() string {
	return fmt.Sprintf("%s:%s:%s:%s",
		config.MustGet().Environment,
		CACHEPREFIX,
		RECIPIENTSKEY,
		time.Now().Format(TIMEFORMAT))
}

func prepareSetterCacheKey(accountId int64) string {
	return fmt.Sprintf("%s:%s:%d:%s",
		config.MustGet().Environment,
		CACHEPREFIX,
		accountId,
		time.Now().Format(TIMEFORMAT))
}
<|endoftext|>


<commit_before>package text

import (
	"strings"
)

type Mutation interface {
	Apply(text string) string
}

type Replace struct {
	Replacement string
}

type Sub struct {
	Search string
	Sub    string
}

type SubWord struct {
	Search Word
	Sub    string
}

type Append struct {
	Appendix string
}

type Prepend struct {
	Prefix string
}

func (r *Replace) Apply(text string) string {
	return r.Replacement
}

func (s *Sub) Apply(text string) string {
	return strings.Replace(text, s.Search, s.Sub, -1)
}

func (s *SubWord) Apply(text string) string {
	if text == s.Search.W {
		return s.Sub
	}

	for {
		i := s.Search.Match(text)

		if i < 0 {
			return text
		}

		sr := []rune(text)

		beginning := sr[:i]
		end := sr[i+len(s.Search.W):]

		text = string(append(append(beginning, []rune(s.Sub)...), end...))
	}

}

func (a *Append) Apply(text string) string {
	if len(a.Appendix) == 0 {
		return text
	}

	return text + a.Appendix
}

func (p *Prepend) Apply(text string) string {
	if len(p.Prefix) == 0 {
		return text
	}

	return p.Prefix + text
}
<commit_msg>Reduce number of casts and allocations in SubWord Mutation application<commit_after>package text

import (
	"strings"
)

type Mutation interface {
	Apply(text string) string
}

type Replace struct {
	Replacement string
}

type Sub struct {
	Search string
	Sub    string
}

type SubWord struct {
	Search Word
	Sub    string
}

type Append struct {
	Appendix string
}

type Prepend struct {
	Prefix string
}

func (r *Replace) Apply(text string) string {
	return r.Replacement
}

func (s *Sub) Apply(text string) string {
	return strings.Replace(text, s.Search, s.Sub, -1)
}

func (s *SubWord) Apply(text string) string {
	if text == s.Search.W {
		return s.Sub
	}
	var subr []rune
	var sr []rune

	for {
		i := s.Search.Match(text)

		if i < 0 {
			return text
		}

		sr = []rune(text)

		beginning := sr[:i]
		end := sr[i+len(s.Search.W):]

		if subr == nil {
			subr = []rune(s.Sub)
		}

		sr = append(append(beginning, subr...), end...)

		text = string(sr)
	}

}

func (a *Append) Apply(text string) string {
	if len(a.Appendix) == 0 {
		return text
	}

	return text + a.Appendix
}

func (p *Prepend) Apply(text string) string {
	if len(p.Prefix) == 0 {
		return text
	}

	return p.Prefix + text
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/cheggaaa/pb"
	"github.com/spf13/viper"
)

var runslice []string
var f *os.File
var err error
var parallel int
var noWrite = false
var isVerbose = false

//Variables , why not struct
type Variables struct {
	server  string
	command string
	config  *viper.Viper
	bar     *pb.ProgressBar
	file    *os.File
	auth    []ssh.AuthMethod
}

func main() {
	configFile := flag.String("config", "config.yaml", "Configuration file to use")
	serverName := flag.String("server", "", "Server or server group to run commands on")
	commandToRun := flag.String("command", "", "Command to run")
	outputFlag := flag.String("output", "", "Result file to write all responses/errors")
	timeoutFlag := flag.Int("timeout", 0, "Maximum allowed time for running the command")
	showNewConfig := flag.Bool("C", false, "Show a config file example")
	flag.Parse()

	conf := viper.New()
	conf.SetConfigFile(*configFile)
	err := conf.ReadInConfig()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Can't read config file: "+*configFile+"\n", err)
		fmt.Println("You can create a new config file template with -C flag.")
		os.Exit(1)
	}

	checkConfigFlag(*showNewConfig)

	wg := sync.WaitGroup{}

	if conf.GetString("parallel") != "" {
		if _, err := strconv.Atoi(conf.GetString("parallel")); err == nil {
			parallel, _ = strconv.Atoi(conf.GetString("parallel"))
		} else {
			fmt.Fprintln(os.Stderr, "Didn't understand 'parallel' line from config file, give me numbers!")
			os.Exit(1)
		}
	} else {
		parallel = 1
	}

	wg.Add(parallel)
	s := make(chan string, parallel)

	if conf.GetBool("verbose") {
		isVerbose = true
	}

	//Should I write it to a file?
	if *outputFlag == "" {
		if conf.GetString("outputfile") == "" {
			noWrite = true
		} else {
			*outputFlag = conf.GetString("outputfile")
		}
	}

	if !noWrite {
		f, err = os.OpenFile(*outputFlag, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			fmt.Fprintln(os.Stderr, "ERROR: Can't write to the output file!\n", err)
			os.Exit(1)
		}
		defer defClose(f)
	}

	if *serverName == "" {
		sendError("server/group name")
	}
	//Checking if it is a group
	if strings.Contains(*serverName, "group:") {
		tempgrp := strings.Replace(*serverName, "group:", "groups.", -1)
		if len(conf.GetStringSlice(tempgrp)) == 0 {
			sendError("valid group name")
		}
		runslice = conf.GetStringSlice(tempgrp)
	} else {
		runslice = append(runslice, *serverName)
	}

	// Is there a command?
	if *commandToRun == "" {
		sendError("command")
	}

	//Configure the status bar
	bar := pb.New(len(runslice))
	if !isVerbose {
		bar.ShowTimeLeft = false
		bar.ShowSpeed = false
		bar.Start()
	}

	//Shit is getting real
	for i := 0; i < parallel; i++ {
		go func() {
			defer wg.Done()
			for server := range s {
				mySet := &Variables{
					server:  server,
					command: *commandToRun,
					config:  conf,
					bar:     nil,
					file:    f,
					auth:    []ssh.AuthMethod{PublicKeyFile(conf.GetString("sshkey"))},
				}
				if isVerbose {
					response, err := checkTimeout(mySet, time.Duration(int64(*timeoutFlag)))
					if err != nil {
						fmt.Println("ERROR -", server, ":", err)
					} else {
						fmt.Fprint(os.Stdout, server, ":", response)
					}
				} else {
					mySet.bar = bar
					checkTimeout(mySet, time.Duration(int64(*timeoutFlag)))
				}
			}
		}()
	}

	//Feed the channel
	for _, server := range runslice {
		s <- server
	}

	close(s)

	wg.Wait()
}

//runCommand ... as name suggests , might be something about running a command
func runCommand(set *Variables) (response string, err error) {

	if set.bar != nil {
		set.bar.Increment()
	}

	sshConfig := &ssh.ClientConfig{
		User: set.config.GetString("username"),
		Auth: set.auth,
	}

	connection, err := ssh.Dial("tcp", set.server+":22", sshConfig)
	if err != nil {
		return "Failed to dial", err
	}

	session, err := connection.NewSession()
	if err != nil {
		return "Failed to create session", err
	}

	var answer []byte

	answer, err = session.Output(set.command)
	if err != nil {
		if !noWrite {
			_, writeErr := set.file.WriteString("ERROR : " + set.server + " : " + err.Error() + "\n")
			if writeErr != nil {
				fmt.Fprintln(os.Stderr, "Couldn't write to output file")
			}
		}
		return "", err
	}

	defer session.Close()

	if !noWrite {
		_, writeErr := set.file.WriteString(set.server + " : " + string(answer))
		if writeErr != nil {
			fmt.Fprintln(os.Stderr, "Couldn't write to output file", writeErr)
		}
	}

	return string(answer), nil
}

//Just close the file please
func defClose(file *os.File) {
	err := file.Close()
	if err != nil {
		log.Fatal(err)
	}
}

//I keep saying the same thing, so why not a func?
func sendError(z string) {
	fmt.Println("Please give a", z)
	os.Exit(1)
}

func checkTimeout(set *Variables, timeout time.Duration) (response string, err error) {

	if timeout == 0 {
		timeout = 999999999999
	}

	var resp string

	timechannel := make(chan bool, 1)
	answerchannel := make(chan bool, 1)
	go func() {
		time.Sleep(time.Second * timeout)
		timechannel <- true
	}()

	go func() {
		resp, err = runCommand(set)
		answerchannel <- true
	}()

	select {
	case t := <-timechannel:
		timeError := errors.New("Taking more than timeout, aborted")
		_ = t
		return "", timeError
	case o := <-answerchannel:
		_ = o
		return resp, err
	}
}

//PublicKeyFile is reading file and shitting as AuthMethod
func PublicKeyFile(file string) ssh.AuthMethod {
	buffer, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}

	key, err := ssh.ParsePrivateKey(buffer)
	if err != nil {
		return nil
	}
	return ssh.PublicKeys(key)
}

func checkConfigFlag(display bool) {
	if display {
		fmt.Println(`
#Config file example
sshkey: /home/user/.ssh/id_rsa
username: root
port: 22
outputfile: results.txt
verbose: true
parallel: 5

groups:
  webservers:
   - 127.0.0.1
   - 127.0.0.2
   - 127.0.0.3
   - 127.0.0.4
  dbservers:
   - 127.0.1.1
   - 127.0.2.2
   - 127.0.3.3`)
		os.Exit(0)
	}
}

<commit_msg>adding license<commit_after>package main

// gurkan.in | license: Apache License 2.0 shared on https://github.com/seqizz/go-tools/blob/master/LICENSE

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/cheggaaa/pb"
	"github.com/spf13/viper"
)

var runslice []string
var f *os.File
var err error
var parallel int
var noWrite = false
var isVerbose = false

//Variables , why not struct
type Variables struct {
	server  string
	command string
	config  *viper.Viper
	bar     *pb.ProgressBar
	file    *os.File
	auth    []ssh.AuthMethod
}

func main() {
	configFile := flag.String("config", "config.yaml", "Configuration file to use")
	serverName := flag.String("server", "", "Server or server group to run commands on")
	commandToRun := flag.String("command", "", "Command to run")
	outputFlag := flag.String("output", "", "Result file to write all responses/errors")
	timeoutFlag := flag.Int("timeout", 0, "Maximum allowed time for running the command")
	showNewConfig := flag.Bool("C", false, "Show a config file example")
	flag.Parse()

	conf := viper.New()
	conf.SetConfigFile(*configFile)
	err := conf.ReadInConfig()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Can't read config file: "+*configFile+"\n", err)
		fmt.Println("You can create a new config file template with -C flag.")
		os.Exit(1)
	}

	checkConfigFlag(*showNewConfig)

	wg := sync.WaitGroup{}

	if conf.GetString("parallel") != "" {
		if _, err := strconv.Atoi(conf.GetString("parallel")); err == nil {
			parallel, _ = strconv.Atoi(conf.GetString("parallel"))
		} else {
			fmt.Fprintln(os.Stderr, "Didn't understand 'parallel' line from config file, give me numbers!")
			os.Exit(1)
		}
	} else {
		parallel = 1
	}

	wg.Add(parallel)
	s := make(chan string, parallel)

	if conf.GetBool("verbose") {
		isVerbose = true
	}

	//Should I write it to a file?
	if *outputFlag == "" {
		if conf.GetString("outputfile") == "" {
			noWrite = true
		} else {
			*outputFlag = conf.GetString("outputfile")
		}
	}

	if !noWrite {
		f, err = os.OpenFile(*outputFlag, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
		if err != nil {
			fmt.Fprintln(os.Stderr, "ERROR: Can't write to the output file!\n", err)
			os.Exit(1)
		}
		defer defClose(f)
	}

	if *serverName == "" {
		sendError("server/group name")
	}
	//Checking if it is a group
	if strings.Contains(*serverName, "group:") {
		tempgrp := strings.Replace(*serverName, "group:", "groups.", -1)
		if len(conf.GetStringSlice(tempgrp)) == 0 {
			sendError("valid group name")
		}
		runslice = conf.GetStringSlice(tempgrp)
	} else {
		runslice = append(runslice, *serverName)
	}

	// Is there a command?
	if *commandToRun == "" {
		sendError("command")
	}

	//Configure the status bar
	bar := pb.New(len(runslice))
	if !isVerbose {
		bar.ShowTimeLeft = false
		bar.ShowSpeed = false
		bar.Start()
	}

	//Shit is getting real
	for i := 0; i < parallel; i++ {
		go func() {
			defer wg.Done()
			for server := range s {
				mySet := &Variables{
					server:  server,
					command: *commandToRun,
					config:  conf,
					bar:     nil,
					file:    f,
					auth:    []ssh.AuthMethod{PublicKeyFile(conf.GetString("sshkey"))},
				}
				if isVerbose {
					response, err := checkTimeout(mySet, time.Duration(int64(*timeoutFlag)))
					if err != nil {
						fmt.Println("ERROR -", server, ":", err)
					} else {
						fmt.Fprint(os.Stdout, server, ":", response)
					}
				} else {
					mySet.bar = bar
					checkTimeout(mySet, time.Duration(int64(*timeoutFlag)))
				}
			}
		}()
	}

	//Feed the channel
	for _, server := range runslice {
		s <- server
	}

	close(s)

	wg.Wait()
}

//runCommand ... as name suggests , might be something about running a command
func runCommand(set *Variables) (response string, err error) {

	if set.bar != nil {
		set.bar.Increment()
	}

	sshConfig := &ssh.ClientConfig{
		User: set.config.GetString("username"),
		Auth: set.auth,
	}

	connection, err := ssh.Dial("tcp", set.server+":22", sshConfig)
	if err != nil {
		return "Failed to dial", err
	}

	session, err := connection.NewSession()
	if err != nil {
		return "Failed to create session", err
	}

	var answer []byte

	answer, err = session.Output(set.command)
	if err != nil {
		if !noWrite {
			_, writeErr := set.file.WriteString("ERROR : " + set.server + " : " + err.Error() + "\n")
			if writeErr != nil {
				fmt.Fprintln(os.Stderr, "Couldn't write to output file")
			}
		}
		return "", err
	}

	defer session.Close()

	if !noWrite {
		_, writeErr := set.file.WriteString(set.server + " : " + string(answer))
		if writeErr != nil {
			fmt.Fprintln(os.Stderr, "Couldn't write to output file", writeErr)
		}
	}

	return string(answer), nil
}

//Just close the file please
func defClose(file *os.File) {
	err := file.Close()
	if err != nil {
		log.Fatal(err)
	}
}

//I keep saying the same thing, so why not a func?
func sendError(z string) {
	fmt.Println("Please give a", z)
	os.Exit(1)
}

func checkTimeout(set *Variables, timeout time.Duration) (response string, err error) {

	if timeout == 0 {
		timeout = 999999999999
	}

	var resp string

	timechannel := make(chan bool, 1)
	answerchannel := make(chan bool, 1)
	go func() {
		time.Sleep(time.Second * timeout)
		timechannel <- true
	}()

	go func() {
		resp, err = runCommand(set)
		answerchannel <- true
	}()

	select {
	case t := <-timechannel:
		timeError := errors.New("Taking more than timeout, aborted")
		_ = t
		return "", timeError
	case o := <-answerchannel:
		_ = o
		return resp, err
	}
}

//PublicKeyFile is reading file and shitting as AuthMethod
func PublicKeyFile(file string) ssh.AuthMethod {
	buffer, err := ioutil.ReadFile(file)
	if err != nil {
		return nil
	}

	key, err := ssh.ParsePrivateKey(buffer)
	if err != nil {
		return nil
	}
	return ssh.PublicKeys(key)
}

func checkConfigFlag(display bool) {
	if display {
		fmt.Println(`
#Config file example
sshkey: /home/user/.ssh/id_rsa
username: root
port: 22
outputfile: results.txt
verbose: true
parallel: 5

groups:
  webservers:
   - 127.0.0.1
   - 127.0.0.2
   - 127.0.0.3
   - 127.0.0.4
  dbservers:
   - 127.0.1.1
   - 127.0.2.2
   - 127.0.3.3`)
		os.Exit(0)
	}
}

<|endoftext|>


<commit_before>package main

/// had to have a `util` file.
/// Because pragmatism.
/// Because irony.

import (
  "os"
  "fmt"
  "time"
  "regexp"
  "runtime"
  "github.com/parnurzeal/gorequest"
  "github.com/olekukonko/tablewriter"
  "github.com/briandowns/spinner"
)

func AugmentRequest(c *gorequest.SuperAgent, cfg *Config) *gorequest.SuperAgent {
  return c.
    AddCookie(cfg.GetAuthCookie()).
    Set("Content-type","application/json").
    Set("User-Agent", "NelsonCLI/0.1."+globalBuildVersion+" ("+runtime.GOOS+")").
    Timeout(15*time.Second).
    SetCurlCommand(false).
    SetDebug(globalEnableDebug)
}

func RenderTableToStdout(headers []string, data [][]string){
  table := tablewriter.NewWriter(os.Stdout)
  table.SetHeader(headers)
  table.SetBorders(tablewriter.Border{Left: false, Top: false, Right: false, Bottom: false})
  table.SetHeaderLine(false)
  table.SetRowLine(false)
  table.SetColWidth(100)
  table.SetColumnSeparator("")
  table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
  table.SetAlignment(tablewriter.ALIGN_LEFT)
  table.AppendBulk(data) // Add Bulk Data
  table.Render()
}

func JavaEpochToDateStr(long int64) string {
  t := time.Unix(0, long*int64(time.Millisecond))
  return t.Format(time.RFC3339)
}

func ProgressIndicator() *spinner.Spinner {
  s := spinner.New(spinner.CharSets[9], 100*time.Millisecond)
  s.Color("green")
  return s
}

func PrintTerminalErrors(errs []error){
  for i, j := 0, len(errs)-1; i < j; i, j = i+1, j-1 {
    errs[i], errs[j] = errs[j], errs[i]
  }

  for _,e := range errs {
    fmt.Println(e)
  }
}

func IsValidGUID(in string) bool {
  match, _ := regexp.MatchString(`^[a-z0-9]{12,12}$`, in)
  return match
}

func isValidCommaDelimitedList(str string) bool {
  match, _ := regexp.MatchString(`^([a-z0-9\\-]+,?)+$`, str)
  return match
}

func CurrentVersion() string {
  if len(globalBuildVersion) == 0 {
    return "devel"
  } else {
    return "v"+globalBuildVersion
  }
}
<commit_msg>switch user agent depending on devel<commit_after>package main

/// had to have a `util` file.
/// Because pragmatism.
/// Because irony.

import (
  "os"
  "fmt"
  "time"
  "regexp"
  "runtime"
  "github.com/parnurzeal/gorequest"
  "github.com/olekukonko/tablewriter"
  "github.com/briandowns/spinner"
)

func AugmentRequest(c *gorequest.SuperAgent, cfg *Config) *gorequest.SuperAgent {
  return c.
    AddCookie(cfg.GetAuthCookie()).
    Set("Content-type","application/json").
    Set("User-Agent", CreateUserAgentString()).
    Timeout(15*time.Second).
    SetCurlCommand(false).
    SetDebug(globalEnableDebug)
}

func RenderTableToStdout(headers []string, data [][]string){
  table := tablewriter.NewWriter(os.Stdout)
  table.SetHeader(headers)
  table.SetBorders(tablewriter.Border{Left: false, Top: false, Right: false, Bottom: false})
  table.SetHeaderLine(false)
  table.SetRowLine(false)
  table.SetColWidth(100)
  table.SetColumnSeparator("")
  table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
  table.SetAlignment(tablewriter.ALIGN_LEFT)
  table.AppendBulk(data) // Add Bulk Data
  table.Render()
}

func JavaEpochToDateStr(long int64) string {
  t := time.Unix(0, long*int64(time.Millisecond))
  return t.Format(time.RFC3339)
}

func ProgressIndicator() *spinner.Spinner {
  s := spinner.New(spinner.CharSets[9], 100*time.Millisecond)
  s.Color("green")
  return s
}

func PrintTerminalErrors(errs []error){
  for i, j := 0, len(errs)-1; i < j; i, j = i+1, j-1 {
    errs[i], errs[j] = errs[j], errs[i]
  }

  for _,e := range errs {
    fmt.Println(e)
  }
}

func IsValidGUID(in string) bool {
  match, _ := regexp.MatchString(`^[a-z0-9]{12,12}$`, in)
  return match
}

func isValidCommaDelimitedList(str string) bool {
  match, _ := regexp.MatchString(`^([a-z0-9\\-]+,?)+$`, str)
  return match
}

func CurrentVersion() string {
  if len(globalBuildVersion) == 0 {
    return "devel"
  } else {
    return "v"+globalBuildVersion
  }
}

func CreateUserAgentString() string {
  var name = "NelsonCLI"
  if len(globalBuildVersion) == 0 {
    return name
  } else {
    return name + "/0.1."+globalBuildVersion+" ("+runtime.GOOS+")"
  }
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"net"
	"os"
)

type ARWServer struct {
	sessions  SessionManager
	listenner net.Listener
	events    ARWEvents
}

func (arw *ARWServer) SendResponceToUser(arwObject *ARWObject, user *User) (err error) {

	return nil
}

func (arw *ARWServer) Initialize() {

	ln, err := net.Listen("tcp", ":8081")

	if err != nil {
		fmt.Println("Error listening...")
		os.Exit(1)
	}
	fmt.Print("Initialize Success... \n\n")

	arw.listenner = ln
}

func (arw *ARWServer) ProcessEvents() {
	for {
		conn, acceptErr := arw.listenner.Accept()

		if acceptErr != nil {
			fmt.Println("Error Accepting :", acceptErr)
			os.Exit(1)
		}

		requestBytes := make([]byte, 1024)

		_, err := conn.Read(requestBytes)
		if err != nil {
			println("Write to server failed:", err.Error())
			os.Exit(1)
		}

		requestBytes = bytes.Trim(requestBytes, "\x00")
		arw.HandleRequests(requestBytes, conn)
	}
}

func (arw *ARWServer) PrivateConnection(conn net.Conn) {
	arw.sessions.StartSession(&conn)
	conn.Write([]byte("ConnectionSuccess"))
}

func (arw *ARWServer) HandleRequests(req []byte, conn net.Conn) {
	var arwObj ARWObject
	arwObj.ExtractToARWObject(req)

	if arwObj.requestName == "ConnectionSuccess" {
		fmt.Println("Connection Success")
		conn.Write(arwObj.CompressToARWObject())
	}
}

func main() {
	var arwServer ARWServer
	arwServer.Initialize()

	arwServer.ProcessEvents()
}
<commit_msg>ARWServer updated<commit_after>package main

import (
	"bytes"
	"fmt"
	"net"
	"os"
)

type ARWServer struct {
	sessions  SessionManager
	listenner net.Listener
	events    ARWEvents
}

func (arw *ARWServer) SendResponceToUser(arwObject *ARWObject, user *User) (err error) {

	return nil
}

func (arw *ARWServer) Initialize() {

	ln, err := net.Listen("tcp", ":8081")

	if err != nil {
		fmt.Println("Error listening...")
		os.Exit(1)
	}
	fmt.Print("Initialize Success... \n\n")

	arw.listenner = ln
}

func (arw *ARWServer) ProcessEvents() {
	for {
		conn, acceptErr := arw.listenner.Accept()

		if acceptErr != nil {
			fmt.Println("Error Accepting :", acceptErr)
			os.Exit(1)
		}

		requestBytes := make([]byte, 1024)

		_, err := conn.Read(requestBytes)
		if err != nil {
			println("Write to server failed:", err.Error())
			os.Exit(1)
		}

		requestBytes = bytes.Trim(requestBytes, "\x00")
		arw.HandleRequests(requestBytes, conn)
	}
}

func (arw *ARWServer) PrivateConnection(conn net.Conn) {
	arw.sessions.StartSession(&conn)
}

func (arw *ARWServer) HandleRequests(req []byte, conn net.Conn) {
	var arwObj ARWObject
	arwObj.ExtractToARWObject(req)

	if arwObj.requestName == "ConnectionSuccess" {
		fmt.Println("Connection Success")
		conn.Write(arwObj.CompressToARWObject())
	}
}

func main() {
	var arwServer ARWServer
	arwServer.Initialize()

	arwServer.ProcessEvents()
}
<|endoftext|>


<commit_before>// Copyright 2014 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package storage

import (
	"errors"
	"reflect"
	"testing"

	"github.com/cockroachdb/cockroach/gossip"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/storage/engine"
	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/stop"
)

func TestStoresAddStore(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))
	store := Store{}
	ls.AddStore(&store)
	if !ls.HasStore(store.Ident.StoreID) {
		t.Errorf("expected local sender to contain storeID=%d", store.Ident.StoreID)
	}
	if ls.HasStore(store.Ident.StoreID + 1) {
		t.Errorf("expected local sender to not contain storeID=%d", store.Ident.StoreID+1)
	}
}

func TestStoresRemoveStore(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))

	storeID := roachpb.StoreID(89)

	ls.AddStore(&Store{Ident: roachpb.StoreIdent{StoreID: storeID}})

	ls.RemoveStore(&Store{Ident: roachpb.StoreIdent{StoreID: storeID}})

	if ls.HasStore(storeID) {
		t.Errorf("expted local sender to remove storeID=%d", storeID)
	}
}

func TestStoresGetStoreCount(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))
	if ls.GetStoreCount() != 0 {
		t.Errorf("expected 0 stores in new local sender")
	}

	expectedCount := 10
	for i := 0; i < expectedCount; i++ {
		ls.AddStore(&Store{Ident: roachpb.StoreIdent{StoreID: roachpb.StoreID(i)}})
	}
	if count := ls.GetStoreCount(); count != expectedCount {
		t.Errorf("expected store count to be %d but was %d", expectedCount, count)
	}
}

func TestStoresVisitStores(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))
	numStores := 10
	for i := 0; i < numStores; i++ {
		ls.AddStore(&Store{Ident: roachpb.StoreIdent{StoreID: roachpb.StoreID(i)}})
	}

	visit := make([]bool, numStores)
	err := ls.VisitStores(func(s *Store) error { visit[s.Ident.StoreID] = true; return nil })
	if err != nil {
		t.Errorf("unexpected error on visit: %s", err.Error())
	}

	for i, visited := range visit {
		if !visited {
			t.Errorf("store %d was not visited", i)
		}
	}

	errBoom := errors.New("boom")
	if err := ls.VisitStores(func(s *Store) error {
		return errBoom
	}); err != errBoom {
		t.Errorf("got unexpected error %v", err)
	}
}

func TestStoresGetStore(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))
	store := Store{}
	replica := roachpb.ReplicaDescriptor{StoreID: store.Ident.StoreID}
	s, pErr := ls.GetStore(replica.StoreID)
	if s != nil || pErr == nil {
		t.Errorf("expected no stores in new local sender")
	}

	ls.AddStore(&store)
	s, pErr = ls.GetStore(replica.StoreID)
	if s == nil {
		t.Errorf("expected store")
	} else if s.Ident.StoreID != store.Ident.StoreID {
		t.Errorf("expected storeID to be %d but was %d",
			s.Ident.StoreID, store.Ident.StoreID)
	} else if pErr != nil {
		t.Errorf("expected no error, instead had err=%s", pErr)
	}
}

func TestStoresLookupReplica(t *testing.T) {
	defer leaktest.AfterTest(t)()
	stopper := stop.NewStopper()
	defer stopper.Stop()
	ctx := TestStoreContext()
	manualClock := hlc.NewManualClock(0)
	ctx.Clock = hlc.NewClock(manualClock.UnixNano)
	ls := NewStores(ctx.Clock)

	// Create two new stores with ranges we care about.
	var e [2]engine.Engine
	var s [2]*Store
	var d [2]*roachpb.RangeDescriptor
	ranges := []struct {
		storeID    roachpb.StoreID
		start, end roachpb.RKey
	}{
		{2, roachpb.RKeyMin, roachpb.RKey("c")},
		{3, roachpb.RKey("x"), roachpb.RKey("z")},
	}
	for i, rng := range ranges {
		e[i] = engine.NewInMem(roachpb.Attributes{}, 1<<20, stopper)
		ctx.Transport = NewDummyRaftTransport()
		s[i] = NewStore(ctx, e[i], &roachpb.NodeDescriptor{NodeID: 1})
		s[i].Ident.StoreID = rng.storeID

		d[i] = &roachpb.RangeDescriptor{
			RangeID:  roachpb.RangeID(i),
			StartKey: rng.start,
			EndKey:   rng.end,
			Replicas: []roachpb.ReplicaDescriptor{{StoreID: rng.storeID}},
		}
		newRng, err := NewReplica(d[i], s[i], 0)
		if err != nil {
			t.Fatal(err)
		}
		if err := s[i].AddReplicaTest(newRng); err != nil {
			t.Error(err)
		}
		ls.AddStore(s[i])
	}

	if _, r, err := ls.LookupReplica(roachpb.RKey("a"), roachpb.RKey("c")); err != nil {
		t.Error(err)
	} else if r.StoreID != s[0].Ident.StoreID {
		t.Errorf("expected store %d; got %d", s[0].Ident.StoreID, r.StoreID)
	}
	if _, r, err := ls.LookupReplica(roachpb.RKey("b"), nil); err != nil {
		t.Error(err)
	} else if r.StoreID != s[0].Ident.StoreID {
		t.Errorf("expected store %d; got %d", s[0].Ident.StoreID, r.StoreID)
	}
	if _, _, err := ls.LookupReplica(roachpb.RKey("b"), roachpb.RKey("d")); !testutils.IsError(err, "outside of bounds of range") {
		t.Errorf("got unexpected error %v", err)
	}
	if _, r, err := ls.LookupReplica(roachpb.RKey("x"), roachpb.RKey("z")); err != nil {
		t.Error(err)
	} else if r.StoreID != s[1].Ident.StoreID {
		t.Errorf("expected store %d; got %d", s[1].Ident.StoreID, r.StoreID)
	}
	if _, r, err := ls.LookupReplica(roachpb.RKey("y"), nil); err != nil {
		t.Error(err)
	} else if r.StoreID != s[1].Ident.StoreID {
		t.Errorf("expected store %d; got %d", s[1].Ident.StoreID, r.StoreID)
	}

	if desc, err := ls.FirstRange(); err != nil {
		t.Error(err)
	} else if !reflect.DeepEqual(desc, d[0]) {
		t.Fatalf("expected first range %+v; got %+v", desc, d[0])
	}
}

var storeIDAlloc roachpb.StoreID

// createStores creates a slice of count stores.
func createStores(count int, t *testing.T) (*hlc.ManualClock, []*Store, *Stores, *stop.Stopper) {
	stopper := stop.NewStopper()
	ctx := TestStoreContext()
	manualClock := hlc.NewManualClock(0)
	ctx.Clock = hlc.NewClock(manualClock.UnixNano)
	ls := NewStores(ctx.Clock)

	// Create two stores with ranges we care about.
	stores := []*Store{}
	for i := 0; i < 2; i++ {
		ctx.Transport = NewDummyRaftTransport()
		s := NewStore(ctx, engine.NewInMem(roachpb.Attributes{}, 1<<20, stopper), &roachpb.NodeDescriptor{NodeID: 1})
		storeIDAlloc++
		s.Ident.StoreID = storeIDAlloc
		stores = append(stores, s)
	}

	return manualClock, stores, ls, stopper
}

// TestStoresGossipStorage verifies reading and writing of bootstrap info.
func TestStoresGossipStorage(t *testing.T) {
	defer leaktest.AfterTest(t)()
	manual, stores, ls, stopper := createStores(2, t)
	defer stopper.Stop()
	ls.AddStore(stores[0])

	manual.Set(1)

	// Verify initial read is empty.
	var bi gossip.BootstrapInfo
	if err := ls.ReadBootstrapInfo(&bi); err != nil {
		t.Fatal(err)
	}
	if len(bi.Addresses) != 0 {
		t.Errorf("expected empty bootstrap info: %+v", bi)
	}

	// Add a fake address and write.
	manual.Increment(1)
	bi.Addresses = append(bi.Addresses, util.MakeUnresolvedAddr("tcp", "127.0.0.1:8001"))
	if err := ls.WriteBootstrapInfo(&bi); err != nil {
		t.Fatal(err)
	}

	// Verify on read.
	manual.Increment(1)
	var newBI gossip.BootstrapInfo
	if err := ls.ReadBootstrapInfo(&newBI); err != nil {
		t.Fatal(err)
	}
	if len(newBI.Addresses) != 1 {
		t.Errorf("expected single bootstrap info address: %+v", newBI)
	}

	// Add another store and verify it has bootstrap info written.
	ls.AddStore(stores[1])

	// Create a new stores object to verify read.
	ls2 := NewStores(ls.clock)
	ls2.AddStore(stores[1])
	var verifyBI gossip.BootstrapInfo
	if err := ls2.ReadBootstrapInfo(&verifyBI); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(bi, verifyBI) {
		t.Errorf("bootstrap info %+v not equal to expected %+v", verifyBI, bi)
	}
}

// TestStoresGossipStorageReadLatest verifies that the latest
// bootstrap info from multiple stores is returned on Read.
func TestStoresGossipStorageReadLatest(t *testing.T) {
	defer leaktest.AfterTest(t)()
	manual, stores, ls, stopper := createStores(2, t)
	defer stopper.Stop()
	ls.AddStore(stores[0])

	// Set clock to 1.
	manual.Set(1)

	// Add a fake address and write.
	var bi gossip.BootstrapInfo
	bi.Addresses = append(bi.Addresses, util.MakeUnresolvedAddr("tcp", "127.0.0.1:8001"))
	if err := ls.WriteBootstrapInfo(&bi); err != nil {
		t.Fatal(err)
	}

	// Now remove store 0 and add store 1.
	ls.RemoveStore(stores[0])
	ls.AddStore(stores[1])

	// Increment clock, add another address and write.
	manual.Increment(1)
	bi.Addresses = append(bi.Addresses, util.MakeUnresolvedAddr("tcp", "127.0.0.1:8002"))
	if err := ls.WriteBootstrapInfo(&bi); err != nil {
		t.Fatal(err)
	}

	// Create a new stores object to freshly read. Should get latest
	// version from store 1.
	manual.Increment(1)
	ls2 := NewStores(ls.clock)
	ls2.AddStore(stores[0])
	ls2.AddStore(stores[1])
	var verifyBI gossip.BootstrapInfo
	if err := ls2.ReadBootstrapInfo(&verifyBI); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(bi, verifyBI) {
		t.Errorf("bootstrap info %+v not equal to expected %+v", verifyBI, bi)
	}

	// Verify that stores[0], which had old info, was updated with
	// latest bootstrap info during the read.
	ls3 := NewStores(ls.clock)
	ls3.AddStore(stores[0])
	verifyBI.Reset()
	if err := ls2.ReadBootstrapInfo(&verifyBI); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(bi, verifyBI) {
		t.Errorf("bootstrap info %+v not equal to expected %+v", verifyBI, bi)
	}
}
<commit_msg>storage: minor cleanup of test<commit_after>// Copyright 2014 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package storage

import (
	"errors"
	"reflect"
	"testing"

	"github.com/cockroachdb/cockroach/gossip"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/storage/engine"
	"github.com/cockroachdb/cockroach/testutils"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/stop"
)

func TestStoresAddStore(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))
	store := Store{}
	ls.AddStore(&store)
	if !ls.HasStore(store.Ident.StoreID) {
		t.Errorf("expected local sender to contain storeID=%d", store.Ident.StoreID)
	}
	if ls.HasStore(store.Ident.StoreID + 1) {
		t.Errorf("expected local sender to not contain storeID=%d", store.Ident.StoreID+1)
	}
}

func TestStoresRemoveStore(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))

	storeID := roachpb.StoreID(89)

	ls.AddStore(&Store{Ident: roachpb.StoreIdent{StoreID: storeID}})

	ls.RemoveStore(&Store{Ident: roachpb.StoreIdent{StoreID: storeID}})

	if ls.HasStore(storeID) {
		t.Errorf("expted local sender to remove storeID=%d", storeID)
	}
}

func TestStoresGetStoreCount(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))
	if ls.GetStoreCount() != 0 {
		t.Errorf("expected 0 stores in new local sender")
	}

	expectedCount := 10
	for i := 0; i < expectedCount; i++ {
		ls.AddStore(&Store{Ident: roachpb.StoreIdent{StoreID: roachpb.StoreID(i)}})
	}
	if count := ls.GetStoreCount(); count != expectedCount {
		t.Errorf("expected store count to be %d but was %d", expectedCount, count)
	}
}

func TestStoresVisitStores(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))
	numStores := 10
	for i := 0; i < numStores; i++ {
		ls.AddStore(&Store{Ident: roachpb.StoreIdent{StoreID: roachpb.StoreID(i)}})
	}

	visit := make([]bool, numStores)
	err := ls.VisitStores(func(s *Store) error { visit[s.Ident.StoreID] = true; return nil })
	if err != nil {
		t.Errorf("unexpected error on visit: %s", err.Error())
	}

	for i, visited := range visit {
		if !visited {
			t.Errorf("store %d was not visited", i)
		}
	}

	errBoom := errors.New("boom")
	if err := ls.VisitStores(func(s *Store) error {
		return errBoom
	}); err != errBoom {
		t.Errorf("got unexpected error %v", err)
	}
}

func TestStoresGetStore(t *testing.T) {
	defer leaktest.AfterTest(t)()
	ls := NewStores(hlc.NewClock(hlc.UnixNano))
	store := Store{}
	replica := roachpb.ReplicaDescriptor{StoreID: store.Ident.StoreID}
	s, pErr := ls.GetStore(replica.StoreID)
	if s != nil || pErr == nil {
		t.Errorf("expected no stores in new local sender")
	}

	ls.AddStore(&store)
	s, pErr = ls.GetStore(replica.StoreID)
	if s == nil {
		t.Errorf("expected store")
	} else if s.Ident.StoreID != store.Ident.StoreID {
		t.Errorf("expected storeID to be %d but was %d",
			s.Ident.StoreID, store.Ident.StoreID)
	} else if pErr != nil {
		t.Errorf("expected no error, instead had err=%s", pErr)
	}
}

func TestStoresLookupReplica(t *testing.T) {
	defer leaktest.AfterTest(t)()
	stopper := stop.NewStopper()
	defer stopper.Stop()
	ctx := TestStoreContext()
	manualClock := hlc.NewManualClock(0)
	ctx.Clock = hlc.NewClock(manualClock.UnixNano)
	ls := NewStores(ctx.Clock)

	// Create two new stores with ranges we care about.
	var e [2]engine.Engine
	var s [2]*Store
	var d [2]*roachpb.RangeDescriptor
	ranges := []struct {
		storeID    roachpb.StoreID
		start, end roachpb.RKey
	}{
		{2, roachpb.RKeyMin, roachpb.RKey("c")},
		{3, roachpb.RKey("x"), roachpb.RKey("z")},
	}
	for i, rng := range ranges {
		e[i] = engine.NewInMem(roachpb.Attributes{}, 1<<20, stopper)
		ctx.Transport = NewDummyRaftTransport()
		s[i] = NewStore(ctx, e[i], &roachpb.NodeDescriptor{NodeID: 1})
		s[i].Ident.StoreID = rng.storeID

		d[i] = &roachpb.RangeDescriptor{
			RangeID:  roachpb.RangeID(i),
			StartKey: rng.start,
			EndKey:   rng.end,
			Replicas: []roachpb.ReplicaDescriptor{{StoreID: rng.storeID}},
		}
		newRng, err := NewReplica(d[i], s[i], 0)
		if err != nil {
			t.Fatal(err)
		}
		if err := s[i].AddReplicaTest(newRng); err != nil {
			t.Error(err)
		}
		ls.AddStore(s[i])
	}

	testCases := []struct {
		start, end roachpb.RKey
		expStoreID roachpb.StoreID
		expError   string
	}{
		{
			start:      roachpb.RKey("a"),
			end:        roachpb.RKey("c"),
			expStoreID: s[0].Ident.StoreID,
		},
		{
			start:      roachpb.RKey("b"),
			end:        nil,
			expStoreID: s[0].Ident.StoreID,
		},
		{
			start:    roachpb.RKey("b"),
			end:      roachpb.RKey("d"),
			expError: "outside of bounds of range",
		},
		{
			start:      roachpb.RKey("x"),
			end:        roachpb.RKey("z"),
			expStoreID: s[1].Ident.StoreID,
		},
		{
			start:      roachpb.RKey("y"),
			end:        nil,
			expStoreID: s[1].Ident.StoreID,
		},
	}
	for testIdx, tc := range testCases {
		_, r, err := ls.LookupReplica(tc.start, tc.end)
		if tc.expError != "" {
			if !testutils.IsError(err, tc.expError) {
				t.Errorf("%d: got error %v (expected %s)", testIdx, err, tc.expError)
			}
		} else if err != nil {
			t.Errorf("%d: %s", testIdx, err)
		} else if r.StoreID != tc.expStoreID {
			t.Errorf("%d: expected store %d; got %d", testIdx, tc.expStoreID, r.StoreID)
		}
	}

	if desc, err := ls.FirstRange(); err != nil {
		t.Error(err)
	} else if !reflect.DeepEqual(desc, d[0]) {
		t.Fatalf("expected first range %+v; got %+v", desc, d[0])
	}
}

var storeIDAlloc roachpb.StoreID

// createStores creates a slice of count stores.
func createStores(count int, t *testing.T) (*hlc.ManualClock, []*Store, *Stores, *stop.Stopper) {
	stopper := stop.NewStopper()
	ctx := TestStoreContext()
	manualClock := hlc.NewManualClock(0)
	ctx.Clock = hlc.NewClock(manualClock.UnixNano)
	ls := NewStores(ctx.Clock)

	// Create two stores with ranges we care about.
	stores := []*Store{}
	for i := 0; i < 2; i++ {
		ctx.Transport = NewDummyRaftTransport()
		s := NewStore(ctx, engine.NewInMem(roachpb.Attributes{}, 1<<20, stopper), &roachpb.NodeDescriptor{NodeID: 1})
		storeIDAlloc++
		s.Ident.StoreID = storeIDAlloc
		stores = append(stores, s)
	}

	return manualClock, stores, ls, stopper
}

// TestStoresGossipStorage verifies reading and writing of bootstrap info.
func TestStoresGossipStorage(t *testing.T) {
	defer leaktest.AfterTest(t)()
	manual, stores, ls, stopper := createStores(2, t)
	defer stopper.Stop()
	ls.AddStore(stores[0])

	manual.Set(1)

	// Verify initial read is empty.
	var bi gossip.BootstrapInfo
	if err := ls.ReadBootstrapInfo(&bi); err != nil {
		t.Fatal(err)
	}
	if len(bi.Addresses) != 0 {
		t.Errorf("expected empty bootstrap info: %+v", bi)
	}

	// Add a fake address and write.
	manual.Increment(1)
	bi.Addresses = append(bi.Addresses, util.MakeUnresolvedAddr("tcp", "127.0.0.1:8001"))
	if err := ls.WriteBootstrapInfo(&bi); err != nil {
		t.Fatal(err)
	}

	// Verify on read.
	manual.Increment(1)
	var newBI gossip.BootstrapInfo
	if err := ls.ReadBootstrapInfo(&newBI); err != nil {
		t.Fatal(err)
	}
	if len(newBI.Addresses) != 1 {
		t.Errorf("expected single bootstrap info address: %+v", newBI)
	}

	// Add another store and verify it has bootstrap info written.
	ls.AddStore(stores[1])

	// Create a new stores object to verify read.
	ls2 := NewStores(ls.clock)
	ls2.AddStore(stores[1])
	var verifyBI gossip.BootstrapInfo
	if err := ls2.ReadBootstrapInfo(&verifyBI); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(bi, verifyBI) {
		t.Errorf("bootstrap info %+v not equal to expected %+v", verifyBI, bi)
	}
}

// TestStoresGossipStorageReadLatest verifies that the latest
// bootstrap info from multiple stores is returned on Read.
func TestStoresGossipStorageReadLatest(t *testing.T) {
	defer leaktest.AfterTest(t)()
	manual, stores, ls, stopper := createStores(2, t)
	defer stopper.Stop()
	ls.AddStore(stores[0])

	// Set clock to 1.
	manual.Set(1)

	// Add a fake address and write.
	var bi gossip.BootstrapInfo
	bi.Addresses = append(bi.Addresses, util.MakeUnresolvedAddr("tcp", "127.0.0.1:8001"))
	if err := ls.WriteBootstrapInfo(&bi); err != nil {
		t.Fatal(err)
	}

	// Now remove store 0 and add store 1.
	ls.RemoveStore(stores[0])
	ls.AddStore(stores[1])

	// Increment clock, add another address and write.
	manual.Increment(1)
	bi.Addresses = append(bi.Addresses, util.MakeUnresolvedAddr("tcp", "127.0.0.1:8002"))
	if err := ls.WriteBootstrapInfo(&bi); err != nil {
		t.Fatal(err)
	}

	// Create a new stores object to freshly read. Should get latest
	// version from store 1.
	manual.Increment(1)
	ls2 := NewStores(ls.clock)
	ls2.AddStore(stores[0])
	ls2.AddStore(stores[1])
	var verifyBI gossip.BootstrapInfo
	if err := ls2.ReadBootstrapInfo(&verifyBI); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(bi, verifyBI) {
		t.Errorf("bootstrap info %+v not equal to expected %+v", verifyBI, bi)
	}

	// Verify that stores[0], which had old info, was updated with
	// latest bootstrap info during the read.
	ls3 := NewStores(ls.clock)
	ls3.AddStore(stores[0])
	verifyBI.Reset()
	if err := ls2.ReadBootstrapInfo(&verifyBI); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(bi, verifyBI) {
		t.Errorf("bootstrap info %+v not equal to expected %+v", verifyBI, bi)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 by Richard A. Wilkes. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, version 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/.
//
// This Source Code Form is "Incompatible With Secondary Licenses", as
// defined by the Mozilla Public License, version 2.0.

package ui

import (
	"fmt"
	"github.com/richardwilkes/geom"
	"github.com/richardwilkes/ui/event"
	"github.com/richardwilkes/ui/keys"
	"syscall"
	"unsafe"
)

// #cgo linux LDFLAGS: -lX11 -lcairo
// #include <X11/Xlib.h>
// #include <X11/keysym.h>
// #include <X11/Xutil.h>
// #include <cairo/cairo.h>
// #include <cairo/cairo-xlib.h>
import "C"

var (
	running             bool
	awaitingQuit        bool
	xWindowCount        int
	xDisplay            *C.Display
	wmProtocolsAtom     C.Atom
	wmDeleteAtom        C.Atom
	goTaskAtom          C.Atom
	quitting            bool
	lastMouseDownWindow platformWindow
	lastMouseDownButton = -1
)

func platformStartUserInterface() {
	C.XInitThreads()
	if xDisplay = C.XOpenDisplay(nil); xDisplay == nil {
		panic("Failed to open the X11 display")
	}
	wmProtocolsAtom = C.XInternAtom(xDisplay, C.CString("WM_PROTOCOLS"), C.False)
	wmDeleteAtom = C.XInternAtom(xDisplay, C.CString("WM_DELETE_WINDOW"), C.False)
	goTaskAtom = C.XInternAtom(xDisplay, C.CString("GoTask"), C.False)
	running = true
	appWillFinishStartup()
	appDidFinishStartup()
	if xWindowCount == 0 && appShouldQuitAfterLastWindowClosed() {
		platformAttemptQuit()
	}
	for running {
		var event C.XEvent
		C.XNextEvent(xDisplay, &event)
		processOneEvent(&event)
	}
}

func processOneEvent(evt *C.XEvent) {
	anyEvent := (*C.XAnyEvent)(unsafe.Pointer(evt))
	window := platformWindow(uintptr(anyEvent.window))
	switch anyEvent._type {
	case C.KeyPress:
		processKeyEvent(evt, window, platformKeyDown)
	case C.KeyRelease:
		processKeyEvent(evt, window, platformKeyUp)
	case C.ButtonPress:
		buttonEvent := (*C.XButtonEvent)(unsafe.Pointer(evt))
		if isScrollWheelButton(buttonEvent.button) {
			var dx, dy float64
			switch buttonEvent.button {
			case 4: // Up
				dy = -1
			case 5: // Down
				dy = 1
			case 6: // Left
				dx = -1
			case 7: // Right
				dx = 1
			}
			handleWindowMouseWheelEvent(window, platformMouseWheel, convertKeyMask(buttonEvent.state), float64(buttonEvent.x), float64(buttonEvent.y), dx, dy)
		} else {
			lastMouseDownButton = getButton(buttonEvent.button)
			lastMouseDownWindow = window
			// RAW: Needs concept of click count
			handleWindowMouseEvent(window, platformMouseDown, convertKeyMask(buttonEvent.state), lastMouseDownButton, 0, float64(buttonEvent.x), float64(buttonEvent.y))
		}
	case C.ButtonRelease:
		buttonEvent := (*C.XButtonEvent)(unsafe.Pointer(evt))
		if !isScrollWheelButton(buttonEvent.button) {
			lastMouseDownButton = -1
			// RAW: Needs concept of click count
			handleWindowMouseEvent(window, platformMouseUp, convertKeyMask(buttonEvent.state), getButton(buttonEvent.button), 0, float64(buttonEvent.x), float64(buttonEvent.y))
		}
	case C.MotionNotify:
		motionEvent := (*C.XMotionEvent)(unsafe.Pointer(evt))
		if lastMouseDownButton != -1 {
			if window != lastMouseDownWindow {
				// RAW: Translate coordinates appropriately
				fmt.Println("need translation for mouse drag")
			}
			handleWindowMouseEvent(lastMouseDownWindow, platformMouseDragged, convertKeyMask(motionEvent.state), lastMouseDownButton, 0, float64(motionEvent.x), float64(motionEvent.y))
		} else {
			handleWindowMouseEvent(window, platformMouseMoved, convertKeyMask(motionEvent.state), 0, 0, float64(motionEvent.x), float64(motionEvent.y))
		}
	case C.EnterNotify:
		crossingEvent := (*C.XCrossingEvent)(unsafe.Pointer(evt))
		handleWindowMouseEvent(window, platformMouseEntered, convertKeyMask(crossingEvent.state), 0, 0, float64(crossingEvent.x), float64(crossingEvent.y))
	case C.LeaveNotify:
		crossingEvent := (*C.XCrossingEvent)(unsafe.Pointer(evt))
		handleWindowMouseEvent(window, platformMouseExited, convertKeyMask(crossingEvent.state), 0, 0, float64(crossingEvent.x), float64(crossingEvent.y))
	case C.FocusIn:
		appWillBecomeActive()
		appDidBecomeActive()
	case C.FocusOut:
		appWillResignActive()
		appDidResignActive()
	case C.Expose:
		if win, ok := windowMap[window]; ok {
			exposeEvent := (*C.XExposeEvent)(unsafe.Pointer(evt))
			gc := C.cairo_create(win.surface)
			C.cairo_set_line_width(gc, 1)
			C.cairo_rectangle(gc, C.double(exposeEvent.x), C.double(exposeEvent.y), C.double(exposeEvent.width), C.double(exposeEvent.height))
			C.cairo_clip(gc)
			drawWindow(window, gc, platformRect{x: C.double(exposeEvent.x), y: C.double(exposeEvent.y), width: C.double(exposeEvent.width), height: C.double(exposeEvent.height)}, false)
			C.cairo_destroy(gc)
		}
	case C.DestroyNotify:
		windowDidClose(window)
		if xWindowCount == 0 {
			if quitting {
				finishQuit()
			}
			if appShouldQuitAfterLastWindowClosed() {
				platformAttemptQuit()
			}
		}
	case C.ConfigureNotify:
		var other C.XEvent
		for C.XCheckTypedWindowEvent(xDisplay, anyEvent.window, C.ConfigureNotify, &other) != 0 {
			// Collect up the last resize event for this window that is already in the queue and use that one instead
			evt = &other
		}
		if win, ok := windowMap[window]; ok {
			win.ignoreRepaint = true
			configEvent := (*C.XConfigureEvent)(unsafe.Pointer(evt))
			lastKnownWindowBounds[window] = geom.Rect{Point: geom.Point{X: float64(configEvent.x), Y: float64(configEvent.y)}, Size: geom.Size{Width: float64(configEvent.width), Height: float64(configEvent.height)}}
			windowResized(window)
			win.root.ValidateLayout()
			win.ignoreRepaint = false
			size := win.ContentFrame().Size
			C.cairo_xlib_surface_set_size(win.surface, C.int(size.Width), C.int(size.Height))
		}
	case C.ClientMessage:
		clientEvent := (*C.XClientMessageEvent)(unsafe.Pointer(evt))
		switch clientEvent.message_type {
		case wmProtocolsAtom:
			if clientEvent.format == 32 {
				data := (*C.Atom)(unsafe.Pointer(&clientEvent.data))
				if *data == wmDeleteAtom {
					if windowShouldClose(window) {
						if win, ok := windowMap[window]; ok {
							win.Close()
						}
					}
				}
			}
		case goTaskAtom:
			data := (*uint64)(unsafe.Pointer(&clientEvent.data))
			dispatchTask(*data)
		}
	}
}

func processKeyEvent(evt *C.XEvent, window platformWindow, eventType platformEventType) {
	keyEvent := (*C.XKeyEvent)(unsafe.Pointer(evt))
	var buffer [5]C.char
	var keySym C.KeySym
	buffer[C.XLookupString(keyEvent, &buffer[0], C.int(len(buffer)-1), &keySym, nil)] = 0
	handleWindowKeyEvent(window, eventType, convertKeyMask(keyEvent.state), int(keySym), &buffer[0], false)
}

func paintWindow(pWindow platformWindow, gc *C.cairo_t, x, y, width, height C.double, future bool) {
	C.cairo_save(gc)
	C.cairo_rectangle(gc, x, y, width, height)
	C.cairo_clip(gc)
	drawWindow(pWindow, gc, platformRect{x: C.double(x), y: C.double(y), width: C.double(width), height: C.double(height)}, false)
	C.cairo_restore(gc)
}

func platformAppName() string {
	// RAW: Implement platformAppName for Linux
	return "<unknown>"
}

func platformHideApp() {
	// RAW: Implement for Linux
}

func platformHideOtherApps() {
	// RAW: Implement for Linux
}

func platformShowAllApps() {
	// RAW: Implement for Linux
}

func platformAttemptQuit() {
	switch appShouldQuit() {
	case QuitCancel:
	case QuitLater:
		awaitingQuit = true
	default:
		initiateQuit()
	}
}

func platformAppMayQuitNow(quit bool) {
	if awaitingQuit {
		awaitingQuit = false
		if quit {
			initiateQuit()
		}
	}
}

func initiateQuit() {
	appWillQuit()
	quitting = true
	if xWindowCount > 0 {
		for _, w := range Windows() {
			w.Close()
		}
	} else {
		finishQuit()
	}
}

func finishQuit() {
	running = false
	C.XCloseDisplay(xDisplay)
	xDisplay = nil
	syscall.Exit(0)
}

func isScrollWheelButton(button C.uint) bool {
	return button > 3 && button < 8
}

func getButton(button C.uint) int {
	if button == 2 {
		return 2
	}
	if button == 3 {
		return 1
	}
	return 0
}

func convertKeyMask(state C.uint) int {
	var keyMask event.KeyMask
	if state&C.LockMask == C.LockMask {
		keyMask |= event.CapsLockKeyMask
	}
	if state&C.ShiftMask == C.ShiftMask {
		keyMask |= event.ShiftKeyMask
	}
	if state&C.ControlMask == C.ControlMask {
		keyMask |= event.ControlKeyMask
	}
	if state&C.Mod1Mask == C.Mod1Mask {
		keyMask |= event.OptionKeyMask
	}
	if state&C.Mod4Mask == C.Mod4Mask {
		keyMask |= event.CommandKeyMask
	}
	return int(keyMask)
}
<commit_msg>Remove unused import<commit_after>// Copyright (c) 2016 by Richard A. Wilkes. All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, version 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/.
//
// This Source Code Form is "Incompatible With Secondary Licenses", as
// defined by the Mozilla Public License, version 2.0.

package ui

import (
	"fmt"
	"github.com/richardwilkes/geom"
	"github.com/richardwilkes/ui/event"
	"syscall"
	"unsafe"
)

// #cgo linux LDFLAGS: -lX11 -lcairo
// #include <X11/Xlib.h>
// #include <X11/keysym.h>
// #include <X11/Xutil.h>
// #include <cairo/cairo.h>
// #include <cairo/cairo-xlib.h>
import "C"

var (
	running             bool
	awaitingQuit        bool
	xWindowCount        int
	xDisplay            *C.Display
	wmProtocolsAtom     C.Atom
	wmDeleteAtom        C.Atom
	goTaskAtom          C.Atom
	quitting            bool
	lastMouseDownWindow platformWindow
	lastMouseDownButton = -1
)

func platformStartUserInterface() {
	C.XInitThreads()
	if xDisplay = C.XOpenDisplay(nil); xDisplay == nil {
		panic("Failed to open the X11 display")
	}
	wmProtocolsAtom = C.XInternAtom(xDisplay, C.CString("WM_PROTOCOLS"), C.False)
	wmDeleteAtom = C.XInternAtom(xDisplay, C.CString("WM_DELETE_WINDOW"), C.False)
	goTaskAtom = C.XInternAtom(xDisplay, C.CString("GoTask"), C.False)
	running = true
	appWillFinishStartup()
	appDidFinishStartup()
	if xWindowCount == 0 && appShouldQuitAfterLastWindowClosed() {
		platformAttemptQuit()
	}
	for running {
		var event C.XEvent
		C.XNextEvent(xDisplay, &event)
		processOneEvent(&event)
	}
}

func processOneEvent(evt *C.XEvent) {
	anyEvent := (*C.XAnyEvent)(unsafe.Pointer(evt))
	window := platformWindow(uintptr(anyEvent.window))
	switch anyEvent._type {
	case C.KeyPress:
		processKeyEvent(evt, window, platformKeyDown)
	case C.KeyRelease:
		processKeyEvent(evt, window, platformKeyUp)
	case C.ButtonPress:
		buttonEvent := (*C.XButtonEvent)(unsafe.Pointer(evt))
		if isScrollWheelButton(buttonEvent.button) {
			var dx, dy float64
			switch buttonEvent.button {
			case 4: // Up
				dy = -1
			case 5: // Down
				dy = 1
			case 6: // Left
				dx = -1
			case 7: // Right
				dx = 1
			}
			handleWindowMouseWheelEvent(window, platformMouseWheel, convertKeyMask(buttonEvent.state), float64(buttonEvent.x), float64(buttonEvent.y), dx, dy)
		} else {
			lastMouseDownButton = getButton(buttonEvent.button)
			lastMouseDownWindow = window
			// RAW: Needs concept of click count
			handleWindowMouseEvent(window, platformMouseDown, convertKeyMask(buttonEvent.state), lastMouseDownButton, 0, float64(buttonEvent.x), float64(buttonEvent.y))
		}
	case C.ButtonRelease:
		buttonEvent := (*C.XButtonEvent)(unsafe.Pointer(evt))
		if !isScrollWheelButton(buttonEvent.button) {
			lastMouseDownButton = -1
			// RAW: Needs concept of click count
			handleWindowMouseEvent(window, platformMouseUp, convertKeyMask(buttonEvent.state), getButton(buttonEvent.button), 0, float64(buttonEvent.x), float64(buttonEvent.y))
		}
	case C.MotionNotify:
		motionEvent := (*C.XMotionEvent)(unsafe.Pointer(evt))
		if lastMouseDownButton != -1 {
			if window != lastMouseDownWindow {
				// RAW: Translate coordinates appropriately
				fmt.Println("need translation for mouse drag")
			}
			handleWindowMouseEvent(lastMouseDownWindow, platformMouseDragged, convertKeyMask(motionEvent.state), lastMouseDownButton, 0, float64(motionEvent.x), float64(motionEvent.y))
		} else {
			handleWindowMouseEvent(window, platformMouseMoved, convertKeyMask(motionEvent.state), 0, 0, float64(motionEvent.x), float64(motionEvent.y))
		}
	case C.EnterNotify:
		crossingEvent := (*C.XCrossingEvent)(unsafe.Pointer(evt))
		handleWindowMouseEvent(window, platformMouseEntered, convertKeyMask(crossingEvent.state), 0, 0, float64(crossingEvent.x), float64(crossingEvent.y))
	case C.LeaveNotify:
		crossingEvent := (*C.XCrossingEvent)(unsafe.Pointer(evt))
		handleWindowMouseEvent(window, platformMouseExited, convertKeyMask(crossingEvent.state), 0, 0, float64(crossingEvent.x), float64(crossingEvent.y))
	case C.FocusIn:
		appWillBecomeActive()
		appDidBecomeActive()
	case C.FocusOut:
		appWillResignActive()
		appDidResignActive()
	case C.Expose:
		if win, ok := windowMap[window]; ok {
			exposeEvent := (*C.XExposeEvent)(unsafe.Pointer(evt))
			gc := C.cairo_create(win.surface)
			C.cairo_set_line_width(gc, 1)
			C.cairo_rectangle(gc, C.double(exposeEvent.x), C.double(exposeEvent.y), C.double(exposeEvent.width), C.double(exposeEvent.height))
			C.cairo_clip(gc)
			drawWindow(window, gc, platformRect{x: C.double(exposeEvent.x), y: C.double(exposeEvent.y), width: C.double(exposeEvent.width), height: C.double(exposeEvent.height)}, false)
			C.cairo_destroy(gc)
		}
	case C.DestroyNotify:
		windowDidClose(window)
		if xWindowCount == 0 {
			if quitting {
				finishQuit()
			}
			if appShouldQuitAfterLastWindowClosed() {
				platformAttemptQuit()
			}
		}
	case C.ConfigureNotify:
		var other C.XEvent
		for C.XCheckTypedWindowEvent(xDisplay, anyEvent.window, C.ConfigureNotify, &other) != 0 {
			// Collect up the last resize event for this window that is already in the queue and use that one instead
			evt = &other
		}
		if win, ok := windowMap[window]; ok {
			win.ignoreRepaint = true
			configEvent := (*C.XConfigureEvent)(unsafe.Pointer(evt))
			lastKnownWindowBounds[window] = geom.Rect{Point: geom.Point{X: float64(configEvent.x), Y: float64(configEvent.y)}, Size: geom.Size{Width: float64(configEvent.width), Height: float64(configEvent.height)}}
			windowResized(window)
			win.root.ValidateLayout()
			win.ignoreRepaint = false
			size := win.ContentFrame().Size
			C.cairo_xlib_surface_set_size(win.surface, C.int(size.Width), C.int(size.Height))
		}
	case C.ClientMessage:
		clientEvent := (*C.XClientMessageEvent)(unsafe.Pointer(evt))
		switch clientEvent.message_type {
		case wmProtocolsAtom:
			if clientEvent.format == 32 {
				data := (*C.Atom)(unsafe.Pointer(&clientEvent.data))
				if *data == wmDeleteAtom {
					if windowShouldClose(window) {
						if win, ok := windowMap[window]; ok {
							win.Close()
						}
					}
				}
			}
		case goTaskAtom:
			data := (*uint64)(unsafe.Pointer(&clientEvent.data))
			dispatchTask(*data)
		}
	}
}

func processKeyEvent(evt *C.XEvent, window platformWindow, eventType platformEventType) {
	keyEvent := (*C.XKeyEvent)(unsafe.Pointer(evt))
	var buffer [5]C.char
	var keySym C.KeySym
	buffer[C.XLookupString(keyEvent, &buffer[0], C.int(len(buffer)-1), &keySym, nil)] = 0
	handleWindowKeyEvent(window, eventType, convertKeyMask(keyEvent.state), int(keySym), &buffer[0], false)
}

func paintWindow(pWindow platformWindow, gc *C.cairo_t, x, y, width, height C.double, future bool) {
	C.cairo_save(gc)
	C.cairo_rectangle(gc, x, y, width, height)
	C.cairo_clip(gc)
	drawWindow(pWindow, gc, platformRect{x: C.double(x), y: C.double(y), width: C.double(width), height: C.double(height)}, false)
	C.cairo_restore(gc)
}

func platformAppName() string {
	// RAW: Implement platformAppName for Linux
	return "<unknown>"
}

func platformHideApp() {
	// RAW: Implement for Linux
}

func platformHideOtherApps() {
	// RAW: Implement for Linux
}

func platformShowAllApps() {
	// RAW: Implement for Linux
}

func platformAttemptQuit() {
	switch appShouldQuit() {
	case QuitCancel:
	case QuitLater:
		awaitingQuit = true
	default:
		initiateQuit()
	}
}

func platformAppMayQuitNow(quit bool) {
	if awaitingQuit {
		awaitingQuit = false
		if quit {
			initiateQuit()
		}
	}
}

func initiateQuit() {
	appWillQuit()
	quitting = true
	if xWindowCount > 0 {
		for _, w := range Windows() {
			w.Close()
		}
	} else {
		finishQuit()
	}
}

func finishQuit() {
	running = false
	C.XCloseDisplay(xDisplay)
	xDisplay = nil
	syscall.Exit(0)
}

func isScrollWheelButton(button C.uint) bool {
	return button > 3 && button < 8
}

func getButton(button C.uint) int {
	if button == 2 {
		return 2
	}
	if button == 3 {
		return 1
	}
	return 0
}

func convertKeyMask(state C.uint) int {
	var keyMask event.KeyMask
	if state&C.LockMask == C.LockMask {
		keyMask |= event.CapsLockKeyMask
	}
	if state&C.ShiftMask == C.ShiftMask {
		keyMask |= event.ShiftKeyMask
	}
	if state&C.ControlMask == C.ControlMask {
		keyMask |= event.ControlKeyMask
	}
	if state&C.Mod1Mask == C.Mod1Mask {
		keyMask |= event.OptionKeyMask
	}
	if state&C.Mod4Mask == C.Mod4Mask {
		keyMask |= event.CommandKeyMask
	}
	return int(keyMask)
}
<|endoftext|>


<commit_before>package memory

import (
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/jamesrwhite/minicached/store"
)

var memoryStore = make(map[string]store.Record)
var lock = sync.RWMutex{}

func init() {
	log.Info("Initialising memory storage engine")
}

func Get(key string) (found bool, record store.Record) {
	lock.RLock()
	defer lock.RUnlock()

	value, ok := memoryStore[key]

	if ok {
		return true, value
	}

	return false, store.Record{}
}

func Set(key, value string, length, flags, ttl int64) {
	lock.Lock()
	defer lock.Unlock()

	memoryStore[key] = store.Record{
		Key: key,
		Value: value,
		Flags: flags,
		Ttl: ttl,
		Length: length,
	}
}

func Delete(key string) {
	lock.Lock()
	defer lock.Unlock()

	delete(memoryStore, key)
}

func Flush() {
	lock.Lock()
	defer lock.Unlock()

	memoryStore = make(map[string]store.Record)
}
<commit_msg>go fmt<commit_after>package memory

import (
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/jamesrwhite/minicached/store"
)

var memoryStore = make(map[string]store.Record)
var lock = sync.RWMutex{}

func init() {
	log.Info("Initialising memory storage engine")
}

func Get(key string) (found bool, record store.Record) {
	lock.RLock()
	defer lock.RUnlock()

	value, ok := memoryStore[key]

	if ok {
		return true, value
	}

	return false, store.Record{}
}

func Set(key, value string, length, flags, ttl int64) {
	lock.Lock()
	defer lock.Unlock()

	memoryStore[key] = store.Record{
		Key:    key,
		Value:  value,
		Flags:  flags,
		Ttl:    ttl,
		Length: length,
	}
}

func Delete(key string) {
	lock.Lock()
	defer lock.Unlock()

	delete(memoryStore, key)
}

func Flush() {
	lock.Lock()
	defer lock.Unlock()

	memoryStore = make(map[string]store.Record)
}
<|endoftext|>