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

// Simple logging package. It defines a type, Logger, with methods
// for formatting output. It also has a predefined 'standard' Logger
// accessible through helper functions Print[f|ln], Exit[f|ln], and
// Panic[f|ln], which are easier to use than creating a Logger manually.
// That logger writes to standard error and prints the date and time
// of each logged message.
// The Exit functions call os.Exit(1) after writing the log message.
// The Panic functions call panic after writing the log message.
package log

import (
	"bytes"
	"fmt"
	"io"
	"runtime"
	"os"
	"time"
	"sync"
)

// These flags define which text to prefix to each log entry generated by the Logger.
const (
	// Bits or'ed together to control what's printed. There is no control over the
	// order they appear (the order listed here) or the format they present (as
	// described in the comments).  A colon appears after these items:
	//	2009/0123 01:23:23.123123 /a/b/c/d.go:23: message
	Ldate         = 1 << iota // the date: 2009/0123
	Ltime                     // the time: 01:23:23
	Lmicroseconds             // microsecond resolution: 01:23:23.123123.  assumes Ltime.
	Llongfile                 // full file name and line number: /a/b/c/d.go:23
	Lshortfile                // final file name element and line number: d.go:23. overrides Llongfile
)

// A Logger represents an active logging object that generates lines of
// output to an io.Writer.  Each logging operation makes a single call to
// the Writer's Write method.  A Logger can be used simultaneously from
// multiple goroutines; it guarantees to serialize access to the Writer.
type Logger struct {
	mu     sync.Mutex // ensures atomic writes
	out    io.Writer  // destination for output
	prefix string     // prefix to write at beginning of each line
	flag   int        // properties
}

// New creates a new Logger.   The out variable sets the
// destination to which log data will be written.
// The prefix appears at the beginning of each generated log line.
// The flag argument defines the logging properties.
func New(out io.Writer, prefix string, flag int) *Logger {
	return &Logger{out: out, prefix: prefix, flag: flag}
}

var std = New(os.Stderr, "", Ldate|Ltime)

// Cheap integer to fixed-width decimal ASCII.  Give a negative width to avoid zero-padding.
// Knows the buffer has capacity.
func itoa(buf *bytes.Buffer, i int, wid int) {
	var u uint = uint(i)
	if u == 0 && wid <= 1 {
		buf.WriteByte('0')
		return
	}

	// Assemble decimal in reverse order.
	var b [32]byte
	bp := len(b)
	for ; u > 0 || wid > 0; u /= 10 {
		bp--
		wid--
		b[bp] = byte(u%10) + '0'
	}

	// avoid slicing b to avoid an allocation.
	for bp < len(b) {
		buf.WriteByte(b[bp])
		bp++
	}
}

func (l *Logger) formatHeader(buf *bytes.Buffer, ns int64, calldepth int) {
	buf.WriteString(l.prefix)
	if l.flag&(Ldate|Ltime|Lmicroseconds) != 0 {
		t := time.SecondsToLocalTime(ns / 1e9)
		if l.flag&Ldate != 0 {
			itoa(buf, int(t.Year), 4)
			buf.WriteByte('/')
			itoa(buf, int(t.Month), 2)
			buf.WriteByte('/')
			itoa(buf, int(t.Day), 2)
			buf.WriteByte(' ')
		}
		if l.flag&(Ltime|Lmicroseconds) != 0 {
			itoa(buf, int(t.Hour), 2)
			buf.WriteByte(':')
			itoa(buf, int(t.Minute), 2)
			buf.WriteByte(':')
			itoa(buf, int(t.Second), 2)
			if l.flag&Lmicroseconds != 0 {
				buf.WriteByte('.')
				itoa(buf, int(ns%1e9)/1e3, 6)
			}
			buf.WriteByte(' ')
		}
	}
	if l.flag&(Lshortfile|Llongfile) != 0 {
		_, file, line, ok := runtime.Caller(calldepth)
		if ok {
			if l.flag&Lshortfile != 0 {
				short := file
				for i := len(file) - 1; i > 0; i-- {
					if file[i] == '/' {
						short = file[i+1:]
						break
					}
				}
				file = short
			}
		} else {
			file = "???"
			line = 0
		}
		buf.WriteString(file)
		buf.WriteByte(':')
		itoa(buf, line, -1)
		buf.WriteString(": ")
	}
}

// Output writes the output for a logging event.  The string s contains
// the text to print after the prefix specified by the flags of the
// Logger.  A newline is appended if the last character of s is not
// already a newline.  Calldepth is used to recover the PC and is
// provided for generality, although at the moment on all pre-defined
// paths it will be 2.
func (l *Logger) Output(calldepth int, s string) os.Error {
	now := time.Nanoseconds() // get this early.
	buf := new(bytes.Buffer)
	l.formatHeader(buf, now, calldepth+1)
	buf.WriteString(s)
	if len(s) > 0 && s[len(s)-1] != '\n' {
		buf.WriteByte('\n')
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	_, err := l.out.Write(buf.Bytes())
	return err
}

// Printf calls l.Output to print to the logger.
// Arguments are handled in the manner of fmt.Printf.
func (l *Logger) Printf(format string, v ...interface{}) {
	l.Output(2, fmt.Sprintf(format, v...))
}

// Print calls l.Output to print to the logger.
// Arguments are handled in the manner of fmt.Print.
func (l *Logger) Print(v ...interface{}) { l.Output(2, fmt.Sprint(v...)) }

// Println calls l.Output to print to the logger.
// Arguments are handled in the manner of fmt.Println.
func (l *Logger) Println(v ...interface{}) { l.Output(2, fmt.Sprintln(v...)) }

// SetOutput sets the output destination for the standard logger.
func SetOutput(w io.Writer) {
	std.out = w
}

// SetFlags sets the output flags for the standard logger.
func SetFlags(flag int) {
	std.flag = flag
}

// SetPrefix sets the output prefix for the standard logger.
func SetPrefix(prefix string) {
	std.prefix = prefix
}

// These functions write to the standard logger.

// Print calls Output to print to the standard logger.
// Arguments are handled in the manner of fmt.Print.
func Print(v ...interface{}) {
	std.Output(2, fmt.Sprint(v...))
}

// Printf calls Output to print to the standard logger.
// Arguments are handled in the manner of fmt.Printf.
func Printf(format string, v ...interface{}) {
	std.Output(2, fmt.Sprintf(format, v...))
}

// Println calls Output to print to the standard logger.
// Arguments are handled in the manner of fmt.Println.
func Println(v ...interface{}) {
	std.Output(2, fmt.Sprintln(v...))
}

// Exit is equivalent to Print() followed by a call to os.Exit(1).
func Exit(v ...interface{}) {
	std.Output(2, fmt.Sprint(v...))
	os.Exit(1)
}

// Exitf is equivalent to Printf() followed by a call to os.Exit(1).
func Exitf(format string, v ...interface{}) {
	std.Output(2, fmt.Sprintf(format, v...))
	os.Exit(1)
}

// Exitln is equivalent to Println() followed by a call to os.Exit(1).
func Exitln(v ...interface{}) {
	std.Output(2, fmt.Sprintln(v...))
	os.Exit(1)
}

// Panic is equivalent to Print() followed by a call to panic().
func Panic(v ...interface{}) {
	s := fmt.Sprint(v...)
	std.Output(2, s)
	panic(s)
}

// Panicf is equivalent to Printf() followed by a call to panic().
func Panicf(format string, v ...interface{}) {
	s := fmt.Sprintf(format, v...)
	std.Output(2, s)
	panic(s)
}

// Panicln is equivalent to Println() followed by a call to panic().
func Panicln(v ...interface{}) {
	s := fmt.Sprintln(v...)
	std.Output(2, s)
	panic(s)
}
<commit_msg>log: add methods for exit and panic There were already functions for the standard logger; this just completes the set.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Simple logging package. It defines a type, Logger, with methods
// for formatting output. It also has a predefined 'standard' Logger
// accessible through helper functions Print[f|ln], Exit[f|ln], and
// Panic[f|ln], which are easier to use than creating a Logger manually.
// That logger writes to standard error and prints the date and time
// of each logged message.
// The Exit functions call os.Exit(1) after writing the log message.
// The Panic functions call panic after writing the log message.
package log

import (
	"bytes"
	"fmt"
	"io"
	"runtime"
	"os"
	"time"
	"sync"
)

// These flags define which text to prefix to each log entry generated by the Logger.
const (
	// Bits or'ed together to control what's printed. There is no control over the
	// order they appear (the order listed here) or the format they present (as
	// described in the comments).  A colon appears after these items:
	//	2009/0123 01:23:23.123123 /a/b/c/d.go:23: message
	Ldate         = 1 << iota // the date: 2009/0123
	Ltime                     // the time: 01:23:23
	Lmicroseconds             // microsecond resolution: 01:23:23.123123.  assumes Ltime.
	Llongfile                 // full file name and line number: /a/b/c/d.go:23
	Lshortfile                // final file name element and line number: d.go:23. overrides Llongfile
)

// A Logger represents an active logging object that generates lines of
// output to an io.Writer.  Each logging operation makes a single call to
// the Writer's Write method.  A Logger can be used simultaneously from
// multiple goroutines; it guarantees to serialize access to the Writer.
type Logger struct {
	mu     sync.Mutex // ensures atomic writes
	out    io.Writer  // destination for output
	prefix string     // prefix to write at beginning of each line
	flag   int        // properties
}

// New creates a new Logger.   The out variable sets the
// destination to which log data will be written.
// The prefix appears at the beginning of each generated log line.
// The flag argument defines the logging properties.
func New(out io.Writer, prefix string, flag int) *Logger {
	return &Logger{out: out, prefix: prefix, flag: flag}
}

var std = New(os.Stderr, "", Ldate|Ltime)

// Cheap integer to fixed-width decimal ASCII.  Give a negative width to avoid zero-padding.
// Knows the buffer has capacity.
func itoa(buf *bytes.Buffer, i int, wid int) {
	var u uint = uint(i)
	if u == 0 && wid <= 1 {
		buf.WriteByte('0')
		return
	}

	// Assemble decimal in reverse order.
	var b [32]byte
	bp := len(b)
	for ; u > 0 || wid > 0; u /= 10 {
		bp--
		wid--
		b[bp] = byte(u%10) + '0'
	}

	// avoid slicing b to avoid an allocation.
	for bp < len(b) {
		buf.WriteByte(b[bp])
		bp++
	}
}

func (l *Logger) formatHeader(buf *bytes.Buffer, ns int64, calldepth int) {
	buf.WriteString(l.prefix)
	if l.flag&(Ldate|Ltime|Lmicroseconds) != 0 {
		t := time.SecondsToLocalTime(ns / 1e9)
		if l.flag&Ldate != 0 {
			itoa(buf, int(t.Year), 4)
			buf.WriteByte('/')
			itoa(buf, int(t.Month), 2)
			buf.WriteByte('/')
			itoa(buf, int(t.Day), 2)
			buf.WriteByte(' ')
		}
		if l.flag&(Ltime|Lmicroseconds) != 0 {
			itoa(buf, int(t.Hour), 2)
			buf.WriteByte(':')
			itoa(buf, int(t.Minute), 2)
			buf.WriteByte(':')
			itoa(buf, int(t.Second), 2)
			if l.flag&Lmicroseconds != 0 {
				buf.WriteByte('.')
				itoa(buf, int(ns%1e9)/1e3, 6)
			}
			buf.WriteByte(' ')
		}
	}
	if l.flag&(Lshortfile|Llongfile) != 0 {
		_, file, line, ok := runtime.Caller(calldepth)
		if ok {
			if l.flag&Lshortfile != 0 {
				short := file
				for i := len(file) - 1; i > 0; i-- {
					if file[i] == '/' {
						short = file[i+1:]
						break
					}
				}
				file = short
			}
		} else {
			file = "???"
			line = 0
		}
		buf.WriteString(file)
		buf.WriteByte(':')
		itoa(buf, line, -1)
		buf.WriteString(": ")
	}
}

// Output writes the output for a logging event.  The string s contains
// the text to print after the prefix specified by the flags of the
// Logger.  A newline is appended if the last character of s is not
// already a newline.  Calldepth is used to recover the PC and is
// provided for generality, although at the moment on all pre-defined
// paths it will be 2.
func (l *Logger) Output(calldepth int, s string) os.Error {
	now := time.Nanoseconds() // get this early.
	buf := new(bytes.Buffer)
	l.formatHeader(buf, now, calldepth+1)
	buf.WriteString(s)
	if len(s) > 0 && s[len(s)-1] != '\n' {
		buf.WriteByte('\n')
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	_, err := l.out.Write(buf.Bytes())
	return err
}

// Printf calls l.Output to print to the logger.
// Arguments are handled in the manner of fmt.Printf.
func (l *Logger) Printf(format string, v ...interface{}) {
	l.Output(2, fmt.Sprintf(format, v...))
}

// Print calls l.Output to print to the logger.
// Arguments are handled in the manner of fmt.Print.
func (l *Logger) Print(v ...interface{}) { l.Output(2, fmt.Sprint(v...)) }

// Println calls l.Output to print to the logger.
// Arguments are handled in the manner of fmt.Println.
func (l *Logger) Println(v ...interface{}) { l.Output(2, fmt.Sprintln(v...)) }

// Exit is equivalent to l.Print() followed by a call to os.Exit(1).
func (l *Logger) Exit(v ...interface{}) {
	l.Output(2, fmt.Sprint(v...))
	os.Exit(1)
}

// Exitf is equivalent to l.Printf() followed by a call to os.Exit(1).
func (l *Logger) Exitf(format string, v ...interface{}) {
	l.Output(2, fmt.Sprintf(format, v...))
	os.Exit(1)
}

// Exitln is equivalent to l.Println() followed by a call to os.Exit(1).
func (l *Logger) Exitln(v ...interface{}) {
	l.Output(2, fmt.Sprintln(v...))
	os.Exit(1)
}

// Panic is equivalent to l.Print() followed by a call to panic().
func (l *Logger) Panic(v ...interface{}) {
	s := fmt.Sprint(v...)
	l.Output(2, s)
	panic(s)
}

// Panicf is equivalent to l.Printf() followed by a call to panic().
func (l *Logger) Panicf(format string, v ...interface{}) {
	s := fmt.Sprintf(format, v...)
	l.Output(2, s)
	panic(s)
}

// Panicln is equivalent to l.Println() followed by a call to panic().
func (l *Logger) Panicln(v ...interface{}) {
	s := fmt.Sprintln(v...)
	l.Output(2, s)
	panic(s)
}

// SetOutput sets the output destination for the standard logger.
func SetOutput(w io.Writer) {
	std.out = w
}

// SetFlags sets the output flags for the standard logger.
func SetFlags(flag int) {
	std.flag = flag
}

// SetPrefix sets the output prefix for the standard logger.
func SetPrefix(prefix string) {
	std.prefix = prefix
}

// These functions write to the standard logger.

// Print calls Output to print to the standard logger.
// Arguments are handled in the manner of fmt.Print.
func Print(v ...interface{}) {
	std.Output(2, fmt.Sprint(v...))
}

// Printf calls Output to print to the standard logger.
// Arguments are handled in the manner of fmt.Printf.
func Printf(format string, v ...interface{}) {
	std.Output(2, fmt.Sprintf(format, v...))
}

// Println calls Output to print to the standard logger.
// Arguments are handled in the manner of fmt.Println.
func Println(v ...interface{}) {
	std.Output(2, fmt.Sprintln(v...))
}

// Exit is equivalent to Print() followed by a call to os.Exit(1).
func Exit(v ...interface{}) {
	std.Output(2, fmt.Sprint(v...))
	os.Exit(1)
}

// Exitf is equivalent to Printf() followed by a call to os.Exit(1).
func Exitf(format string, v ...interface{}) {
	std.Output(2, fmt.Sprintf(format, v...))
	os.Exit(1)
}

// Exitln is equivalent to Println() followed by a call to os.Exit(1).
func Exitln(v ...interface{}) {
	std.Output(2, fmt.Sprintln(v...))
	os.Exit(1)
}

// Panic is equivalent to Print() followed by a call to panic().
func Panic(v ...interface{}) {
	s := fmt.Sprint(v...)
	std.Output(2, s)
	panic(s)
}

// Panicf is equivalent to Printf() followed by a call to panic().
func Panicf(format string, v ...interface{}) {
	s := fmt.Sprintf(format, v...)
	std.Output(2, s)
	panic(s)
}

// Panicln is equivalent to Println() followed by a call to panic().
func Panicln(v ...interface{}) {
	s := fmt.Sprintln(v...)
	std.Output(2, s)
	panic(s)
}
<|endoftext|>


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

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/scanner"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

var (
	// main operation modes
	list        = flag.Bool("l", false, "list files whose formatting differs from gofmt's")
	write       = flag.Bool("w", false, "write result to (source) file instead of stdout")
	rewriteRule = flag.String("r", "", "rewrite rule (e.g., 'α[β:len(α)] -> α[β:]')")
	simplifyAST = flag.Bool("s", false, "simplify code")

	// layout control
	comments  = flag.Bool("comments", true, "print comments")
	tabWidth  = flag.Int("tabwidth", 8, "tab width")
	tabIndent = flag.Bool("tabindent", true, "indent with tabs independent of -spaces")
	useSpaces = flag.Bool("spaces", true, "align with spaces instead of tabs")
)

var (
	fset        = token.NewFileSet()
	exitCode    = 0
	rewrite     func(*ast.File) *ast.File
	parserMode  uint
	printerMode uint
)

func report(err os.Error) {
	scanner.PrintError(os.Stderr, err)
	exitCode = 2
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: gofmt [flags] [path ...]\n")
	flag.PrintDefaults()
	os.Exit(2)
}

func initParserMode() {
	parserMode = uint(0)
	if *comments {
		parserMode |= parser.ParseComments
	}
}

func initPrinterMode() {
	printerMode = uint(0)
	if *tabIndent {
		printerMode |= printer.TabIndent
	}
	if *useSpaces {
		printerMode |= printer.UseSpaces
	}
}

func isGoFile(f *os.FileInfo) bool {
	// ignore non-Go files
	return f.IsRegular() && !strings.HasPrefix(f.Name, ".") && strings.HasSuffix(f.Name, ".go")
}

func processFile(f *os.File) os.Error {
	src, err := ioutil.ReadAll(f)
	if err != nil {
		return err
	}

	file, err := parser.ParseFile(fset, f.Name(), src, parserMode)

	if err != nil {
		return err
	}

	if rewrite != nil {
		file = rewrite(file)
	}

	if *simplifyAST {
		simplify(file)
	}

	var buf bytes.Buffer
	_, err = (&printer.Config{printerMode, *tabWidth}).Fprint(&buf, fset, file)
	if err != nil {
		return err
	}
	res := buf.Bytes()

	if !bytes.Equal(src, res) {
		// formatting has changed
		if *list {
			fmt.Fprintln(os.Stdout, f.Name())
		}
		if *write {
			err = ioutil.WriteFile(f.Name(), res, 0)
			if err != nil {
				return err
			}
		}
	}

	if !*list && !*write {
		_, err = os.Stdout.Write(res)
	}

	return err
}

func processFileByName(filename string) os.Error {
	file, err := os.Open(filename, os.O_RDONLY, 0)
	if err != nil {
		return err
	}
	defer file.Close()
	return processFile(file)
}

type fileVisitor chan os.Error

func (v fileVisitor) VisitDir(path string, f *os.FileInfo) bool {
	return true
}

func (v fileVisitor) VisitFile(path string, f *os.FileInfo) {
	if isGoFile(f) {
		v <- nil // synchronize error handler
		if err := processFileByName(path); err != nil {
			v <- err
		}
	}
}

func walkDir(path string) {
	// start an error handler
	done := make(chan bool)
	v := make(fileVisitor)
	go func() {
		for err := range v {
			if err != nil {
				report(err)
			}
		}
		done <- true
	}()
	// walk the tree
	filepath.Walk(path, v, v)
	close(v) // terminate error handler loop
	<-done   // wait for all errors to be reported
}

func main() {
	flag.Usage = usage
	flag.Parse()
	if *tabWidth < 0 {
		fmt.Fprintf(os.Stderr, "negative tabwidth %d\n", *tabWidth)
		os.Exit(2)
	}

	initParserMode()
	initPrinterMode()
	initRewrite()

	if flag.NArg() == 0 {
		if err := processFile(os.Stdin); err != nil {
			report(err)
		}
	}

	for i := 0; i < flag.NArg(); i++ {
		path := flag.Arg(i)
		switch dir, err := os.Stat(path); {
		case err != nil:
			report(err)
		case dir.IsRegular():
			if err := processFileByName(path); err != nil {
				report(err)
			}
		case dir.IsDirectory():
			walkDir(path)
		}
	}

	os.Exit(exitCode)
}
<commit_msg>gofmt: simpler walkDir<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/scanner"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

var (
	// main operation modes
	list        = flag.Bool("l", false, "list files whose formatting differs from gofmt's")
	write       = flag.Bool("w", false, "write result to (source) file instead of stdout")
	rewriteRule = flag.String("r", "", "rewrite rule (e.g., 'α[β:len(α)] -> α[β:]')")
	simplifyAST = flag.Bool("s", false, "simplify code")

	// layout control
	comments  = flag.Bool("comments", true, "print comments")
	tabWidth  = flag.Int("tabwidth", 8, "tab width")
	tabIndent = flag.Bool("tabindent", true, "indent with tabs independent of -spaces")
	useSpaces = flag.Bool("spaces", true, "align with spaces instead of tabs")
)

var (
	fset        = token.NewFileSet()
	exitCode    = 0
	rewrite     func(*ast.File) *ast.File
	parserMode  uint
	printerMode uint
)

func report(err os.Error) {
	scanner.PrintError(os.Stderr, err)
	exitCode = 2
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: gofmt [flags] [path ...]\n")
	flag.PrintDefaults()
	os.Exit(2)
}

func initParserMode() {
	parserMode = uint(0)
	if *comments {
		parserMode |= parser.ParseComments
	}
}

func initPrinterMode() {
	printerMode = uint(0)
	if *tabIndent {
		printerMode |= printer.TabIndent
	}
	if *useSpaces {
		printerMode |= printer.UseSpaces
	}
}

func isGoFile(f *os.FileInfo) bool {
	// ignore non-Go files
	return f.IsRegular() && !strings.HasPrefix(f.Name, ".") && strings.HasSuffix(f.Name, ".go")
}

func processFile(f *os.File) os.Error {
	src, err := ioutil.ReadAll(f)
	if err != nil {
		return err
	}

	file, err := parser.ParseFile(fset, f.Name(), src, parserMode)

	if err != nil {
		return err
	}

	if rewrite != nil {
		file = rewrite(file)
	}

	if *simplifyAST {
		simplify(file)
	}

	var buf bytes.Buffer
	_, err = (&printer.Config{printerMode, *tabWidth}).Fprint(&buf, fset, file)
	if err != nil {
		return err
	}
	res := buf.Bytes()

	if !bytes.Equal(src, res) {
		// formatting has changed
		if *list {
			fmt.Fprintln(os.Stdout, f.Name())
		}
		if *write {
			err = ioutil.WriteFile(f.Name(), res, 0)
			if err != nil {
				return err
			}
		}
	}

	if !*list && !*write {
		_, err = os.Stdout.Write(res)
	}

	return err
}

func processFileByName(filename string) os.Error {
	file, err := os.Open(filename, os.O_RDONLY, 0)
	if err != nil {
		return err
	}
	defer file.Close()
	return processFile(file)
}

type fileVisitor chan os.Error

func (v fileVisitor) VisitDir(path string, f *os.FileInfo) bool {
	return true
}

func (v fileVisitor) VisitFile(path string, f *os.FileInfo) {
	if isGoFile(f) {
		v <- nil // synchronize error handler
		if err := processFileByName(path); err != nil {
			v <- err
		}
	}
}

func walkDir(path string) {
	v := make(fileVisitor)
	go func() {
		filepath.Walk(path, v, v)
		close(v)
	}()
	for err := range v {
		if err != nil {
			report(err)
		}
	}
}

func main() {
	flag.Usage = usage
	flag.Parse()
	if *tabWidth < 0 {
		fmt.Fprintf(os.Stderr, "negative tabwidth %d\n", *tabWidth)
		os.Exit(2)
	}

	initParserMode()
	initPrinterMode()
	initRewrite()

	if flag.NArg() == 0 {
		if err := processFile(os.Stdin); err != nil {
			report(err)
		}
	}

	for i := 0; i < flag.NArg(); i++ {
		path := flag.Arg(i)
		switch dir, err := os.Stat(path); {
		case err != nil:
			report(err)
		case dir.IsRegular():
			if err := processFileByName(path); err != nil {
				report(err)
			}
		case dir.IsDirectory():
			walkDir(path)
		}
	}

	os.Exit(exitCode)
}
<|endoftext|>


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

package main

import (
	"bufio"
	"crypto/x509"
	"errors"
	"flag"
	"encoding/json"
        "github.com/sideshow/apns2"
	"github.com/sideshow/apns2/certificate"
	"log"
	"net"
	"os"
	"strings"
)

const Version = "2.0"

type command struct {
	name string
	args map[string]interface{}
}

type Payload struct {
	Aps Aps `json:"aps"`
}

type Aps struct {
	AccountID string `json:"account-id"`
}

func (cmd *command) getStringArg(name string) (string, bool) {
	arg, ok := cmd.args[name].(string)
	return arg, ok
}

func (cmd *command) getListArg(name string) ([]string, bool) {
	arg, ok := cmd.args[name].([]string)
	return arg, ok
}

func parseListValue(value string) ([]string, error) {
	list := []string{}
	values := strings.Split(value[1:len(value)-1], ",")
	for _, value := range values {
		stringValue, err := parseStringValue(value)
		if err != nil {
			return nil, err
		}
		list = append(list, stringValue)
	}
	return list, nil
}

func parseStringValue(value string) (string, error) {
	return value[1 : len(value)-1], nil // TODO Escaping!
}

func parseCommand(line string) (command, error) {
	cmd := command{args: make(map[string]interface{})}

	parts := strings.SplitN(line, " ", 2)
	if len(parts) != 2 {
		return cmd, errors.New("Failed to parse: no name found")
	}

	cmd.name = parts[0]

	for _, pair := range strings.Split(parts[1], "\t") {
		nameAndValue := strings.SplitN(pair, "=", 2)
		if len(nameAndValue) != 2 {
			return cmd, errors.New("Failed to parse: no name/value pair found")
		}

		switch {
		case strings.HasPrefix(nameAndValue[1], `"`) && strings.HasSuffix(nameAndValue[1], `"`):
			value, err := parseStringValue(nameAndValue[1])
			if err != nil {
				return cmd, err
			}
			cmd.args[nameAndValue[0]] = value
		case strings.HasPrefix(nameAndValue[1], "(") && strings.HasSuffix(nameAndValue[1], ")"):
			value, err := parseListValue(nameAndValue[1])
			if err != nil {
				return cmd, err
			}
			cmd.args[nameAndValue[0]] = value
		default:
			return cmd, errors.New("Failed to parse: invalid value in key/value pair")
		}
	}

	return cmd, nil
}

var debug = flag.Bool("debug", false, "enable debug logging")
var socket = flag.String("socket", "/var/run/xapsd/xapsd.sock", "path to the socket for Dovecot")
var database = flag.String("database", "/var/lib/xapsd/database.json", "path to the database file")
var certfile = flag.String("certificate", "/etc/xapsd/certificate.pem", "path to the pem/p12 file containing the key and certificate")

func topicFromCertificate(cert *x509.Certificate) (string, error) {
	if len(cert.Subject.Names) == 0 {
		return "", errors.New("Subject.Names is empty")
	}

	oidUid := []int{0, 9, 2342, 19200300, 100, 1, 1}
	if !cert.Subject.Names[0].Type.Equal(oidUid) {
		return "", errors.New("Did not find a Subject.Names[0] with type 0.9.2342.19200300.100.1.1")
	}

	return cert.Subject.Names[0].Value.(string), nil
}

func SetAccountID(accountid string) []byte {
        p := Payload{}
        p.Aps.AccountID = accountid
        aps, _ := json.Marshal(p)
        return aps
}

func main() {
	flag.Parse()

	if *debug {
		log.Println("[DEBUG] Opening database at", *database)
	}

	db, err := newDatabase(*database)
	if err != nil {
		log.Fatal("Cannot open database: ", *database, err.Error())
	}

	// Delete the socket is it already exists
	if _, err := os.Stat(*socket); err == nil {
		if err := os.Remove(*socket); err != nil {
			log.Fatal("Could not delete existing socket: ", *socket, err.Error())
		}
	}

	if *debug {
		log.Println("[DEBUG] Listening on UNIX socket at", *socket)
	}

	listener, err := net.Listen("unix", *socket)
	if err != nil {
		log.Fatal("Could not create socket: ", err.Error())
	}
	defer os.Remove(*socket)

	// TODO What is the proper way to limit Dovecot to this socket
	if err := os.Chmod(*socket, 0777); err != nil {
		log.Fatal("Could not chmod socket: ", err.Error())
	}

	if *debug {
		log.Println("[DEBUG] Parsing", *certfile, "to obtain APNS Topic")
	}

        cert, err := certificate.FromPemFile(*certfile, "")
        if err != nil {
		log.Println("PEM Certificate Loading Error: ", err)
		cert, err = certificate.FromP12File(*certfile, "")
		if err != nil {
			log.Fatal("P12 Certificate Loading Error: ", err)
		}
	}

	topic, err := topicFromCertificate(cert.Leaf)
	if err != nil {
		log.Fatal("Could not parse apns topic from certificate: ", err.Error())
	}

	if *debug {
		log.Println("[DEBUG] Topic is", topic)
	}

	if *debug {
		log.Println("[DEBUG] Creating APNS client to", apns2.HostProduction)
	}

	c := apns2.NewClient(cert).Production()

	log.Printf("Starting xapsd %s on %s", Version, *socket)

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Failed to accept connection: ", err.Error())
			os.Exit(1)
		}

		if *debug {
			log.Println("[DEBUG] Accepted a connection")
		}

		go handleRequest(conn, c, db, topic)
	}
}

func handleRequest(conn net.Conn, client *apns2.Client, db *Database, topic string) {
	defer conn.Close()

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		if *debug {
			log.Println("[DEBUG] Received request:", scanner.Text())
		}

		command, err := parseCommand(scanner.Text())
		if err != nil {
			log.Println("Reading froms socket: ", err)
		}

		switch command.name {
		case "REGISTER":
			handleRegister(conn, command, client, db, topic)
		case "NOTIFY":
			handleNotify(conn, command, client, db)
		default:
			writeError(conn, "Unknown command")
		}
	}

	if err := scanner.Err(); err != nil {
		log.Println("Reading froms socket: ", err)
	}
}

//
// Handle the REGISTER command. It looks as follows:
//
//  REGISTER aps-account-id="AAA" aps-device-token="BBB"
//     aps-subtopic="com.apple.mobilemail"
//     dovecot-username="stefan"
//     dovecot-mailboxes=("Inbox","Notes")
//
// The command returns the aps-topic, which is the common name of
// the certificate issued by OS X Server for email push
// notifications.
//

func handleRegister(conn net.Conn, cmd command, client *apns2.Client, db *Database, topic string) {
	// Make sure the subtopic is ok
	subtopic, ok := cmd.getStringArg("aps-subtopic")
	if !ok {
		writeError(conn, "Missing apis-subtopic argument")
		return
	}
	if subtopic != "com.apple.mobilemail" {
		writeError(conn, "Unknown aps-subtopic")
		return
	}

	// Make sure we got the required parameters
	accountId, ok := cmd.getStringArg("aps-account-id")
	if !ok {
		writeError(conn, "Missing aps-account-id argument")
		return
	}
	deviceToken, ok := cmd.getStringArg("aps-device-token")
	if !ok {
		writeError(conn, "Missing aps-device-token argument")
		return
	}
	username, ok := cmd.getStringArg("dovecot-username")
	if !ok {
		writeError(conn, "Missing dovecot-username argument")
		return
	}
	mailboxes, ok := cmd.getListArg("dovecot-mailboxes")
	if !ok {
		writeError(conn, "Missing dovecot-mailboxes argument")
		return
	}

	// Register this email/account-id/device-token combination
	err := db.addRegistration(username, accountId, deviceToken, mailboxes)
	if err != nil {
		writeError(conn, "Failed to register client: "+err.Error())
		return
	}

	writeSuccess(conn, topic)
}

//
// Handle the NOTIFY command. It looks as follows:
//
//  NOTIFY dovecot-username="stefan" dovecot-mailbox="Inbox"
//
// See if the the username has devices registered. If it has, loop
// over them to find the ones that are interested in the named
// mailbox and send those a push notificiation.
//
// The push notification looks like this:
//
//  { "aps": { "account-id": aps-account-id } }
//

func handleNotify(conn net.Conn, cmd command, client *apns2.Client, db *Database) {
	// Make sure we got the required arguments
	username, ok := cmd.getStringArg("dovecot-username")
	if !ok {
		writeError(conn, "Missing dovecot-username argument")
		return
	}
	mailbox, ok := cmd.getStringArg("dovecot-mailbox")
	if !ok {
		writeError(conn, "Missing dovecot-mailbox argument")
		return
	}

	// Find all the devices registered for this mailbox event
	registrations, err := db.findRegistrations(username, mailbox)
	if err != nil {
		writeError(conn, "Cannot lookup registrations: "+err.Error())
		return
	}

	// Send a notification to all registered devices. We ignore failures
	// because there is not a lot we can do.
	for _, registration := range registrations {
		if *debug {
			log.Println("[DEBUG] Sending notification to", registration.AccountId, "/", registration.DeviceToken)
		}
		sendNotification(registration, client)
	}

	writeSuccess(conn, "")
}

func sendNotification(reg Registration, client *apns2.Client) {
	notification := &apns2.Notification{}
	notification.Payload = SetAccountID(reg.AccountId)
	notification.DeviceToken = reg.DeviceToken
	res, err := client.Push(notification)

	if err != nil {
		if *debug {
			log.Println("[DEBUG] Error: ", err)
		}
        }
	if *debug {
		log.Printf("[DEBUG] %v %v %v\n", res.StatusCode, res.ApnsID, res.Reason)
        }
}

func writeError(conn net.Conn, msg string) {
	if *debug {
		log.Println("[DEBUG] Returning failure:", msg)
	}
	conn.Write([]byte("ERROR" + " " + msg + "\n"))
}

func writeSuccess(conn net.Conn, msg string) {
	if *debug {
		log.Println("[DEBUG] Returning success:", msg)
	}
	conn.Write([]byte("OK" + " " + msg + "\n"))
}
<commit_msg>Support SQL database Add support external configuration file to use with database Add support for correctly shutting down server if it is killed with control + c and kill signal Add support deleting registrations<commit_after>//
// The MIT License (MIT)
//
// Copyright (c) 2015 Stefan Arentz <stefan@arentz.ca>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

package main

import (
	"bufio"
	"crypto/x509"
	"errors"
	"flag"
	"encoding/json"
        "github.com/sideshow/apns2"
	"github.com/sideshow/apns2/certificate"
	"github.com/jinzhu/configor"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
	"strings"
)

const Version = "2.0"

type command struct {
	name string
	args map[string]interface{}
}

type Payload struct {
	Aps Aps `json:"aps"`
}

type Aps struct {
	AccountID string `json:"account-id"`
}

func (cmd *command) getStringArg(name string) (string, bool) {
	arg, ok := cmd.args[name].(string)
	return arg, ok
}

func (cmd *command) getListArg(name string) ([]string, bool) {
	arg, ok := cmd.args[name].([]string)
	return arg, ok
}

func parseListValue(value string) ([]string, error) {
	list := []string{}
	values := strings.Split(value[1:len(value)-1], ",")
	for _, value := range values {
		stringValue, err := parseStringValue(value)
		if err != nil {
			return nil, err
		}
		list = append(list, stringValue)
	}
	return list, nil
}

func parseStringValue(value string) (string, error) {
	return value[1 : len(value)-1], nil // TODO Escaping!
}

func parseCommand(line string) (command, error) {
	cmd := command{args: make(map[string]interface{})}

	parts := strings.SplitN(line, " ", 2)
	if len(parts) != 2 {
		return cmd, errors.New("Failed to parse: no name found")
	}

	cmd.name = parts[0]

	for _, pair := range strings.Split(parts[1], "\t") {
		nameAndValue := strings.SplitN(pair, "=", 2)
		if len(nameAndValue) != 2 {
			return cmd, errors.New("Failed to parse: no name/value pair found")
		}

		switch {
		case strings.HasPrefix(nameAndValue[1], `"`) && strings.HasSuffix(nameAndValue[1], `"`):
			value, err := parseStringValue(nameAndValue[1])
			if err != nil {
				return cmd, err
			}
			cmd.args[nameAndValue[0]] = value
		case strings.HasPrefix(nameAndValue[1], "(") && strings.HasSuffix(nameAndValue[1], ")"):
			value, err := parseListValue(nameAndValue[1])
			if err != nil {
				return cmd, err
			}
			cmd.args[nameAndValue[0]] = value
		default:
			return cmd, errors.New("Failed to parse: invalid value in key/value pair")
		}
	}

	return cmd, nil
}

var debug = flag.Bool("debug", false, "enable debug logging")

func topicFromCertificate(cert *x509.Certificate) (string, error) {
	if len(cert.Subject.Names) == 0 {
		return "", errors.New("Subject.Names is empty")
	}

	oidUid := []int{0, 9, 2342, 19200300, 100, 1, 1}
	if !cert.Subject.Names[0].Type.Equal(oidUid) {
		return "", errors.New("Did not find a Subject.Names[0] with type 0.9.2342.19200300.100.1.1")
	}

	return cert.Subject.Names[0].Value.(string), nil
}

func SetAccountID(accountid string) []byte {
        p := Payload{}
        p.Aps.AccountID = accountid
        aps, _ := json.Marshal(p)
        return aps
}

type SQLQueries struct {
    Sql	string
}

var Config = struct {
	Certificate string `default:"/etc/xapsd/certificate.pem"`
	Socket      string `default:"/var/run/xapsd/xapsd.sock"`

	DB struct {
		Host	 string
		Port	 uint16 `default:"3306"`
		Socket	 string
		Name	 string
		User	 string `default:"root"`
		Password string `required:"true"`
		Options  string `default:"timeout=5s&collation=utf8mb4_unicode_ci"`
		Queries map[string]SQLQueries
	}
}{}

func main() {
	config := flag.String("config", "/etc/xapsd.toml", "path to configuration file")
	socket := flag.String("socket", "", "path to the socket for Dovecot")
	certfile := flag.String("certificate", "", "path to the pem/p12 file containing the key and certificate")
	flag.Parse()

	configor.Load(&Config, *config)

	if *certfile != "" {
		Config.Certificate = *certfile
	}

        if *socket != "" {
                Config.Socket = *socket
        }

	db, err := connectDatabase()
	if err != nil {
		log.Fatal(err)
	}
	defer db.conn.Close()

	for name, sql := range db.queries {
		defer sql.Close()
        }

	// Delete the socket if it already exists
	if _, err := os.Stat(Config.Socket); err == nil {
		if err := os.Remove(Config.Socket); err != nil {
			log.Fatal("Could not delete existing socket: ", Config.Socket, err.Error())
		}
	}

	if *debug {
		log.Println("[DEBUG] Listening on UNIX socket at", Config.Socket)
	}

	listener, err := net.Listen("unix", Config.Socket)
	if err != nil {
		log.Fatal("Could not create socket: ", err.Error())
	}
	defer listener.Close()
	defer os.Remove(Config.Socket)

	// TODO What is the proper way to limit Dovecot to this socket
	if err := os.Chmod(Config.Socket, 0777); err != nil {
		log.Fatal("Could not chmod socket: ", err.Error())
	}

	if *debug {
		log.Println("[DEBUG] Parsing", Config.Certificate, "to obtain APNS Topic")
	}

        cert, err := certificate.FromPemFile(Config.Certificate, "")
        if err != nil {
		log.Println("PEM Certificate Loading Error: ", err)
		cert, err = certificate.FromP12File(Config.Certificate, "")
		if err != nil {
			log.Fatal("P12 Certificate Loading Error: ", err)
		}
	}

	topic, err := topicFromCertificate(cert.Leaf)
	if err != nil {
		log.Fatal("Could not parse apns topic from certificate: ", err.Error())
	}

	if *debug {
		log.Println("[DEBUG] Topic is", topic)
	}

	if *debug {
		log.Println("[DEBUG] Creating APNS client to", apns2.HostProduction)
	}

	c := apns2.NewClient(cert).Production()

	signalChannel := make(chan os.Signal, 2)
	quit := make(chan bool)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-signalChannel
		close(quit)
		listener.Close()
	}()

	log.Printf("Starting xapsd %s on %s", Version, Config.Socket)

	Accept:
	for {
		conn, err := listener.Accept()
		if err != nil {
			select {
			case <-quit:
				log.Printf("Shutting Down xapsd %s", Version)
				break Accept
			default:
				log.Println("Failed to accept connection: ", err.Error())
				os.Exit(1)
			}
			continue
		}

		if *debug {
			log.Println("[DEBUG] Accepted a connection")
		}

		go handleRequest(conn, c, db, topic)

	}
}

func handleRequest(conn net.Conn, client *apns2.Client, db *Database, topic string) {
	defer conn.Close()

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		if *debug {
			log.Println("[DEBUG] Received request:", scanner.Text())
		}

		command, err := parseCommand(scanner.Text())
		if err != nil {
			log.Println("Reading from socket: ", err)
		}

		switch command.name {
		case "REGISTER":
			handleRegister(conn, command, client, db, topic)
		case "NOTIFY":
			handleNotify(conn, command, client, db)
		default:
			writeError(conn, "Unknown command")
		}
	}

	if err := scanner.Err(); err != nil {
		log.Println("Reading from socket: ", err)
	}
}

//
// Handle the REGISTER command. It looks as follows:
//
//  REGISTER aps-account-id="AAA" aps-device-token="BBB"
//     aps-subtopic="com.apple.mobilemail"
//     dovecot-username="stefan"
//     dovecot-mailboxes=("Inbox","Notes")
//
// The command returns the aps-topic, which is the common name of
// the certificate issued by OS X Server for email push
// notifications.
//

func handleRegister(conn net.Conn, cmd command, client *apns2.Client, db *Database, topic string) {
	// Make sure the subtopic is ok
	subtopic, ok := cmd.getStringArg("aps-subtopic")
	if !ok {
		writeError(conn, "Missing apis-subtopic argument")
		return
	}
	if subtopic != "com.apple.mobilemail" {
		writeError(conn, "Unknown aps-subtopic")
		return
	}

	// Make sure we got the required parameters
	accountId, ok := cmd.getStringArg("aps-account-id")
	if !ok {
		writeError(conn, "Missing aps-account-id argument")
		return
	}
	deviceToken, ok := cmd.getStringArg("aps-device-token")
	if !ok {
		writeError(conn, "Missing aps-device-token argument")
		return
	}
	username, ok := cmd.getStringArg("dovecot-username")
	if !ok {
		writeError(conn, "Missing dovecot-username argument")
		return
	}
	mailboxes, ok := cmd.getListArg("dovecot-mailboxes")
	if !ok {
		writeError(conn, "Missing dovecot-mailboxes argument")
		return
	}

	// Register this email/account-id/device-token combination
	err := db.addRegistration(username, accountId, deviceToken, mailboxes)
	if err != nil {
		writeError(conn, "Failed to register client: "+err.Error())
		return
	}

	writeSuccess(conn, topic)
}

//
// Handle the NOTIFY command. It looks as follows:
//
//  NOTIFY dovecot-username="stefan" dovecot-mailbox="Inbox"
//
// See if the the username has devices registered. If it has, loop
// over them to find the ones that are interested in the named
// mailbox and send those a push notificiation.
//
// The push notification looks like this:
//
//  { "aps": { "account-id": aps-account-id } }
//

func handleNotify(conn net.Conn, cmd command, client *apns2.Client, db *Database) {
	// Make sure we got the required arguments
	username, ok := cmd.getStringArg("dovecot-username")
	if !ok {
		writeError(conn, "Missing dovecot-username argument")
		return
	}
	mailbox, ok := cmd.getStringArg("dovecot-mailbox")
	if !ok {
		writeError(conn, "Missing dovecot-mailbox argument")
		return
	}

	// Find all the devices registered for this mailbox event
	registrations, err := db.findRegistrations(username, mailbox)
	if err != nil {
		writeError(conn, "Cannot lookup registrations: "+err.Error())
		return
	}

	// Send a notification to all registered devices. We ignore failures
	// because there is not a lot we can do. We do delete registrations
	// if Apple servers respond with 410 error.
	for _, registration := range registrations {
		if *debug {
			log.Println("[DEBUG] Sending notification to", registration.AccountId, "/", registration.DeviceToken)
		}
		res := sendNotification(registration, client)
		if res.StatusCode == 410 {
			if *debug {
				log.Printf("[DEBUG] Device %v (DB: %v) is no longer registered. APN-Status: %v (%v)\n", registration.AccountId, registration.DbId, res.StatusCode, res.Reason)
			}
			db.deleteRegistration(registration)
		}
	}

	writeSuccess(conn, "")
}

func sendNotification(reg Registration, client *apns2.Client) (*apns2.Response) {
	notification := &apns2.Notification{}
	notification.Payload = SetAccountID(reg.AccountId)
	notification.DeviceToken = reg.DeviceToken
	res, err := client.Push(notification)

	if err != nil {
		log.Println("Sending Notification failed: ", err)
		return nil
        }

	if *debug {
		log.Printf("[DEBUG] %v %v %v\n", res.StatusCode, res.ApnsID, res.Reason)
        }
	return res
}

func writeError(conn net.Conn, msg string) {
	if *debug {
		log.Println("[DEBUG] Returning failure:", msg)
	}
	conn.Write([]byte("ERROR" + " " + msg + "\n"))
}

func writeSuccess(conn net.Conn, msg string) {
	if *debug {
		log.Println("[DEBUG] Returning success:", msg)
	}
	conn.Write([]byte("OK" + " " + msg + "\n"))
}
<|endoftext|>


<commit_before>package controllers

import (
	"encoding/json"
	"strconv"

	"github.com/astaxie/beego"
	"github.com/cosminrentea/franz/models"
)

type MessageController struct {
	beego.Controller
}

// Example:
//
//   req: GET /message/
//   res: 200 {"Messages": [
//          {"ID": 1, "Title": "Learn Go", "Done": false},
//          {"ID": 2, "Title": "Buy bread", "Done": true}
//        ]}
func (this *MessageController) ListMessages() {
	res := struct{ Messages []*models.Message }{models.DefaultMessageList.All()}
	this.Data["json"] = res
	this.ServeJSON()
}

// Examples:
//
//   req: POST /message/ {"Title": ""}
//   res: 400 empty title
//
//   req: POST /message/ {"Title": "Buy bread"}
//   res: 200
func (this *MessageController) NewMessage() {
	req := struct{ Title string }{}
	if err := json.Unmarshal(this.Ctx.Input.RequestBody, &req); err != nil {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte("empty title"))
		return
	}
	t, err := models.NewMessage(req.Title)
	if err != nil {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte(err.Error()))
		return
	}
	err = models.DefaultMessageList.Send(t)
	if err != nil {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte(err.Error()))
		return
	}
	models.DefaultMessageList.Save(t)
}

// Examples:
//
//   req: GET /message/1
//   res: 200 {"ID": 1, "Title": "Buy bread", "Done": true}
//
//   req: GET /message/42
//   res: 404 message not found
func (this *MessageController) GetMessage() {
	id := this.Ctx.Input.Param(":id")
	beego.Info("Message is ", id)
	intid, _ := strconv.ParseInt(id, 10, 64)
	t, ok := models.DefaultMessageList.Find(intid)
	beego.Info("Found", ok)
	if !ok {
		this.Ctx.Output.SetStatus(404)
		this.Ctx.Output.Body([]byte("message not found"))
		return
	}
	this.Data["json"] = t
	this.ServeJSON()
}

// Example:
//
//   req: PUT /message/1 {"ID": 1, "Title": "Learn Go", "Done": true}
//   res: 200
//
//   req: PUT /message/2 {"ID": 2, "Title": "Learn Go", "Done": true}
//   res: 400 inconsistent message IDs
func (this *MessageController) UpdateMessage() {
	id := this.Ctx.Input.Param(":id")
	beego.Info("Message is ", id)
	intid, _ := strconv.ParseInt(id, 10, 64)
	var t models.Message
	if err := json.Unmarshal(this.Ctx.Input.RequestBody, &t); err != nil {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte(err.Error()))
		return
	}
	if t.ID != intid {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte("inconsistent message IDs"))
		return
	}
	if _, ok := models.DefaultMessageList.Find(intid); !ok {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte("message not found"))
		return
	}
	models.DefaultMessageList.Save(&t)
}
<commit_msg>refactoring<commit_after>package controllers

import (
	"encoding/json"
	"strconv"

	"github.com/astaxie/beego"
	"github.com/cosminrentea/franz/models"
)

type MessageController struct {
	beego.Controller
}

// Example:
//
//   req: GET /message/
//   res: 200 {"Messages": [
//          {"ID": 1, "Title": "Learn Go", "Done": false},
//          {"ID": 2, "Title": "Buy bread", "Done": true}
//        ]}
func (this *MessageController) ListMessages() {
	res := struct{ Messages []*models.Message }{models.DefaultMessageList.All()}
	this.Data["json"] = res
	this.ServeJSON()
}

// Examples:
//
//   req: POST /message/ {"Title": ""}
//   res: 400 empty title
//
//   req: POST /message/ {"Title": "Buy bread"}
//   res: 200
func (this *MessageController) NewMessage() {
	req := struct{ Title string }{}
	if err := json.Unmarshal(this.Ctx.Input.RequestBody, &req); err != nil {
		this.outputError(400, "empty title")
		return
	}
	t, err := models.NewMessage(req.Title)
	if err != nil {
		this.outputError(400, err.Error())
		return
	}
	err = models.DefaultMessageList.Send(t)
	if err != nil {
		this.outputError(400, err.Error())
		return
	}
	models.DefaultMessageList.Save(t)
}

func (this *MessageController) outputError(status int, msg string) {
	this.Ctx.Output.SetStatus(status)
	this.Ctx.Output.Body([]byte(msg))
}

// Examples:
//
//   req: GET /message/1
//   res: 200 {"ID": 1, "Title": "Buy bread", "Done": true}
//
//   req: GET /message/42
//   res: 404 message not found
func (this *MessageController) GetMessage() {
	id := this.Ctx.Input.Param(":id")
	beego.Info("Message is ", id)
	intid, _ := strconv.ParseInt(id, 10, 64)
	t, ok := models.DefaultMessageList.Find(intid)
	beego.Info("Found", ok)
	if !ok {
		this.Ctx.Output.SetStatus(404)
		this.Ctx.Output.Body([]byte("message not found"))
		return
	}
	this.Data["json"] = t
	this.ServeJSON()
}

// Example:
//
//   req: PUT /message/1 {"ID": 1, "Title": "Learn Go", "Done": true}
//   res: 200
//
//   req: PUT /message/2 {"ID": 2, "Title": "Learn Go", "Done": true}
//   res: 400 inconsistent message IDs
func (this *MessageController) UpdateMessage() {
	id := this.Ctx.Input.Param(":id")
	beego.Info("Message is ", id)
	intid, _ := strconv.ParseInt(id, 10, 64)
	var t models.Message
	if err := json.Unmarshal(this.Ctx.Input.RequestBody, &t); err != nil {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte(err.Error()))
		return
	}
	if t.ID != intid {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte("inconsistent message IDs"))
		return
	}
	if _, ok := models.DefaultMessageList.Find(intid); !ok {
		this.Ctx.Output.SetStatus(400)
		this.Ctx.Output.Body([]byte("message not found"))
		return
	}
	models.DefaultMessageList.Save(&t)
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package router

import (
	"encoding"
	"encoding/binary"
	"fmt"
	"sync"
	"time"

	dbutil "github.com/TheThingsNetwork/ttn/core/storage"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/readwriter"
)

var dbBrokers = []byte("brokers")

// BrkStorage gives a facade to manipulate the router's brokers
type BrkStorage interface {
	read(devAddr []byte) ([]brkEntry, error)
	create(entry brkEntry) error
	//remove(entry brkEntry) error
	done() error
}

type brkEntry struct {
	DevAddr     []byte
	BrokerIndex uint16
	until       time.Time
}

type brkStorage struct {
	sync.Mutex
	ExpiryDelay time.Duration
	db          dbutil.Interface
}

// NewBrkStorage creates a new internal storage for the router
func NewBrkStorage(name string, delay time.Duration) (BrkStorage, error) {
	itf, err := dbutil.New(name)
	if err != nil {
		return nil, errors.New(errors.Operational, err)
	}
	return &brkStorage{db: itf, ExpiryDelay: delay}, nil
}

// read implements the router.BrkStorage interface
func (s *brkStorage) read(devAddr []byte) ([]brkEntry, error) {
	s.Lock()
	defer s.Unlock()
	itf, err := s.db.Read(devAddr, &brkEntry{}, dbBrokers)
	if err != nil {
		return nil, err
	}
	entries := itf.([]brkEntry)

	if s.ExpiryDelay != 0 {
		// Get rid of expired entries
		var newEntries []encoding.BinaryMarshaler
		var filtered []brkEntry
		for _, e := range entries {
			if e.until.After(time.Now()) {
				newEntry := new(brkEntry)
				*newEntry = e
				newEntries = append(newEntries, newEntry)
				filtered = append(filtered, e)
			}
		}
		// Replace filtered entries
		if err := s.db.Update(devAddr, newEntries, dbBrokers); err != nil {
			return nil, errors.New(errors.Operational, err)
		}
		entries = filtered
	}

	if len(entries) == 0 {
		return nil, errors.New(errors.NotFound, fmt.Sprintf("No entry for: %v", devAddr))
	}
	return entries, nil
}

// create implements the router.BrkStorage interface
func (s *brkStorage) create(entry brkEntry) error {
	s.Lock()
	defer s.Unlock()
	entry.until = time.Now().Add(s.ExpiryDelay)
	return s.db.Append(entry.DevAddr, []encoding.BinaryMarshaler{entry}, dbBrokers)
}

// done implements the router.BrkStorage interface
func (s *brkStorage) done() error {
	return s.db.Close()
}

// MarshalBinary implements the encoding.BinaryMarshaler
func (e brkEntry) MarshalBinary() ([]byte, error) {
	data, err := e.until.MarshalBinary()
	if err != nil {
		return nil, errors.New(errors.Structural, err)
	}
	rw := readwriter.New(nil)
	rw.Write(e.BrokerIndex)
	rw.Write(e.DevAddr)
	rw.Write(data)
	return rw.Bytes()
}

// UnmarshalBinary implements the encoding.BinaryUnmarshaler
func (e *brkEntry) UnmarshalBinary(data []byte) error {
	rw := readwriter.New(data)
	rw.Read(func(data []byte) { e.BrokerIndex = binary.BigEndian.Uint16(data) })
	rw.Read(func(data []byte) {
		e.DevAddr = make([]byte, len(data))
		copy(e.DevAddr, data)
	})
	rw.TryRead(func(data []byte) error { return e.until.UnmarshalBinary(data) })
	return rw.Err()
}
<commit_msg>Fix concurrent access to the router storage, causing a same broker to be stored twiced in some cases<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package router

import (
	"encoding"
	"encoding/binary"
	"fmt"
	"sync"
	"time"

	dbutil "github.com/TheThingsNetwork/ttn/core/storage"
	"github.com/TheThingsNetwork/ttn/utils/errors"
	"github.com/TheThingsNetwork/ttn/utils/readwriter"
)

var dbBrokers = []byte("brokers")

// BrkStorage gives a facade to manipulate the router's brokers
type BrkStorage interface {
	read(devAddr []byte) ([]brkEntry, error)
	create(entry brkEntry) error
	//remove(entry brkEntry) error
	done() error
}

type brkEntry struct {
	DevAddr     []byte
	BrokerIndex uint16
	until       time.Time
}

type brkStorage struct {
	sync.Mutex
	ExpiryDelay time.Duration
	db          dbutil.Interface
}

// NewBrkStorage creates a new internal storage for the router
func NewBrkStorage(name string, delay time.Duration) (BrkStorage, error) {
	itf, err := dbutil.New(name)
	if err != nil {
		return nil, errors.New(errors.Operational, err)
	}
	return &brkStorage{db: itf, ExpiryDelay: delay}, nil
}

// read implements the router.BrkStorage interface
func (s *brkStorage) read(devAddr []byte) ([]brkEntry, error) {
	return s._read(devAddr, true)
}

// _read implements the router.BrkStorage interface logic but gives control over mutex to the caller
func (s *brkStorage) _read(devAddr []byte, shouldLock bool) ([]brkEntry, error) {
	if shouldLock {
		s.Lock()
		defer s.Unlock()
	}
	itf, err := s.db.Read(devAddr, &brkEntry{}, dbBrokers)
	if err != nil {
		return nil, err
	}
	entries := itf.([]brkEntry)

	if s.ExpiryDelay != 0 {
		// Get rid of expired entries
		var newEntries []encoding.BinaryMarshaler
		var filtered []brkEntry
		for _, e := range entries {
			if e.until.After(time.Now()) {
				newEntry := new(brkEntry)
				*newEntry = e
				newEntries = append(newEntries, newEntry)
				filtered = append(filtered, e)
			}
		}
		// Replace filtered entries
		if err := s.db.Update(devAddr, newEntries, dbBrokers); err != nil {
			return nil, errors.New(errors.Operational, err)
		}
		entries = filtered
	}

	if len(entries) == 0 {
		return nil, errors.New(errors.NotFound, fmt.Sprintf("No entry for: %v", devAddr))
	}
	return entries, nil
}

// create implements the router.BrkStorage interface
func (s *brkStorage) create(entry brkEntry) error {
	s.Lock()
	defer s.Unlock()

	entries, err := s._read(entry.DevAddr, false)
	if err != nil && err.(errors.Failure).Nature != errors.NotFound {
		return err
	}

	var updates []encoding.BinaryMarshaler
	until, found := time.Now().Add(s.ExpiryDelay), false
	for i, e := range entries {
		if entry.BrokerIndex == e.BrokerIndex {
			// Entry already there, just update the TTL
			entries[i].until = until
			found = true
		}
		updates = append(updates, e)
	}

	if found { // The entry was already existing
		return s.db.Update(entry.DevAddr, updates, dbBrokers)
	}

	// Otherwise, we just happend it
	entry.until = until
	return s.db.Append(entry.DevAddr, []encoding.BinaryMarshaler{entry}, dbBrokers)
}

// done implements the router.BrkStorage interface
func (s *brkStorage) done() error {
	return s.db.Close()
}

// MarshalBinary implements the encoding.BinaryMarshaler
func (e brkEntry) MarshalBinary() ([]byte, error) {
	data, err := e.until.MarshalBinary()
	if err != nil {
		return nil, errors.New(errors.Structural, err)
	}
	rw := readwriter.New(nil)
	rw.Write(e.BrokerIndex)
	rw.Write(e.DevAddr)
	rw.Write(data)
	return rw.Bytes()
}

// UnmarshalBinary implements the encoding.BinaryUnmarshaler
func (e *brkEntry) UnmarshalBinary(data []byte) error {
	rw := readwriter.New(data)
	rw.Read(func(data []byte) { e.BrokerIndex = binary.BigEndian.Uint16(data) })
	rw.Read(func(data []byte) {
		e.DevAddr = make([]byte, len(data))
		copy(e.DevAddr, data)
	})
	rw.TryRead(func(data []byte) error { return e.until.UnmarshalBinary(data) })
	return rw.Err()
}
<|endoftext|>


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

package elasticsearch

import (
	"encoding/json"
	"fmt"

	"github.com/venicegeo/pz-gocommon"

	"gopkg.in/olivere/elastic.v2"
)

type Index struct {
	esClient *Client
	lib      *elastic.Client
	index    string
}

func NewIndex(es *Client, index string) *Index {
	esi := &Index{
		esClient: es,
		lib:      es.lib,
		index:    index + es.indexSuffix,
	}
	return esi
}

func (esi *Index) IndexName() string {
	return esi.index
}

func (esi *Index) IndexExists() bool {
	ok, err := esi.lib.IndexExists(esi.index).Do()
	if err != nil {
		return false
	}
	return ok
}

func (esi *Index) TypeExists(typ string) bool {
	ok := esi.IndexExists()
	if !ok {
		return false
	}

	ok, err := esi.lib.TypeExists().Index(esi.index).Type(typ).Do()
	if err != nil {
		return false
	}
	return ok
}

func (esi *Index) ItemExists(typ string, id string) bool {
	ok := esi.TypeExists(typ)
	if !ok {
		return false
	}

	ok, err := esi.lib.Exists().Index(esi.index).Type(typ).Id(id).Do()
	if err != nil {
		return false
	}
	return ok
}

// if index already exists, does nothing
func (esi *Index) Create() error {

	ok := esi.IndexExists()
	if ok {
		return fmt.Errorf("Index %s already exists", esi.index)
	}

	createIndex, err := esi.lib.CreateIndex(esi.index).Do()
	if err != nil {
		return err
	}

	if !createIndex.Acknowledged {
		return fmt.Errorf("Elasticsearch: create index not acknowledged!")
	}

	return nil
}

// if index doesn't already exist, does nothing
func (esi *Index) Close() error {

	// TODO: the caller should enforce this instead
	ok := esi.IndexExists()
	if !ok {
		return fmt.Errorf("Index %s does not already exist", esi.index)
	}

	_, err := esi.lib.CloseIndex(esi.index).Do()
	if err != nil {
		return err
	}

	return nil
}

// if index doesn't already exist, does nothing
func (esi *Index) Delete() error {

	ok := esi.IndexExists()
	if !ok {
		return fmt.Errorf("Index %s does not exist", esi.index)
	}

	deleteIndex, err := esi.lib.DeleteIndex(esi.index).Do()
	if err != nil {
		return err
	}

	// TODO: is this check needed? should it also be on Create(), etc?
	if !deleteIndex.Acknowledged {
		return fmt.Errorf("Elasticsearch: delete index not acknowledged!")
	}
	return nil
}

// TODO: how often should we do this?
func (esi *Index) Flush() error {
	// TODO: the caller should enforce this instead
	ok := esi.IndexExists()
	if !ok {
		return fmt.Errorf("Index %s does not exist", esi.index)
	}

	_, err := esi.lib.Flush().Index(esi.index).Do()
	if err != nil {
		return err
	}

	return nil
}

func (esi *Index) PostData(typ string, id string, obj interface{}) (*elastic.IndexResult, error) {
	/*ok := esi.IndexExists()
	if !ok {
		log.Printf("Index %s does not exist", esi.index)
		return nil, errors.New(fmt.Sprintf("Index %s does not exist", esi.index))
	}
	ok = esi.TypeExists(typ)
	if !ok {
		log.Printf("Index %s or type %s does not exist", esi.index, typ)
		return nil, errors.New(fmt.Sprintf("Index %s or type %s does not exist", esi.index, typ))
	}*/

	indexResult, err := esi.lib.Index().
		Index(esi.index).
		Type(typ).
		Id(id).
		BodyJson(obj).
		Do()

	return indexResult, err
}

func (esi *Index) GetByID(typ string, id string) (*elastic.GetResult, error) {
	// TODO: the caller should enforce this instead (here and elsewhere)
	ok := esi.ItemExists(typ, id)
	if !ok {
		return nil, fmt.Errorf("Item %s in index %s and type %s does not exist", id, esi.index, typ)
	}

	getResult, err := esi.lib.Get().Index(esi.index).Type(typ).Id(id).Do()
	return getResult, err
}

func (esi *Index) DeleteByID(typ string, id string) (*elastic.DeleteResult, error) {
	ok := esi.ItemExists(typ, id)
	if !ok {
		return nil, fmt.Errorf("Item %s in index %s and type %s does not exist", id, esi.index, typ)
	}

	deleteResult, err := esi.lib.Delete().
		Index(esi.index).
		Type(typ).
		Id(id).
		Do()
	return deleteResult, err
}

func (esi *Index) FilterByMatchAll(typ string) (*elastic.SearchResult, error) {
	//q := elastic.NewBoolFilter()
	//q.Must(elastic.NewTermFilter("a", 1))

	// TODO: the caller should enforce this instead
	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	q := elastic.NewMatchAllFilter()
	searchResult, err := esi.lib.Search().
		Index(esi.index).
		Type(typ).
		Query(q).
		//Sort("id", true).
		Do()
	return searchResult, err
}

func (esi *Index) FilterByTermQuery(typ string, name string, value interface{}) (*elastic.SearchResult, error) {

	// TODO: the caller should enforce this instead
	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	termQuery := elastic.NewTermFilter(name, value)
	searchResult, err := esi.lib.Search().
		Index(esi.index).
		Type(typ).
		Query(&termQuery).
		//Sort("id", true).
		Do()
	return searchResult, err
}

func (esi *Index) SearchByJSON(typ string, jsn string) (*elastic.SearchResult, error) {

	// TODO: the caller should enforce this instead
	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	var obj interface{}
	err := json.Unmarshal([]byte(jsn), &obj)
	if err != nil {
		return nil, err
	}

	searchResult, err := esi.lib.Search().
		Index(esi.index).
		Type(typ).
		Source(obj).Do()

	return searchResult, err
}

func (esi *Index) SetMapping(typename string, jsn piazza.JsonString) error {

	ok := esi.IndexExists()
	if !ok {
		return fmt.Errorf("Index %s does not exist", esi.index)
	}

	putresp, err := esi.lib.PutMapping().Index(esi.index).Type(typename).BodyString(string(jsn)).Do()
	if err != nil {
		return fmt.Errorf("expected put mapping to succeed; got: %v", err)
	}
	if putresp == nil {
		return fmt.Errorf("expected put mapping response; got: %v", putresp)
	}
	if !putresp.Acknowledged {
		return fmt.Errorf("expected put mapping ack; got: %v", putresp.Acknowledged)
	}

	return nil
}

func (esi *Index) GetTypes() ([]string, error) {

	ok := esi.IndexExists()
	if !ok {
		return nil, fmt.Errorf("Index %s does not exist", esi.index)
	}

	getresp, err := esi.lib.IndexGet().Feature("_mappings").Index(esi.index).Do()
	if err != nil {
		return nil, err
	}

	typs := (*getresp[esi.index]).Mappings
	result := make([]string, len(typs))

	i := 0
	for k := range typs {
		result[i] = k
		i++
	}

	return result, nil
}

func (esi *Index) GetMapping(typ string) (interface{}, error) {

	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	getresp, err := esi.lib.GetMapping().Index(esi.index).Type(typ).Do()
	if err != nil {
		return nil, fmt.Errorf("expected get mapping to succeed; got: %v", err)
	}
	if getresp == nil {
		return nil, fmt.Errorf("expected get mapping response; got: %v", getresp)
	}
	props, ok := getresp[esi.index]
	if !ok {
		return nil, fmt.Errorf("expected JSON root to be of type map[string]interface{}; got: %s -- %#v", esi.index, getresp)
	}

	props2 := props.(map[string]interface{})

	return props2["mappings"], nil
}

func (esi *Index) AddPercolationQuery(id string, query piazza.JsonString) (*elastic.IndexResult, error) {

	ok := esi.IndexExists()
	if !ok {
		return nil, fmt.Errorf("Index %s does not exist", esi.index)
	}

	indexResponse, err := esi.lib.
		Index().
		Index(esi.index).
		Type(".percolator").
		Id(id).
		BodyString(string(query)).
		Do()
	if err != nil {
		return nil, err
	}

	return indexResponse, nil
}

func (esi *Index) DeletePercolationQuery(id string) (*elastic.DeleteResult, error) {
	typ := ".percolator"
	ok := esi.ItemExists(typ, id)
	if !ok {
		return nil, fmt.Errorf("Item %s in index %s and type %s does not exist", id, esi.index, typ)
	}

	deleteResult, err := esi.lib.Delete().
		Index(esi.index).
		Type(".percolator").
		Id(id).
		Do()
	if err != nil {
		return nil, err
	}

	return deleteResult, nil
}

func (esi *Index) AddPercolationDocument(typ string, doc interface{}) (*elastic.PercolateResponse, error) {
	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	percolateResponse, err := esi.lib.
		Percolate().
		Index(esi.index).
		Type(typ).
		Doc(doc).
		//Pretty(true).
		Do()
	if err != nil {
		return nil, err
	}

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

package elasticsearch

import (
	"encoding/json"
	"fmt"
	"log"

	"github.com/venicegeo/pz-gocommon"

	"gopkg.in/olivere/elastic.v2"
)

type Index struct {
	esClient *Client
	lib      *elastic.Client
	index    string
}

func NewIndex(es *Client, index string) *Index {
	esi := &Index{
		esClient: es,
		lib:      es.lib,
		index:    index + es.indexSuffix,
	}
	return esi
}

func (esi *Index) IndexName() string {
	return esi.index
}

func (esi *Index) IndexExists() bool {
	ok, err := esi.lib.IndexExists(esi.index).Do()
	if err != nil {
		return false
	}
	return ok
}

func (esi *Index) TypeExists(typ string) bool {
	ok := esi.IndexExists()
	if !ok {
		return false
	}

	ok, err := esi.lib.TypeExists().Index(esi.index).Type(typ).Do()
	if err != nil {
		return false
	}
	return ok
}

func (esi *Index) ItemExists(typ string, id string) bool {
	ok := esi.TypeExists(typ)
	if !ok {
		return false
	}

	ok, err := esi.lib.Exists().Index(esi.index).Type(typ).Id(id).Do()
	if err != nil {
		return false
	}
	return ok
}

// if index already exists, does nothing
func (esi *Index) Create() error {

	ok := esi.IndexExists()
	if ok {
		return fmt.Errorf("Index %s already exists", esi.index)
	}

	createIndex, err := esi.lib.CreateIndex(esi.index).Do()
	if err != nil {
		return err
	}

	if !createIndex.Acknowledged {
		return fmt.Errorf("Elasticsearch: create index not acknowledged!")
	}

	return nil
}

// if index doesn't already exist, does nothing
func (esi *Index) Close() error {

	// TODO: the caller should enforce this instead
	ok := esi.IndexExists()
	if !ok {
		return fmt.Errorf("Index %s does not already exist", esi.index)
	}

	_, err := esi.lib.CloseIndex(esi.index).Do()
	if err != nil {
		return err
	}

	return nil
}

// if index doesn't already exist, does nothing
func (esi *Index) Delete() error {

	ok := esi.IndexExists()
	if !ok {
		return fmt.Errorf("Index %s does not exist", esi.index)
	}

	deleteIndex, err := esi.lib.DeleteIndex(esi.index).Do()
	if err != nil {
		return err
	}

	// TODO: is this check needed? should it also be on Create(), etc?
	if !deleteIndex.Acknowledged {
		return fmt.Errorf("Elasticsearch: delete index not acknowledged!")
	}
	return nil
}

// TODO: how often should we do this?
func (esi *Index) Flush() error {
	// TODO: the caller should enforce this instead
	ok := esi.IndexExists()
	if !ok {
		return fmt.Errorf("Index %s does not exist", esi.index)
	}

	_, err := esi.lib.Flush().Index(esi.index).Do()
	if err != nil {
		return err
	}

	return nil
}

func (esi *Index) PostData(typ string, id string, obj interface{}) (*elastic.IndexResult, error) {
	/*ok := esi.IndexExists()
	if !ok {
		log.Printf("Index %s does not exist", esi.index)
		return nil, errors.New(fmt.Sprintf("Index %s does not exist", esi.index))
	}
	ok = esi.TypeExists(typ)
	if !ok {
		log.Printf("Index %s or type %s does not exist", esi.index, typ)
		return nil, errors.New(fmt.Sprintf("Index %s or type %s does not exist", esi.index, typ))
	}*/

	indexResult, err := esi.lib.Index().
		Index(esi.index).
		Type(typ).
		Id(id).
		BodyJson(obj).
		Do()

	return indexResult, err
}

func (esi *Index) GetByID(typ string, id string) (*elastic.GetResult, error) {
	// TODO: the caller should enforce this instead (here and elsewhere)
	ok := esi.ItemExists(typ, id)
	if !ok {
		return nil, fmt.Errorf("Item %s in index %s and type %s does not exist", id, esi.index, typ)
	}

	svc := esi.lib.Get().Index(esi.index).Type(typ).Id(id)
	log.Printf("Index.GetByID: %s", svc.String())
	getResult, err := svc.Do()
	if err != nil {
		log.Printf("Index.GetByID failed: %s", err)
		return nil, err
	}

	return getResult, nil
}

func (esi *Index) DeleteByID(typ string, id string) (*elastic.DeleteResult, error) {
	ok := esi.ItemExists(typ, id)
	if !ok {
		return nil, fmt.Errorf("Item %s in index %s and type %s does not exist", id, esi.index, typ)
	}

	deleteResult, err := esi.lib.Delete().
		Index(esi.index).
		Type(typ).
		Id(id).
		Do()
	return deleteResult, err
}

func (esi *Index) FilterByMatchAll(typ string) (*elastic.SearchResult, error) {
	//q := elastic.NewBoolFilter()
	//q.Must(elastic.NewTermFilter("a", 1))

	// TODO: the caller should enforce this instead
	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	q := elastic.NewMatchAllFilter()
	searchResult, err := esi.lib.Search().
		Index(esi.index).
		Type(typ).
		Query(q).
		//Sort("id", true).
		Do()
	return searchResult, err
}

func (esi *Index) FilterByTermQuery(typ string, name string, value interface{}) (*elastic.SearchResult, error) {

	// TODO: the caller should enforce this instead
	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	termQuery := elastic.NewTermFilter(name, value)
	searchResult, err := esi.lib.Search().
		Index(esi.index).
		Type(typ).
		Query(&termQuery).
		//Sort("id", true).
		Do()
	return searchResult, err
}

func (esi *Index) SearchByJSON(typ string, jsn string) (*elastic.SearchResult, error) {

	// TODO: the caller should enforce this instead
	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	var obj interface{}
	err := json.Unmarshal([]byte(jsn), &obj)
	if err != nil {
		return nil, err
	}

	searchResult, err := esi.lib.Search().
		Index(esi.index).
		Type(typ).
		Source(obj).Do()

	return searchResult, err
}

func (esi *Index) SetMapping(typename string, jsn piazza.JsonString) error {

	ok := esi.IndexExists()
	if !ok {
		return fmt.Errorf("Index %s does not exist", esi.index)
	}

	putresp, err := esi.lib.PutMapping().Index(esi.index).Type(typename).BodyString(string(jsn)).Do()
	if err != nil {
		return fmt.Errorf("expected put mapping to succeed; got: %v", err)
	}
	if putresp == nil {
		return fmt.Errorf("expected put mapping response; got: %v", putresp)
	}
	if !putresp.Acknowledged {
		return fmt.Errorf("expected put mapping ack; got: %v", putresp.Acknowledged)
	}

	return nil
}

func (esi *Index) GetTypes() ([]string, error) {

	ok := esi.IndexExists()
	if !ok {
		return nil, fmt.Errorf("Index %s does not exist", esi.index)
	}

	getresp, err := esi.lib.IndexGet().Feature("_mappings").Index(esi.index).Do()
	if err != nil {
		return nil, err
	}

	typs := (*getresp[esi.index]).Mappings
	result := make([]string, len(typs))

	i := 0
	for k := range typs {
		result[i] = k
		i++
	}

	return result, nil
}

func (esi *Index) GetMapping(typ string) (interface{}, error) {

	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	getresp, err := esi.lib.GetMapping().Index(esi.index).Type(typ).Do()
	if err != nil {
		return nil, fmt.Errorf("expected get mapping to succeed; got: %v", err)
	}
	if getresp == nil {
		return nil, fmt.Errorf("expected get mapping response; got: %v", getresp)
	}
	props, ok := getresp[esi.index]
	if !ok {
		return nil, fmt.Errorf("expected JSON root to be of type map[string]interface{}; got: %s -- %#v", esi.index, getresp)
	}

	props2 := props.(map[string]interface{})

	return props2["mappings"], nil
}

func (esi *Index) AddPercolationQuery(id string, query piazza.JsonString) (*elastic.IndexResult, error) {

	ok := esi.IndexExists()
	if !ok {
		return nil, fmt.Errorf("Index %s does not exist", esi.index)
	}

	indexResponse, err := esi.lib.
		Index().
		Index(esi.index).
		Type(".percolator").
		Id(id).
		BodyString(string(query)).
		Do()
	if err != nil {
		return nil, err
	}

	return indexResponse, nil
}

func (esi *Index) DeletePercolationQuery(id string) (*elastic.DeleteResult, error) {
	typ := ".percolator"
	ok := esi.ItemExists(typ, id)
	if !ok {
		return nil, fmt.Errorf("Item %s in index %s and type %s does not exist", id, esi.index, typ)
	}

	deleteResult, err := esi.lib.Delete().
		Index(esi.index).
		Type(".percolator").
		Id(id).
		Do()
	if err != nil {
		return nil, err
	}

	return deleteResult, nil
}

func (esi *Index) AddPercolationDocument(typ string, doc interface{}) (*elastic.PercolateResponse, error) {
	ok := esi.TypeExists(typ)
	if !ok {
		return nil, fmt.Errorf("Type %s in index %s does not exist", typ, esi.index)
	}

	percolateResponse, err := esi.lib.
		Percolate().
		Index(esi.index).
		Type(typ).
		Doc(doc).
		//Pretty(true).
		Do()
	if err != nil {
		return nil, err
	}

	return percolateResponse, nil
}
<|endoftext|>


<commit_before>package dataStore

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

type DataStore interface {
	Clear() error
	Query(url string) bool
	// for now lets assume urls will be uploaded to us
	// via a comma separated list as one string
	Upload(urls string) bool
	// data stores that are accessing databases should
	// pull authentication from env vars instead of
	// by passing params to this function
	Initialize() error
	String()
}

type LocalDataStore struct {
	storage map[string]bool
}

func (data LocalDataStore) Clear() error {
	data.storage = make(map[string]bool)
	return nil
}

// we just want to know "does this url exist in the list"
// from the description we don't actually need to know the value
// any url in the list has the property the consumer wants to know about
func (data LocalDataStore) Query(url string) bool {
	log.Println("checking data store for url: ", url)
	_, ok := data.storage[url]
	return ok
}

// lets not care about the case where we're sent duplicate urls
// or urls the store already knows about for now
// it's not entirely clear if that's a case we actually care about
func (data LocalDataStore) Upload(urls string) bool {
	eachUrl := strings.Split(urls, ",")
	for _, url := range eachUrl {
		log.Println("adding url ", url, "to data store.")
		data.storage[url] = true
	}
	// remote data stores (redis) may fail for some reason?
	return true
}

func (data LocalDataStore) Initialize() error {
	// don't really need this function for local data store
	return nil
}

// for convenience/testing
func (data LocalDataStore) String() {
	for key, _ := range data.storage {
		fmt.Println(key)
	}
}

func NewLocalDataStore() LocalDataStore {
	var data LocalDataStore
	data.storage = make(map[string]bool)
	return data
}
<commit_msg>use log instead of fmt when printing storage<commit_after>package dataStore

import (
	"log"
	"strings"
)

type DataStore interface {
	Clear() error
	Query(url string) bool
	// for now lets assume urls will be uploaded to us
	// via a comma separated list as one string
	Upload(urls string) bool
	// data stores that are accessing databases should
	// pull authentication from env vars instead of
	// by passing params to this function
	Initialize() error
	String()
}

type LocalDataStore struct {
	storage map[string]bool
}

func (data LocalDataStore) Clear() error {
	data.storage = make(map[string]bool)
	return nil
}

// we just want to know "does this url exist in the list"
// from the description we don't actually need to know the value
// any url in the list has the property the consumer wants to know about
func (data LocalDataStore) Query(url string) bool {
	log.Println("checking data store for url: ", url)
	_, ok := data.storage[url]
	return ok
}

// lets not care about the case where we're sent duplicate urls
// or urls the store already knows about for now
// it's not entirely clear if that's a case we actually care about
func (data LocalDataStore) Upload(urls string) bool {
	eachUrl := strings.Split(urls, ",")
	for _, url := range eachUrl {
		log.Println("adding url ", url, "to data store.")
		data.storage[url] = true
	}
	// remote data stores (redis) may fail for some reason?
	return true
}

func (data LocalDataStore) Initialize() error {
	// don't really need this function for local data store
	return nil
}

// for convenience/testing
func (data LocalDataStore) String() {
	for key, _ := range data.storage {
		log.Println(key)
	}
}

func NewLocalDataStore() LocalDataStore {
	var data LocalDataStore
	data.storage = make(map[string]bool)
	return data
}
<|endoftext|>


<commit_before>package cmap

import (
	"strconv"
	"testing"

	"github.com/ivan1993spb/snake-server/engine"
)

const keyA = 0

func BenchmarkItems(b *testing.B) {
	m, _ := New(defaultShardCount)

	// Insert 100 elements.
	for i := 0; i < 10000; i++ {
		m.Set(uint16(i), Animal{strconv.Itoa(i)})
	}
	for i := 0; i < b.N; i++ {
		m.Items()
	}
}

func BenchmarkMarshalJson(b *testing.B) {
	m, _ := New(defaultShardCount)

	// Insert 100 elements.
	for i := 0; i < 10000; i++ {
		m.Set(uint16(i), Animal{strconv.Itoa(i)})
	}
	for i := 0; i < b.N; i++ {
		m.MarshalJSON()
	}
}

func BenchmarkStrconv(b *testing.B) {
	for i := 0; i < b.N; i++ {
		strconv.Itoa(i)
	}
}

func BenchmarkSingleInsertAbsent(b *testing.B) {
	m, _ := New(defaultShardCount)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Set(uint16(i), "value")
	}
}

func Benchmark_ConcurrentMap_MSet_MRemove(b *testing.B) {
	const dotsCount = 10
	const dotsPadding = dotsCount

	a := engine.MustArea(255, 255)
	m, _ := New(defaultShardCount)

	object := "value"

	b.ResetTimer()

	b.StopTimer()

	for i := 0; i < b.N; i += dotsPadding {
		dotsToBeRemoved := make([]uint16, 0, dotsCount)
		for j := 0; j < dotsCount; j++ {
			index := uint16(i+j) % a.Size()
			dot := engine.Dot{
				X: uint8(index % uint16(a.Width())),
				Y: uint8(index / uint16(a.Width())),
			}
			dotsToBeRemoved = append(dotsToBeRemoved, dot.Hash())
		}

		dotsToBeSet := make(map[uint16]interface{})
		for j := 0; j < dotsCount; j++ {
			index := uint16(i+j+dotsPadding) % a.Size()
			dot := engine.Dot{
				X: uint8(index % uint16(a.Width())),
				Y: uint8(index / uint16(a.Width())),
			}
			dotsToBeSet[dot.Hash()] = object
		}

		b.StartTimer()

		m.MRemove(dotsToBeRemoved)
		m.MSet(dotsToBeSet)

		b.StopTimer()
	}
}

func BenchmarkSingleInsertPresent(b *testing.B) {
	m, _ := New(defaultShardCount)
	m.Set(keyA, "value")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Set(keyA, "value")
	}
}

func benchmarkMultiInsertDifferent(b *testing.B, shardCount int) {
	m, _ := New(shardCount)
	finished := make(chan struct{}, b.N)
	_, set := GetSet(m, finished)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set(uint16(i), "value")
	}
	for i := 0; i < b.N; i++ {
		<-finished
	}
}

func BenchmarkMultiInsertDifferent_1_Shard(b *testing.B) {
	benchmarkMultiInsertDifferent(b, 1)
}
func BenchmarkMultiInsertDifferent_16_Shard(b *testing.B) {
	benchmarkMultiInsertDifferent(b, 16)
}
func BenchmarkMultiInsertDifferent_32_Shard(b *testing.B) {
	benchmarkMultiInsertDifferent(b, 32)
}
func BenchmarkMultiInsertDifferent_256_Shard(b *testing.B) {
	benchmarkMultiInsertDifferent(b, 256)
}

func BenchmarkMultiInsertSame(b *testing.B) {
	m, _ := New(defaultShardCount)
	finished := make(chan struct{}, b.N)
	_, set := GetSet(m, finished)
	m.Set(keyA, "value")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set(keyA, "value")
	}
	for i := 0; i < b.N; i++ {
		<-finished
	}
}

func BenchmarkMultiGetSame(b *testing.B) {
	m, _ := New(defaultShardCount)
	finished := make(chan struct{}, b.N)
	get, _ := GetSet(m, finished)
	m.Set(keyA, "value")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		get(keyA, "value")
	}
	for i := 0; i < b.N; i++ {
		<-finished
	}
}

func benchmarkMultiGetSetDifferent(b *testing.B, shardCount int) {
	m, _ := New(shardCount)
	finished := make(chan struct{}, 2*b.N)
	get, set := GetSet(m, finished)
	m.Set(0, "value")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set(uint16(i), "value")
		get(uint16(i+1), "value")
	}
	for i := 0; i < 2*b.N; i++ {
		<-finished
	}
}

func BenchmarkMultiGetSetDifferent_1_Shard(b *testing.B) {
	benchmarkMultiGetSetDifferent(b, 1)
}
func BenchmarkMultiGetSetDifferent_16_Shard(b *testing.B) {
	benchmarkMultiGetSetDifferent(b, 16)
}
func BenchmarkMultiGetSetDifferent_32_Shard(b *testing.B) {
	benchmarkMultiGetSetDifferent(b, 32)
}
func BenchmarkMultiGetSetDifferent_256_Shard(b *testing.B) {
	benchmarkMultiGetSetDifferent(b, 256)
}

func benchmarkMultiGetSetBlock(b *testing.B, shardCount int) {
	m, _ := New(shardCount)
	finished := make(chan struct{}, 2*b.N)
	get, set := GetSet(m, finished)
	for i := 0; i < b.N; i++ {
		m.Set(uint16(i%100), "value")
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set(uint16(i%100), "value")
		get(uint16(i%100), "value")
	}
	for i := 0; i < 2*b.N; i++ {
		<-finished
	}
}

func BenchmarkMultiGetSetBlock_1_Shard(b *testing.B) {
	benchmarkMultiGetSetBlock(b, 1)
}
func BenchmarkMultiGetSetBlock_16_Shard(b *testing.B) {
	benchmarkMultiGetSetBlock(b, 16)
}
func BenchmarkMultiGetSetBlock_32_Shard(b *testing.B) {
	benchmarkMultiGetSetBlock(b, 32)
}
func BenchmarkMultiGetSetBlock_256_Shard(b *testing.B) {
	benchmarkMultiGetSetBlock(b, 256)
}

func GetSet(m *ConcurrentMap, finished chan struct{}) (set func(key uint16, value string), get func(key uint16, value string)) {
	return func(key uint16, value string) {
			for i := 0; i < 10; i++ {
				m.Get(key)
			}
			finished <- struct{}{}
		}, func(key uint16, value string) {
			for i := 0; i < 10; i++ {
				m.Set(key, value)
			}
			finished <- struct{}{}
		}
}

func BenchmarkKeys(b *testing.B) {
	m, _ := New(defaultShardCount)

	// Insert 100 elements.
	for i := 0; i < 10000; i++ {
		m.Set(uint16(i), Animal{strconv.Itoa(i)})
	}
	for i := 0; i < b.N; i++ {
		m.Keys()
	}
}
<commit_msg>issue #43 Create benchmarks for ConcurrentMap<commit_after>package cmap

import (
	"math/rand"
	"strconv"
	"testing"
	"time"
)

const keyA = 0

func getHash(x, y uint8) uint16 {
	return uint16(x)<<8 | uint16(y)
}

func BenchmarkItems(b *testing.B) {
	m, _ := New(defaultShardCount)

	// Insert 100 elements.
	for i := 0; i < 10000; i++ {
		m.Set(uint16(i), Animal{strconv.Itoa(i)})
	}
	for i := 0; i < b.N; i++ {
		m.Items()
	}
}

func BenchmarkMarshalJson(b *testing.B) {
	m, _ := New(defaultShardCount)

	// Insert 100 elements.
	for i := 0; i < 10000; i++ {
		m.Set(uint16(i), Animal{strconv.Itoa(i)})
	}
	for i := 0; i < b.N; i++ {
		m.MarshalJSON()
	}
}

func BenchmarkStrconv(b *testing.B) {
	for i := 0; i < b.N; i++ {
		strconv.Itoa(i)
	}
}

func BenchmarkSingleInsertAbsent(b *testing.B) {
	m, _ := New(defaultShardCount)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Set(uint16(i), "value")
	}
}

func rawBenchmarkConcurrentMapSet(b *testing.B, shardCount int) {
	b.ReportAllocs()

	rand.Seed(time.Now().UTC().UnixNano())

	m, _ := New(shardCount)
	object := "value"

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		b.StopTimer()
		hash := uint16(rand.Int())
		b.StartTimer()

		m.Set(hash, object)
	}
}

func Benchmark_ConcurrentMap_Set_sc2(b *testing.B) {
	const shardCount = 2

	rawBenchmarkConcurrentMapSet(b, shardCount)
}

func Benchmark_ConcurrentMap_Set_sc32(b *testing.B) {
	const shardCount = 32

	rawBenchmarkConcurrentMapSet(b, shardCount)
}

func Benchmark_ConcurrentMap_Set_sc64(b *testing.B) {
	const shardCount = 64

	rawBenchmarkConcurrentMapSet(b, shardCount)
}

func rawBenchmarkConcurrentMapGet(b *testing.B, width, height uint8, shardCount int) {
	b.ReportAllocs()

	rand.Seed(time.Now().UTC().UnixNano())

	m, _ := New(shardCount)

	for y := uint8(0); y < height; y++ {
		for x := uint8(0); x < width; x++ {
			hash := getHash(x, y)

			if hash&1 == 1 {
				m.Set(hash, "object")
			}
		}
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		b.StopTimer()
		x := uint8(rand.Intn(int(width)))
		y := uint8(rand.Intn(int(height)))
		hash := getHash(x, y)
		b.StartTimer()

		m.Get(hash)
	}
}

func Benchmark_ConcurrentMap_Get_8x8_sc2(b *testing.B) {
	const (
		width  = 8
		height = 8

		shardCount = 2
	)

	rawBenchmarkConcurrentMapGet(b, width, height, shardCount)
}

func Benchmark_ConcurrentMap_Get_128x128_sc32(b *testing.B) {
	const (
		width  = 128
		height = 128

		shardCount = 32
	)

	rawBenchmarkConcurrentMapGet(b, width, height, shardCount)
}

func Benchmark_ConcurrentMap_Get_255x255_sc64(b *testing.B) {
	const (
		width  = 255
		height = 255

		shardCount = 64
	)

	rawBenchmarkConcurrentMapGet(b, width, height, shardCount)
}

func rawBenchmarkConcurrentMapMSetMRemove(b *testing.B, width, height uint8, shardCount, dotsCount int) {
	b.ReportAllocs()

	rand.Seed(time.Now().UTC().UnixNano())

	size := uint16(width) * uint16(height)
	m, _ := New(shardCount)

	object := "value"

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		b.StopTimer()

		dotIndex := rand.Int()

		dotsToBeRemoved := make([]uint16, 0, dotsCount)
		dotsToBeSet := make(map[uint16]interface{})
		for j := 0; j < dotsCount; j++ {
			index := uint16(dotIndex+j) % size
			hash := getHash(uint8(index%uint16(width)), uint8(index/uint16(width)))
			dotsToBeRemoved = append(dotsToBeRemoved, hash)
			dotsToBeSet[hash] = object
		}

		b.StartTimer()

		m.MSet(dotsToBeSet)
		m.MRemove(dotsToBeRemoved)
	}
}

func Benchmark_ConcurrentMap_MSet_MRemove_64x64_sc16_d12(b *testing.B) {
	const (
		width  = 64
		height = 64

		shardCount = 16

		dotsCount = 12
	)

	rawBenchmarkConcurrentMapMSetMRemove(b, width, height, shardCount, dotsCount)
}

func Benchmark_ConcurrentMap_MSet_MRemove_128x128_sc32_d32(b *testing.B) {
	const (
		width  = 128
		height = 128

		shardCount = 32

		dotsCount = 32
	)

	rawBenchmarkConcurrentMapMSetMRemove(b, width, height, shardCount, dotsCount)
}

func Benchmark_ConcurrentMap_MSet_MRemove_255x255_sc64_d64(b *testing.B) {
	const (
		width  = 255
		height = 255

		shardCount = 64

		dotsCount = 64
	)

	rawBenchmarkConcurrentMapMSetMRemove(b, width, height, shardCount, dotsCount)
}

func Benchmark_ConcurrentMap_MSet_MRemove_255x255_sc64_d256(b *testing.B) {
	const (
		width  = 255
		height = 255

		shardCount = 64

		dotsCount = 256
	)

	rawBenchmarkConcurrentMapMSetMRemove(b, width, height, shardCount, dotsCount)
}

func rawBenchmarkConcurrentMapMSetIfAbsentMRemoveCb(b *testing.B, width, height uint8, shardCount, dotsCount int) {
	b.ReportAllocs()

	rand.Seed(time.Now().UTC().UnixNano())

	size := uint16(width) * uint16(height)
	m, _ := New(shardCount)

	object := "value"

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		b.StopTimer()

		dotIndex := rand.Int()

		dotsToBeRemoved := make([]uint16, 0, dotsCount)
		dotsToBeSet := make(map[uint16]interface{})

		for j := 0; j < dotsCount; j++ {
			index := uint16(dotIndex+j) % size
			hash := getHash(uint8(index%uint16(width)), uint8(index/uint16(width)))
			dotsToBeRemoved = append(dotsToBeRemoved, hash)
			dotsToBeSet[hash] = object
		}

		b.StartTimer()

		m.MSetIfAbsent(dotsToBeSet)
		m.MRemoveCb(dotsToBeRemoved, func(key uint16, v interface{}, exists bool) bool {
			return exists && v == object
		})
	}
}

func Benchmark_ConcurrentMap_MSetIfAbsent_MRemoveCb_64x64_sc16_d12(b *testing.B) {
	const (
		width  = 64
		height = 64

		shardCount = 16

		dotsCount = 12
	)

	rawBenchmarkConcurrentMapMSetIfAbsentMRemoveCb(b, width, height, dotsCount, shardCount)
}

func Benchmark_ConcurrentMap_MSetIfAbsent_MRemoveCb_128x128_sc32_d32(b *testing.B) {
	const (
		width  = 128
		height = 128

		shardCount = 32

		dotsCount = 32
	)

	rawBenchmarkConcurrentMapMSetIfAbsentMRemoveCb(b, width, height, dotsCount, shardCount)
}

func Benchmark_ConcurrentMap_MSetIfAbsent_MRemoveCb_255x255_sc64_d64(b *testing.B) {
	const (
		width  = 255
		height = 255

		shardCount = 64

		dotsCount = 64
	)

	rawBenchmarkConcurrentMapMSetIfAbsentMRemoveCb(b, width, height, dotsCount, shardCount)
}

func Benchmark_ConcurrentMap_MSetIfAbsent_MRemoveCb_255x255_sc64_d256(b *testing.B) {
	const (
		width  = 255
		height = 255

		shardCount = 64

		dotsCount = 256
	)

	rawBenchmarkConcurrentMapMSetIfAbsentMRemoveCb(b, width, height, dotsCount, shardCount)
}

func rawBenchmarkConcurrentMapSetMSetIfAllAbsentRemoveCb(b *testing.B, width, height uint8, shardCount, dotsCount int) {
	b.ReportAllocs()

	rand.Seed(time.Now().UTC().UnixNano())

	size := uint16(width) * uint16(height)
	m, _ := New(shardCount)

	object := "value"

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		b.StopTimer()

		dotIndex := rand.Int()

		dotsToBeRemoved := make([]uint16, 0, dotsCount)
		dotsToBeSet := make(map[uint16]interface{})
		for j := 0; j < dotsCount; j++ {
			index := uint16(dotIndex+j) % size
			hash := getHash(uint8(index%uint16(width)), uint8(index/uint16(width)))
			dotsToBeRemoved = append(dotsToBeRemoved, hash)
			dotsToBeSet[hash] = object
		}

		spoiler := dotsToBeRemoved[dotsCount-1]

		b.StartTimer()

		m.Set(spoiler, object)
		m.MSetIfAllAbsent(dotsToBeSet)
		m.MRemoveCb(dotsToBeRemoved, func(key uint16, v interface{}, exists bool) bool {
			return exists && v == object
		})
	}
}

func Benchmark_ConcurrentMap_Set_MSetIfAllAbsent_RemoveCb_64x64_sc16_d12(b *testing.B) {
	const (
		width  = 64
		height = 64

		shardCount = 16

		dotsCount = 12
	)

	rawBenchmarkConcurrentMapSetMSetIfAllAbsentRemoveCb(b, width, height, shardCount, dotsCount)
}

func Benchmark_ConcurrentMap_Set_MSetIfAllAbsent_RemoveCb_128x128_sc32_d32(b *testing.B) {
	const (
		width  = 128
		height = 128

		shardCount = 32

		dotsCount = 32
	)

	rawBenchmarkConcurrentMapSetMSetIfAllAbsentRemoveCb(b, width, height, shardCount, dotsCount)
}

func Benchmark_ConcurrentMap_Set_MSetIfAllAbsent_RemoveCb_255x255_sc64_d64(b *testing.B) {
	const (
		width  = 255
		height = 255

		shardCount = 64

		dotsCount = 64
	)

	rawBenchmarkConcurrentMapSetMSetIfAllAbsentRemoveCb(b, width, height, shardCount, dotsCount)
}

func Benchmark_ConcurrentMap_Set_MSetIfAllAbsent_RemoveCb_255x255_sc64_d256(b *testing.B) {
	const (
		width  = 255
		height = 255

		shardCount = 64

		dotsCount = 256
	)

	rawBenchmarkConcurrentMapSetMSetIfAllAbsentRemoveCb(b, width, height, shardCount, dotsCount)
}

func BenchmarkSingleInsertPresent(b *testing.B) {
	m, _ := New(defaultShardCount)
	m.Set(keyA, "value")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		m.Set(keyA, "value")
	}
}

func benchmarkMultiInsertDifferent(b *testing.B, shardCount int) {
	m, _ := New(shardCount)
	finished := make(chan struct{}, b.N)
	_, set := GetSet(m, finished)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set(uint16(i), "value")
	}
	for i := 0; i < b.N; i++ {
		<-finished
	}
}

func BenchmarkMultiInsertDifferent_1_Shard(b *testing.B) {
	benchmarkMultiInsertDifferent(b, 1)
}
func BenchmarkMultiInsertDifferent_16_Shard(b *testing.B) {
	benchmarkMultiInsertDifferent(b, 16)
}
func BenchmarkMultiInsertDifferent_32_Shard(b *testing.B) {
	benchmarkMultiInsertDifferent(b, 32)
}
func BenchmarkMultiInsertDifferent_256_Shard(b *testing.B) {
	benchmarkMultiInsertDifferent(b, 256)
}

func BenchmarkMultiInsertSame(b *testing.B) {
	m, _ := New(defaultShardCount)
	finished := make(chan struct{}, b.N)
	_, set := GetSet(m, finished)
	m.Set(keyA, "value")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set(keyA, "value")
	}
	for i := 0; i < b.N; i++ {
		<-finished
	}
}

func BenchmarkMultiGetSame(b *testing.B) {
	m, _ := New(defaultShardCount)
	finished := make(chan struct{}, b.N)
	get, _ := GetSet(m, finished)
	m.Set(keyA, "value")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		get(keyA, "value")
	}
	for i := 0; i < b.N; i++ {
		<-finished
	}
}

func benchmarkMultiGetSetDifferent(b *testing.B, shardCount int) {
	m, _ := New(shardCount)
	finished := make(chan struct{}, 2*b.N)
	get, set := GetSet(m, finished)
	m.Set(0, "value")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set(uint16(i), "value")
		get(uint16(i+1), "value")
	}
	for i := 0; i < 2*b.N; i++ {
		<-finished
	}
}

func BenchmarkMultiGetSetDifferent_1_Shard(b *testing.B) {
	benchmarkMultiGetSetDifferent(b, 1)
}
func BenchmarkMultiGetSetDifferent_16_Shard(b *testing.B) {
	benchmarkMultiGetSetDifferent(b, 16)
}
func BenchmarkMultiGetSetDifferent_32_Shard(b *testing.B) {
	benchmarkMultiGetSetDifferent(b, 32)
}
func BenchmarkMultiGetSetDifferent_256_Shard(b *testing.B) {
	benchmarkMultiGetSetDifferent(b, 256)
}

func benchmarkMultiGetSetBlock(b *testing.B, shardCount int) {
	m, _ := New(shardCount)
	finished := make(chan struct{}, 2*b.N)
	get, set := GetSet(m, finished)
	for i := 0; i < b.N; i++ {
		m.Set(uint16(i%100), "value")
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		set(uint16(i%100), "value")
		get(uint16(i%100), "value")
	}
	for i := 0; i < 2*b.N; i++ {
		<-finished
	}
}

func BenchmarkMultiGetSetBlock_1_Shard(b *testing.B) {
	benchmarkMultiGetSetBlock(b, 1)
}
func BenchmarkMultiGetSetBlock_16_Shard(b *testing.B) {
	benchmarkMultiGetSetBlock(b, 16)
}
func BenchmarkMultiGetSetBlock_32_Shard(b *testing.B) {
	benchmarkMultiGetSetBlock(b, 32)
}
func BenchmarkMultiGetSetBlock_256_Shard(b *testing.B) {
	benchmarkMultiGetSetBlock(b, 256)
}

func GetSet(m *ConcurrentMap, finished chan struct{}) (set func(key uint16, value string), get func(key uint16, value string)) {
	return func(key uint16, value string) {
			for i := 0; i < 10; i++ {
				m.Get(key)
			}
			finished <- struct{}{}
		}, func(key uint16, value string) {
			for i := 0; i < 10; i++ {
				m.Set(key, value)
			}
			finished <- struct{}{}
		}
}

func BenchmarkKeys(b *testing.B) {
	m, _ := New(defaultShardCount)

	// Insert 100 elements.
	for i := 0; i < 10000; i++ {
		m.Set(uint16(i), Animal{strconv.Itoa(i)})
	}
	for i := 0; i < b.N; i++ {
		m.Keys()
	}
}
<|endoftext|>