<commit_before>/*
Copyright 2016 Tamás Gulácsi

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 oracall

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	fstructs "github.com/fatih/structs"
	"github.com/pkg/errors"
)

var Gogo bool
var NumberAsString bool

//go:generate sh ./download-protoc.sh
//go:generate go get -u github.com/gogo/protobuf/protoc-gen-gofast

// build: protoc --gofast_out=plugins=grpc:. my.proto
// build: protoc --go_out=plugins=grpc:. my.proto

func SaveProtobuf(dst io.Writer, functions []Function, pkg string) error {
	var err error
	w := errWriter{Writer: dst, err: &err}

	io.WriteString(w, `syntax = "proto3";`+"\n\n")

	if pkg != "" {
		fmt.Fprintf(w, "package %s;\n", pkg)
	}
	if Gogo {
		io.WriteString(w, `
	import "github.com/gogo/protobuf/gogoproto/gogo.proto";
`)
	}
	seen := make(map[string]struct{}, 16)

	services := make([]string, 0, len(functions))

FunLoop:
	for _, fun := range functions {
		fName := strings.ToLower(fun.name)
		if err := fun.SaveProtobuf(w, seen); err != nil {
			if errors.Cause(err) == ErrMissingTableOf {
				Log("msg", "SKIP function, missing TableOf info", "function", fName)
				continue FunLoop
			}
			return err
		}
		var streamQual string
		if fun.HasCursorOut() {
			streamQual = "stream "
		}
		name := CamelCase(dot2D.Replace(fName))
		services = append(services,
			fmt.Sprintf(`rpc %s (%s) returns (%s%s) {}`,
				name,
				CamelCase(fun.getStructName(false, false)),
				streamQual,
				CamelCase(fun.getStructName(true, false)),
			),
		)
	}

	fmt.Fprintf(w, "\nservice %s {\n", CamelCase(pkg))
	for _, s := range services {
		fmt.Fprintf(w, "\t%s\n", s)
	}
	w.Write([]byte("}"))

	return nil
}

func (f Function) SaveProtobuf(dst io.Writer, seen map[string]struct{}) error {
	var buf bytes.Buffer
	if err := f.saveProtobufDir(&buf, seen, false); err != nil {
		return errors.Wrap(err, "input")
	}
	if err := f.saveProtobufDir(&buf, seen, true); err != nil {
		return errors.Wrap(err, "output")
	}
	_, err := dst.Write(buf.Bytes())
	return err
}
func (f Function) saveProtobufDir(dst io.Writer, seen map[string]struct{}, out bool) error {
	dirmap, dirname := uint8(DIR_IN), "input"
	if out {
		dirmap, dirname = DIR_OUT, "output"
	}
	args := make([]Argument, 0, len(f.Args)+1)
	for _, arg := range f.Args {
		if arg.Direction&dirmap > 0 {
			args = append(args, arg)
		}
	}
	// return variable for function out structs
	if out && f.Returns != nil {
		args = append(args, *f.Returns)
	}

	return protoWriteMessageTyp(dst,
		CamelCase(dot2D.Replace(strings.ToLower(f.name))+"__"+dirname),
		seen, args...)
}

var dot2D = strings.NewReplacer(".", "__")

func protoWriteMessageTyp(dst io.Writer, msgName string, seen map[string]struct{}, args ...Argument) error {
	for _, arg := range args {
		if arg.Flavor == FLAVOR_TABLE && arg.TableOf == nil {
			return errors.Wrapf(ErrMissingTableOf, "no table of data for %s.%s (%v)", msgName, arg, arg)
		}
	}

	var err error
	w := errWriter{Writer: dst, err: &err}
	fmt.Fprintf(w, "\nmessage %s {\n", msgName)

	buf := buffers.Get()
	defer buffers.Put(buf)
	for i, arg := range args {
		var rule string
		if strings.HasSuffix(arg.Name, "#") {
			arg.Name = replHidden(arg.Name)
		}
		if arg.Flavor == FLAVOR_TABLE {
			if arg.TableOf == nil {
				return errors.Wrapf(ErrMissingTableOf, "no table of data for %s.%s (%v)", msgName, arg, arg)
			}
			rule = "repeated "
		}
		aName := arg.Name
		got := arg.goType(false)
		if strings.HasPrefix(got, "*") {
			got = got[1:]
		}
		if strings.HasPrefix(got, "[]") {
			rule = "repeated "
			got = got[2:]
		}
		if strings.HasPrefix(got, "*") {
			got = got[1:]
		}
		if got == "" {
			got = mkRecTypName(arg.Name)
		}
		typ, pOpts := protoType(got, arg.Name)
		var optS string
		if s := pOpts.String(); s != "" {
			optS = " " + s
		}
		if arg.Flavor == FLAVOR_SIMPLE || arg.Flavor == FLAVOR_TABLE && arg.TableOf.Flavor == FLAVOR_SIMPLE {
			fmt.Fprintf(w, "\t%s%s %s = %d%s;\t// %s\n", rule, typ, aName, i+1, optS, arg.AbsType)
			continue
		}
		typ = CamelCase(typ)
		if _, ok := seen[typ]; !ok {
			//lName := strings.ToLower(arg.Name)
			subArgs := make([]Argument, 0, 16)
			if arg.TableOf != nil {
				if arg.TableOf.RecordOf == nil {
					subArgs = append(subArgs, *arg.TableOf)
				} else {
					for _, v := range arg.TableOf.RecordOf {
						subArgs = append(subArgs, v.Argument)
					}
				}
			} else {
				for _, v := range arg.RecordOf {
					subArgs = append(subArgs, v.Argument)
				}
			}
			if err := protoWriteMessageTyp(buf, typ, seen, subArgs...); err != nil {
				Log("msg", "protoWriteMessageTyp", "error", err)
				return err
			}
			seen[typ] = struct{}{}
		}
		fmt.Fprintf(w, "\t%s%s %s = %d%s;\n", rule, typ, aName, i+1, optS)
	}
	io.WriteString(w, "}\n")
	w.Write(buf.Bytes())

	return err
}

func protoType(got, aName string) (string, protoOptions) {
	switch trimmed := strings.ToLower(strings.TrimPrefix(strings.TrimPrefix(got, "[]"), "*")); trimmed {
	case "ora.time", "time.time":
		return "string", nil
	case "ora.string":
		return "string", nil

	case "int32", "ora.int32":
		if NumberAsString {
			return "sint32", protoOptions{
				"gogoproto.jsontag": aName + ",string,omitempty",
			}
		}
		return "sint32", nil
	case "float64", "ora.float64":
		if NumberAsString {
			return "double", protoOptions{
				"gogoproto.jsontag": aName + ",string,omitempty",
			}
		}
		return "double", nil

	case "ora.date", "custom.date":
		return "string", nil
	case "n", "ora.n":
		return "string", nil
	case "ora.lob":
		return "bytes", nil
	default:
		return trimmed, nil
	}
}

type protoOptions map[string]interface{}

func (opts protoOptions) String() string {
	if len(opts) == 0 {
		return ""
	}
	var buf bytes.Buffer
	buf.WriteByte('[')
	for k, v := range opts {
		if buf.Len() != 1 {
			buf.WriteString(", ")
		}
		fmt.Fprintf(&buf, "(%s)=", k)
		switch v.(type) {
		case bool:
			fmt.Fprintf(&buf, "%t", v)
		default:
			fmt.Fprintf(&buf, "%q", v)
		}
	}
	buf.WriteByte(']')
	return buf.String()
}

func CopyStruct(dest interface{}, src interface{}) error {
	ds := fstructs.New(dest)
	ss := fstructs.New(src)
	snames := ss.Names()
	svalues := ss.Values()
	for _, df := range ds.Fields() {
		dnm := df.Name()
		for i, snm := range snames {
			if snm == dnm || dnm == CamelCase(snm) || CamelCase(dnm) == snm {
				svalue := svalues[i]
				if err := df.Set(svalue); err != nil {
					return errors.Wrapf(err, "set %q to %q (%v %T)", dnm, snm, svalue, svalue)
				}
			}
		}
	}
	return nil
}
func mkRecTypName(name string) string { return strings.ToLower(name) + "_rek_typ" }
<commit_msg>move comment before field<commit_after>/*
Copyright 2016 Tamás Gulácsi

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 oracall

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	fstructs "github.com/fatih/structs"
	"github.com/pkg/errors"
)

var Gogo bool
var NumberAsString bool

//go:generate sh ./download-protoc.sh
//go:generate go get -u github.com/gogo/protobuf/protoc-gen-gofast

// build: protoc --gofast_out=plugins=grpc:. my.proto
// build: protoc --go_out=plugins=grpc:. my.proto

func SaveProtobuf(dst io.Writer, functions []Function, pkg string) error {
	var err error
	w := errWriter{Writer: dst, err: &err}

	io.WriteString(w, `syntax = "proto3";`+"\n\n")

	if pkg != "" {
		fmt.Fprintf(w, "package %s;\n", pkg)
	}
	if Gogo {
		io.WriteString(w, `
	import "github.com/gogo/protobuf/gogoproto/gogo.proto";
`)
	}
	seen := make(map[string]struct{}, 16)

	services := make([]string, 0, len(functions))

FunLoop:
	for _, fun := range functions {
		fName := strings.ToLower(fun.name)
		if err := fun.SaveProtobuf(w, seen); err != nil {
			if errors.Cause(err) == ErrMissingTableOf {
				Log("msg", "SKIP function, missing TableOf info", "function", fName)
				continue FunLoop
			}
			return err
		}
		var streamQual string
		if fun.HasCursorOut() {
			streamQual = "stream "
		}
		name := CamelCase(dot2D.Replace(fName))
		services = append(services,
			fmt.Sprintf(`rpc %s (%s) returns (%s%s) {}`,
				name,
				CamelCase(fun.getStructName(false, false)),
				streamQual,
				CamelCase(fun.getStructName(true, false)),
			),
		)
	}

	fmt.Fprintf(w, "\nservice %s {\n", CamelCase(pkg))
	for _, s := range services {
		fmt.Fprintf(w, "\t%s\n", s)
	}
	w.Write([]byte("}"))

	return nil
}

func (f Function) SaveProtobuf(dst io.Writer, seen map[string]struct{}) error {
	var buf bytes.Buffer
	if err := f.saveProtobufDir(&buf, seen, false); err != nil {
		return errors.Wrap(err, "input")
	}
	if err := f.saveProtobufDir(&buf, seen, true); err != nil {
		return errors.Wrap(err, "output")
	}
	_, err := dst.Write(buf.Bytes())
	return err
}
func (f Function) saveProtobufDir(dst io.Writer, seen map[string]struct{}, out bool) error {
	dirmap, dirname := uint8(DIR_IN), "input"
	if out {
		dirmap, dirname = DIR_OUT, "output"
	}
	args := make([]Argument, 0, len(f.Args)+1)
	for _, arg := range f.Args {
		if arg.Direction&dirmap > 0 {
			args = append(args, arg)
		}
	}
	// return variable for function out structs
	if out && f.Returns != nil {
		args = append(args, *f.Returns)
	}

	return protoWriteMessageTyp(dst,
		CamelCase(dot2D.Replace(strings.ToLower(f.name))+"__"+dirname),
		seen, args...)
}

var dot2D = strings.NewReplacer(".", "__")

func protoWriteMessageTyp(dst io.Writer, msgName string, seen map[string]struct{}, args ...Argument) error {
	for _, arg := range args {
		if arg.Flavor == FLAVOR_TABLE && arg.TableOf == nil {
			return errors.Wrapf(ErrMissingTableOf, "no table of data for %s.%s (%v)", msgName, arg, arg)
		}
	}

	var err error
	w := errWriter{Writer: dst, err: &err}
	fmt.Fprintf(w, "\nmessage %s {\n", msgName)

	buf := buffers.Get()
	defer buffers.Put(buf)
	for i, arg := range args {
		var rule string
		if strings.HasSuffix(arg.Name, "#") {
			arg.Name = replHidden(arg.Name)
		}
		if arg.Flavor == FLAVOR_TABLE {
			if arg.TableOf == nil {
				return errors.Wrapf(ErrMissingTableOf, "no table of data for %s.%s (%v)", msgName, arg, arg)
			}
			rule = "repeated "
		}
		aName := arg.Name
		got := arg.goType(false)
		if strings.HasPrefix(got, "*") {
			got = got[1:]
		}
		if strings.HasPrefix(got, "[]") {
			rule = "repeated "
			got = got[2:]
		}
		if strings.HasPrefix(got, "*") {
			got = got[1:]
		}
		if got == "" {
			got = mkRecTypName(arg.Name)
		}
		typ, pOpts := protoType(got, arg.Name)
		var optS string
		if s := pOpts.String(); s != "" {
			optS = " " + s
		}
		if arg.Flavor == FLAVOR_SIMPLE || arg.Flavor == FLAVOR_TABLE && arg.TableOf.Flavor == FLAVOR_SIMPLE {
			fmt.Fprintf(w, "\t// %s\n\t%s%s %s = %d%s;\n", arg.AbsType, rule, typ, aName, i+1, optS)
			continue
		}
		typ = CamelCase(typ)
		if _, ok := seen[typ]; !ok {
			//lName := strings.ToLower(arg.Name)
			subArgs := make([]Argument, 0, 16)
			if arg.TableOf != nil {
				if arg.TableOf.RecordOf == nil {
					subArgs = append(subArgs, *arg.TableOf)
				} else {
					for _, v := range arg.TableOf.RecordOf {
						subArgs = append(subArgs, v.Argument)
					}
				}
			} else {
				for _, v := range arg.RecordOf {
					subArgs = append(subArgs, v.Argument)
				}
			}
			if err := protoWriteMessageTyp(buf, typ, seen, subArgs...); err != nil {
				Log("msg", "protoWriteMessageTyp", "error", err)
				return err
			}
			seen[typ] = struct{}{}
		}
		fmt.Fprintf(w, "\t%s%s %s = %d%s;\n", rule, typ, aName, i+1, optS)
	}
	io.WriteString(w, "}\n")
	w.Write(buf.Bytes())

	return err
}

func protoType(got, aName string) (string, protoOptions) {
	switch trimmed := strings.ToLower(strings.TrimPrefix(strings.TrimPrefix(got, "[]"), "*")); trimmed {
	case "ora.time", "time.time":
		return "string", nil
	case "ora.string":
		return "string", nil

	case "int32", "ora.int32":
		if NumberAsString {
			return "sint32", protoOptions{
				"gogoproto.jsontag": aName + ",string,omitempty",
			}
		}
		return "sint32", nil
	case "float64", "ora.float64":
		if NumberAsString {
			return "double", protoOptions{
				"gogoproto.jsontag": aName + ",string,omitempty",
			}
		}
		return "double", nil

	case "ora.date", "custom.date":
		return "string", nil
	case "n", "ora.n":
		return "string", nil
	case "ora.lob":
		return "bytes", nil
	default:
		return trimmed, nil
	}
}

type protoOptions map[string]interface{}

func (opts protoOptions) String() string {
	if len(opts) == 0 {
		return ""
	}
	var buf bytes.Buffer
	buf.WriteByte('[')
	for k, v := range opts {
		if buf.Len() != 1 {
			buf.WriteString(", ")
		}
		fmt.Fprintf(&buf, "(%s)=", k)
		switch v.(type) {
		case bool:
			fmt.Fprintf(&buf, "%t", v)
		default:
			fmt.Fprintf(&buf, "%q", v)
		}
	}
	buf.WriteByte(']')
	return buf.String()
}

func CopyStruct(dest interface{}, src interface{}) error {
	ds := fstructs.New(dest)
	ss := fstructs.New(src)
	snames := ss.Names()
	svalues := ss.Values()
	for _, df := range ds.Fields() {
		dnm := df.Name()
		for i, snm := range snames {
			if snm == dnm || dnm == CamelCase(snm) || CamelCase(dnm) == snm {
				svalue := svalues[i]
				if err := df.Set(svalue); err != nil {
					return errors.Wrapf(err, "set %q to %q (%v %T)", dnm, snm, svalue, svalue)
				}
			}
		}
	}
	return nil
}
func mkRecTypName(name string) string { return strings.ToLower(name) + "_rek_typ" }
<|endoftext|>


<commit_before>// Copyright 2016-2017 The psh Authors. All rights reserved.

package psh

import (
	"bytes"
	"fmt"
	"os"
	"path"
)

const (
	// SegmentVirtualEnvBackground is the background color to use
	SegmentVirtualEnvBackground = 22 // #585858
)

// SegmentVirtualEnv implements the python virtual env partial of the prompt.
//
// Renders the current virtalenv python environment name if it is available.
type SegmentVirtualEnv struct {
	Data []byte
}

// NewSegmentVirtualEnv creates an instace of SegmentVirtualEnv type.
func NewSegmentVirtualEnv() *SegmentVirtualEnv {
	segment := &SegmentVirtualEnv{}
	return segment
}

// Compile collects the data for this segment.
func (s *SegmentVirtualEnv) Compile() {
	s.Data = []byte("")
	virtualEnvPath, ok := os.LookupEnv("VIRTUAL_ENV")
	if ok {
		virtualEnvName := path.Base(virtualEnvPath)
		s.Data = []byte(virtualEnvName)
	}
}

// Render renders the segment results.
func (s *SegmentVirtualEnv) Render() []byte {
	var b bytes.Buffer
	if len(s.Data) > 0 {
		b.Write(SetBackground(SegmentVirtualEnvBackground))
		fmt.Fprint(&b, " ")
		b.Write(s.Data)
		fmt.Fprint(&b, " ")
	}
	return b.Bytes()
}
<commit_msg>Added python powerline icon to the virtualenv segment<commit_after>// Copyright 2016-2017 The psh Authors. All rights reserved.

package psh

import (
	"bytes"
	"fmt"
	"os"
	"path"
)

const (
	// SegmentVirtualEnvBackground is the background color to use
	SegmentVirtualEnvBackground = 22 // #585858
)

// SegmentVirtualEnv implements the python virtual env partial of the prompt.
//
// Renders the current virtalenv python environment name if it is available.
type SegmentVirtualEnv struct {
	Data []byte
}

// NewSegmentVirtualEnv creates an instace of SegmentVirtualEnv type.
func NewSegmentVirtualEnv() *SegmentVirtualEnv {
	segment := &SegmentVirtualEnv{}
	return segment
}

// Compile collects the data for this segment.
func (s *SegmentVirtualEnv) Compile() {
	s.Data = []byte("")
	virtualEnvPath, ok := os.LookupEnv("VIRTUAL_ENV")
	if ok {
		virtualEnvName := path.Base(virtualEnvPath)
		s.Data = []byte(virtualEnvName)
	}
}

// Render renders the segment results.
func (s *SegmentVirtualEnv) Render() []byte {
	var b bytes.Buffer
	if len(s.Data) > 0 {
		b.Write(SetBackground(SegmentVirtualEnvBackground))
		fmt.Fprint(&b, " ")
		fmt.Fprint(&b, " ")
		b.Write(s.Data)
		fmt.Fprint(&b, " ")
	}
	return b.Bytes()
}
<|endoftext|>


<commit_before>package layout

import (
	"time"

	"github.com/goccy/go-yaml"
	"github.com/spf13/cast"
	tele "gopkg.in/tucnak/telebot.v3"
)

// Config represents typed map interface related to the "config" section in layout.
type Config struct {
	v map[string]interface{}
}

// Unmarshal parses the config into the out value. It's useful when you want to
// describe and to pre-define the fields in your custom configuration struct.
func (c *Config) Unmarshal(v interface{}) error {
	data, err := yaml.Marshal(c.v)
	if err != nil {
		return err
	}
	return yaml.Unmarshal(data, v)
}

// Get returns child map field wrapped into Config.
// If the field isn't map, returns nil.
func (c *Config) Get(k string) *Config {
	v, ok := c.v[k].(map[string]interface{})
	if !ok {
		return nil
	}
	return &Config{v: v}
}

// String returns a field casted to the string.
func (c *Config) String(k string) string {
	return cast.ToString(c.v[k])
}

// Int returns a field casted to the int.
func (c *Config) Int(k string) int {
	return cast.ToInt(c.v[k])
}

// Int64 returns a field casted to the int64.
func (c *Config) Int64(k string) int64 {
	return cast.ToInt64(c.v[k])
}

// Float returns a field casted to the float64.
func (c *Config) Float(k string) float64 {
	return cast.ToFloat64(c.v[k])
}

// Float returns a field casted to the bool.
func (c *Config) Bool(k string) bool {
	return cast.ToBool(c.v[k])
}

// Duration returns a field casted to the time.Duration.
// Accepts number-represented duration or a string in 0nsuµmh format.
func (c *Config) Duration(k string) time.Duration {
	return cast.ToDuration(c.v[k])
}

// ChatID returns a field casted to the ChatID.
// The value must be an integer.
func (c *Config) ChatID(k string) tele.ChatID {
	return tele.ChatID(c.Int64(k))
}
<commit_msg>layout: add config functions for slices<commit_after>package layout

import (
	"time"

	"github.com/goccy/go-yaml"
	"github.com/spf13/cast"
	tele "gopkg.in/tucnak/telebot.v3"
)

// Config represents typed map interface related to the "config" section in layout.
type Config struct {
	v map[string]interface{}
}

// Unmarshal parses the config into the out value. It's useful when you want to
// describe and to pre-define the fields in your custom configuration struct.
func (c *Config) Unmarshal(v interface{}) error {
	data, err := yaml.Marshal(c.v)
	if err != nil {
		return err
	}
	return yaml.Unmarshal(data, v)
}

// Get returns child map field wrapped into Config.
// If the field isn't map, returns nil.
func (c *Config) Get(k string) *Config {
	v, ok := c.v[k].(map[string]interface{})
	if !ok {
		return nil
	}
	return &Config{v: v}
}

// String returns a field casted to the string.
func (c *Config) String(k string) string {
	return cast.ToString(c.v[k])
}

// Int returns a field casted to the int.
func (c *Config) Int(k string) int {
	return cast.ToInt(c.v[k])
}

// Int64 returns a field casted to the int64.
func (c *Config) Int64(k string) int64 {
	return cast.ToInt64(c.v[k])
}

// Float returns a field casted to the float64.
func (c *Config) Float(k string) float64 {
	return cast.ToFloat64(c.v[k])
}

// Float returns a field casted to the bool.
func (c *Config) Bool(k string) bool {
	return cast.ToBool(c.v[k])
}

// Duration returns a field casted to the time.Duration.
// Accepts number-represented duration or a string in 0nsuµmh format.
func (c *Config) Duration(k string) time.Duration {
	return cast.ToDuration(c.v[k])
}

// ChatID returns a field casted to the ChatID.
// The value must be an integer.
func (c *Config) ChatID(k string) tele.ChatID {
	return tele.ChatID(c.Int64(k))
}

// Strings returns a field casted to the string slice.
func (c *Config) Strings(k string) []string {
	return cast.ToStringSlice(c.v[k])
}

// Ints returns a field casted to the int slice.
func (c *Config) Ints(k string) []int {
	return cast.ToIntSlice(c.v[k])
}

// Floats returns a field casted to the float slice.
func (c *Config) Floats(k string) []float64 {
	slice := cast.ToSlice(c.v[k])

	fs := make([]float64, len(slice))
	for i, a := range slice {
		fs[i] = cast.ToFloat64(a)
	}

	return fs
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime/pprof"
	"time"

	"github.com/mrfuxi/neural"
	"github.com/petar/GoMNIST"
)

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	nnSaveFile = flag.String("save-file", "", "Save neural network to file")
	nnLoadFile = flag.String("load-file", "", "Load neural network to file")
	inputSize  = GoMNIST.Width * GoMNIST.Height
)

func prepareMnistData(rawData *GoMNIST.Set) []neural.TrainExample {
	trainData := make([]neural.TrainExample, rawData.Count())
	for i := range trainData {
		image, label := rawData.Get(i)
		trainData[i].Input = make([]float64, inputSize, inputSize)
		trainData[i].Output = make([]float64, 10, 10)
		for j, pix := range image {
			trainData[i].Input[j] = (float64(pix) / 255)
		}

		for j := range trainData[i].Output {
			trainData[i].Output[j] = 0
		}
		trainData[i].Output[label] = 1
	}
	return trainData
}

func loadTestData() ([]neural.TrainExample, []neural.TrainExample) {
	train, test, err := GoMNIST.Load("./data")
	if err != nil {
		panic(err)
	}

	trainData := prepareMnistData(train)
	testData := prepareMnistData(test)
	return trainData, testData
}

func epocheCallback(nn neural.Evaluator, cost neural.Cost, trainData, testData []neural.TrainExample) neural.EpocheCallback {
	return func(epoche int, dt time.Duration) {
		avgCost, errors := neural.CalculateCorrectness(nn, cost, testData)
		fmt.Printf("%v,%v,%v\n", epoche, avgCost, errors)
	}
}

func main() {
	trainData, testData := loadTestData()

	activator := neural.NewSigmoidActivator()
	outActivator := neural.NewSoftmaxFunction()
	nn := neural.NewNeuralNetwork(
		[]int{inputSize, 100, 10},
		neural.NewFullyConnectedLayer(activator),
		neural.NewFullyConnectedLayer(outActivator),
	)

	flag.Parse()

	if *nnLoadFile != "" {
		fn, err := os.Open(*nnLoadFile)
		if err != nil {
			log.Fatalln(err)
		}
		if err := neural.Load(nn, fn); err != nil {
			log.Fatalln(err)
		}
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	cost := neural.NewLogLikelihoodCost()
	options := neural.TrainOptions{
		Epochs:         40,
		MiniBatchSize:  10,
		LearningRate:   0.4,
		Regularization: 5,
		TrainerFactory: neural.NewBackpropagationTrainer,
		EpocheCallback: epocheCallback(nn, cost, trainData, testData),
		Cost:           cost,
	}

	t0 := time.Now()
	neural.Train(nn, trainData, options)
	dt := time.Since(t0)

	fmt.Println("Training complete in", dt)

	if *nnSaveFile != "" {
		fn, err := os.OpenFile(*nnSaveFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
		if err != nil {
			log.Fatalln(err)
		}
		if err := neural.Save(nn, fn); err != nil {
			log.Fatalln(err)
		}
	}
}
<commit_msg>Split MNIST data into train, validation and test. Report accuracy over validation and test data after every epoche<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"runtime/pprof"
	"time"

	"github.com/mrfuxi/neural"
	"github.com/petar/GoMNIST"
)

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	nnSaveFile = flag.String("save-file", "", "Save neural network to file")
	nnLoadFile = flag.String("load-file", "", "Load neural network to file")
	inputSize  = GoMNIST.Width * GoMNIST.Height
)

func prepareMnistData(rawData *GoMNIST.Set) []neural.TrainExample {
	trainData := make([]neural.TrainExample, rawData.Count())
	for i := range trainData {
		image, label := rawData.Get(i)
		trainData[i].Input = make([]float64, inputSize, inputSize)
		trainData[i].Output = make([]float64, 10, 10)
		for j, pix := range image {
			trainData[i].Input[j] = (float64(pix) / 255)
		}

		for j := range trainData[i].Output {
			trainData[i].Output[j] = 0
		}
		trainData[i].Output[label] = 1
	}
	return trainData
}

func loadTestData() ([]neural.TrainExample, []neural.TrainExample, []neural.TrainExample) {
	train, test, err := GoMNIST.Load("./data")
	if err != nil {
		panic(err)
	}

	tmp := prepareMnistData(train)
	trainData := tmp[:40000]
	validationData := tmp[40000:]
	testData := prepareMnistData(test)
	return trainData, validationData, testData
}

func epocheCallback(nn neural.Evaluator, cost neural.Cost, validationData, testData []neural.TrainExample) neural.EpocheCallback {
	return func(epoche int, dt time.Duration) {
		_, validationErrors := neural.CalculateCorrectness(nn, cost, validationData)
		_, testErrors := neural.CalculateCorrectness(nn, cost, testData)
		if epoche == 1 {
			fmt.Println("epoche,validation error,test error")
		}
		fmt.Printf("%v,%v,%v\n", epoche, validationErrors, testErrors)
	}
}

func main() {
	trainData, validationData, testData := loadTestData()

	activator := neural.NewSigmoidActivator()
	outActivator := neural.NewSoftmaxFunction()
	nn := neural.NewNeuralNetwork(
		[]int{inputSize, 100, 10},
		neural.NewFullyConnectedLayer(activator),
		neural.NewFullyConnectedLayer(outActivator),
	)

	flag.Parse()

	if *nnLoadFile != "" {
		fn, err := os.Open(*nnLoadFile)
		if err != nil {
			log.Fatalln(err)
		}
		if err := neural.Load(nn, fn); err != nil {
			log.Fatalln(err)
		}
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	cost := neural.NewLogLikelihoodCost()
	options := neural.TrainOptions{
		Epochs:         100,
		MiniBatchSize:  10,
		LearningRate:   0.4,
		Regularization: 5,
		TrainerFactory: neural.NewBackpropagationTrainer,
		EpocheCallback: epocheCallback(nn, cost, validationData, testData),
		Cost:           cost,
	}

	t0 := time.Now()
	neural.Train(nn, trainData, options)
	dt := time.Since(t0)

	fmt.Println("Training complete in", dt)

	if *nnSaveFile != "" {
		fn, err := os.OpenFile(*nnSaveFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
		if err != nil {
			log.Fatalln(err)
		}
		if err := neural.Save(nn, fn); err != nil {
			log.Fatalln(err)
		}
	}
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"sync/atomic"
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/security/securitytest"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/util/log"
)

var dbName = flag.String("db-name", "", "Name/URL of the distributed database backend.")
var useTransaction = flag.Bool("use-transaction", true, "Turn off to disable transaction.")

// These two flags configure a range of accounts over which the program functions.
var firstAccount = flag.Int("first-account", 0, "First account in the account range.")
var numAccounts = flag.Int("num-accounts", 1000, "Number of accounts in the account range.")

var numParallelTransfers = flag.Int("num-parallel-transfers", 100, "Number of parallel transfers.")

// Bank stores all the bank related state.
type Bank struct {
	db *client.DB
	// First account in the account range.
	firstAccount int
	// Total number of accounts.
	numAccounts  int
	numTransfers int32
}

// Account holds all the customers account information
type Account struct {
	Balance int64
}

func (a Account) encode() ([]byte, error) {
	return json.Marshal(a)
}

func (a *Account) decode(b []byte) error {
	return json.Unmarshal(b, a)
}

// Makes an id string from an id int.
func (bank *Bank) makeAccountID(id int) []byte {
	return []byte(fmt.Sprintf("%09d", bank.firstAccount+id))
}

// Read the balances in all the accounts and return them.
func (bank *Bank) sumAllAccounts() int64 {
	var result int64
	err := bank.db.Tx(func(tx *client.Tx) error {
		scan, err := tx.Scan(bank.makeAccountID(0), bank.makeAccountID(bank.numAccounts), int64(bank.numAccounts))
		if err != nil {
			return err
		}
		if len(scan.Rows) != bank.numAccounts {
			return fmt.Errorf("Could only read %d of %d rows of the database.\n", len(scan.Rows), bank.numAccounts)
		}
		// Sum up the balances.
		for i := 0; i < bank.numAccounts; i++ {
			account := &Account{}
			err := account.decode(scan.Rows[i].ValueBytes())
			if err != nil {
				return err
			}
			// fmt.Printf("Account %d contains %d$\n", bank.firstAccount+i, account.Balance)
			result += account.Balance
		}
		return nil
	})
	if err != nil {
		log.Fatal(err)
	}
	return result
}

// continuouslyTransferMoney() keeps moving random amounts between
// random accounts.
func (bank *Bank) continuouslyTransferMoney(cash int64) {
	for {
		from := bank.makeAccountID(rand.Intn(bank.numAccounts))
		to := bank.makeAccountID(rand.Intn(bank.numAccounts))
		// Continue when from == to.
		if bytes.Equal(from, to) {
			continue
		}
		exchangeAmount := rand.Int63n(cash)
		// transferMoney transfers exchangeAmount between the two accounts
		transferMoney := func(runner client.Runner) error {
			batchRead := &client.Batch{}
			batchRead.Get(from, to)
			if err := runner.Run(batchRead); err != nil {
				return err
			}
			// Read from value.
			fromAccount := &Account{}
			err := fromAccount.decode(batchRead.Results[0].Rows[0].ValueBytes())
			if err != nil {
				return err
			}
			// Ensure there is enough cash.
			if fromAccount.Balance < exchangeAmount {
				return nil
			}
			// Read to value.
			toAccount := &Account{}
			errRead := toAccount.decode(batchRead.Results[0].Rows[1].ValueBytes())
			if errRead != nil {
				return errRead
			}
			// Update both accounts.
			batchWrite := &client.Batch{}
			fromAccount.Balance -= exchangeAmount
			toAccount.Balance += exchangeAmount
			if fromValue, err := fromAccount.encode(); err != nil {
				return err
			} else if toValue, err := toAccount.encode(); err != nil {
				return err
			} else {
				batchWrite.Put(from, fromValue).Put(to, toValue)
			}
			return runner.Run(batchWrite)
		}
		var err error
		if *useTransaction {
			err = bank.db.Tx(func(tx *client.Tx) error { return transferMoney(tx) })
		} else {
			err = transferMoney(bank.db)
		}
		if err != nil {
			log.Fatal(err)
		}
		atomic.AddInt32(&bank.numTransfers, 1)
	}
}

// Initialize all the bank accounts with cash.  When multiple
// instances of the bank app are running, only one gets to initialize
// the bank.
func (bank *Bank) initBankAccounts(cash int64) {
	if err := bank.db.Tx(func(tx *client.Tx) error {
		// Check if the accounts have been initialized by another instance.
		scan, err := tx.Scan(bank.makeAccountID(0), bank.makeAccountID(bank.numAccounts), int64(bank.numAccounts))
		if err != nil {
			return err
		}
		// Determine existing accounts.
		accts := map[string]struct{}{}
		for i := range scan.Rows {
			accts[string(scan.Rows[i].Key)] = struct{}{}
		}
		// Let's initialize all the accounts.
		batch := &client.Batch{}
		account := Account{Balance: cash}
		value, err := account.encode()
		if err != nil {
			return err
		}
		for i := 0; i < bank.numAccounts; i++ {
			id := bank.makeAccountID(i)
			if _, ok := accts[string(id)]; !ok {
				batch.Put(id, value)
			}
		}
		return tx.Run(batch)
	}); err != nil {
		log.Fatal(err)
	}
	log.Info("done initializing all accounts")
}

func (bank *Bank) periodicallyCheckBalances(initCash int64) {
	var lastNumTransfers int32
	for {
		// Sleep for a bit to allow money transfers to happen in the background.
		time.Sleep(time.Second)
		numTransfers := atomic.LoadInt32(&bank.numTransfers)
		fmt.Printf("%d transfers were executed at %d/second.\n", bank.numTransfers, numTransfers-lastNumTransfers)
		lastNumTransfers = numTransfers
		// Check that all the money is accounted for.
		totalAmount := bank.sumAllAccounts()
		if totalAmount != int64(bank.numAccounts)*initCash {
			log.Fatalf("\nTotal cash in the bank = %d.\n", totalAmount)
		}
		fmt.Printf("The bank is in good order.\n")
	}
}

func main() {
	fmt.Printf("A simple program that keeps moving money between bank accounts.\n\n")
	flag.Parse()
	if *numAccounts < 2 {
		fmt.Fprintf(os.Stderr, "At least two accounts are required to transfer money.\n")
		os.Exit(1)
	}
	if *numParallelTransfers < 1 {
		fmt.Fprintf(os.Stderr, "At least one transfer routine must be active.\n")
		os.Exit(1)
	}
	if !*useTransaction {
		fmt.Printf("Use of transactions has been disabled.\n")
	}
	// Initialize the bank.
	var bank Bank
	bank.firstAccount = *firstAccount
	bank.numAccounts = *numAccounts
	if *dbName == "" {
		// Run a test cockroach instance to represent the bank.
		security.SetReadFileFn(securitytest.Asset)
		serv := server.StartTestServer(nil)
		defer serv.Stop()
		*dbName = "https://root@" + serv.ServingAddr() + "?certs=test_certs"
	}
	// Create a database handle.
	db, err := client.Open(*dbName)
	if err != nil {
		log.Fatal(err)
	}
	bank.db = db
	// Initialize all the bank accounts.
	const initCash = 1000
	bank.initBankAccounts(initCash)

	// Start all the money transfer routines.
	for i := 0; i < *numParallelTransfers; i++ {
		// Keep transferring upto 10% of initCash between accounts.
		go bank.continuouslyTransferMoney(initCash / 10)
	}

	bank.periodicallyCheckBalances(initCash)
}
<commit_msg>further improvements, in particular to determine actual initial amount of cash between accounts.<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Vivek Menezes (vivek.menezes@gmail.com)

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"sync/atomic"
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/security/securitytest"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/util/log"
)

var dbName = flag.String("db-name", "", "Name/URL of the distributed database backend.")
var useTransaction = flag.Bool("use-transaction", true, "Turn off to disable transaction.")

// These two flags configure a range of accounts over which the program functions.
var firstAccount = flag.Int("first-account", 0, "First account in the account range.")
var numAccounts = flag.Int("num-accounts", 1000, "Number of accounts in the account range.")

var numParallelTransfers = flag.Int("num-parallel-transfers", 100, "Number of parallel transfers.")

// Bank stores all the bank related state.
type Bank struct {
	db *client.DB
	// First account in the account range.
	firstAccount int
	// Total number of accounts.
	numAccounts  int
	numTransfers int32
}

// Account holds all the customers account information
type Account struct {
	Balance int64
}

func (a Account) encode() ([]byte, error) {
	return json.Marshal(a)
}

func (a *Account) decode(b []byte) error {
	return json.Unmarshal(b, a)
}

// Makes an id string from an id int.
func (bank *Bank) makeAccountID(id int) []byte {
	return []byte(fmt.Sprintf("%09d", bank.firstAccount+id))
}

// Read the balances in all the accounts and return them.
func (bank *Bank) sumAllAccounts() int64 {
	var result int64
	err := bank.db.Tx(func(tx *client.Tx) error {
		scan, err := tx.Scan(bank.makeAccountID(0), bank.makeAccountID(bank.numAccounts), int64(bank.numAccounts))
		if err != nil {
			return err
		}
		if len(scan.Rows) != bank.numAccounts {
			return fmt.Errorf("Could only read %d of %d rows of the database.\n", len(scan.Rows), bank.numAccounts)
		}
		// Sum up the balances.
		for i := 0; i < bank.numAccounts; i++ {
			account := &Account{}
			err := account.decode(scan.Rows[i].ValueBytes())
			if err != nil {
				return err
			}
			// fmt.Printf("Account %d contains %d$\n", bank.firstAccount+i, account.Balance)
			result += account.Balance
		}
		return nil
	})
	if err != nil {
		log.Fatal(err)
	}
	return result
}

// continuouslyTransferMoney() keeps moving random amounts between
// random accounts.
func (bank *Bank) continuouslyTransferMoney(cash int64) {
	for {
		from := bank.makeAccountID(rand.Intn(bank.numAccounts))
		to := bank.makeAccountID(rand.Intn(bank.numAccounts))
		// Continue when from == to.
		if bytes.Equal(from, to) {
			continue
		}
		exchangeAmount := rand.Int63n(cash)
		// transferMoney transfers exchangeAmount between the two accounts
		transferMoney := func(runner client.Runner) error {
			batchRead := &client.Batch{}
			batchRead.Get(from, to)
			if err := runner.Run(batchRead); err != nil {
				return err
			}
			// Read from value.
			fromAccount := &Account{}
			err := fromAccount.decode(batchRead.Results[0].Rows[0].ValueBytes())
			if err != nil {
				return err
			}
			// Ensure there is enough cash.
			if fromAccount.Balance < exchangeAmount {
				return nil
			}
			// Read to value.
			toAccount := &Account{}
			errRead := toAccount.decode(batchRead.Results[0].Rows[1].ValueBytes())
			if errRead != nil {
				return errRead
			}
			// Update both accounts.
			batchWrite := &client.Batch{}
			fromAccount.Balance -= exchangeAmount
			toAccount.Balance += exchangeAmount
			if fromValue, err := fromAccount.encode(); err != nil {
				return err
			} else if toValue, err := toAccount.encode(); err != nil {
				return err
			} else {
				batchWrite.Put(from, fromValue).Put(to, toValue)
			}
			return runner.Run(batchWrite)
		}
		var err error
		if *useTransaction {
			err = bank.db.Tx(func(tx *client.Tx) error { return transferMoney(tx) })
		} else {
			err = transferMoney(bank.db)
		}
		if err != nil {
			log.Fatal(err)
		}
		atomic.AddInt32(&bank.numTransfers, 1)
	}
}

// Initialize all the bank accounts with cash.  When multiple
// instances of the bank app are running, only one gets to initialize
// the bank.
func (bank *Bank) initBankAccounts(cash int64) int64 {
	var initCash int64
	if err := bank.db.Tx(func(tx *client.Tx) error {
		// Check if the accounts have been initialized by another instance.
		scan, err := tx.Scan(bank.makeAccountID(0), bank.makeAccountID(bank.numAccounts), int64(bank.numAccounts))
		if err != nil {
			return err
		}
		// Determine existing accounts.
		existAcct := &Account{}
		accts := map[string]struct{}{}
		for i := range scan.Rows {
			accts[string(scan.Rows[i].Key)] = struct{}{}
			if err := existAcct.decode(scan.Rows[i].ValueBytes()); err != nil {
				log.Fatal("error decoding existing account %s: %s", scan.Rows[i].Key, err)
			}
			initCash += existAcct.Balance
		}
		// Let's initialize all the accounts.
		batch := &client.Batch{}
		account := Account{Balance: cash}
		value, err := account.encode()
		if err != nil {
			return err
		}
		for i := 0; i < bank.numAccounts; i++ {
			id := bank.makeAccountID(i)
			if _, ok := accts[string(id)]; !ok {
				batch.Put(id, value)
				initCash += cash
			}
		}
		return tx.Run(batch)
	}); err != nil {
		log.Fatal(err)
	}
	log.Info("done initializing all accounts")
	return initCash
}

func (bank *Bank) periodicallyCheckBalances(initCash int64) {
	var lastNumTransfers int32
	lastNow := time.Now()
	for {
		// Sleep for a bit to allow money transfers to happen in the background.
		time.Sleep(time.Second)
		now := time.Now()
		elapsed := now.Sub(lastNow)
		numTransfers := atomic.LoadInt32(&bank.numTransfers)
		fmt.Printf("%d transfers were executed at %.1f/second.\n", bank.numTransfers,
			float64(numTransfers-lastNumTransfers)/(float64(elapsed.Nanoseconds())/1E9))
		lastNumTransfers = numTransfers
		lastNow = now
		// Check that all the money is accounted for.
		totalAmount := bank.sumAllAccounts()
		if totalAmount != initCash {
			log.Fatalf("\nTotal cash in the bank $%d; expected $%d.\n", totalAmount, initCash)
		}
		fmt.Printf("The bank is in good order.\n")
	}
}

func main() {
	fmt.Printf("A simple program that keeps moving money between bank accounts.\n\n")
	flag.Parse()
	if *numAccounts < 2 {
		fmt.Fprintf(os.Stderr, "At least two accounts are required to transfer money.\n")
		os.Exit(1)
	}
	if *numParallelTransfers < 1 {
		fmt.Fprintf(os.Stderr, "At least one transfer routine must be active.\n")
		os.Exit(1)
	}
	if !*useTransaction {
		fmt.Printf("Use of transactions has been disabled.\n")
	}
	// Initialize the bank.
	var bank Bank
	bank.firstAccount = *firstAccount
	bank.numAccounts = *numAccounts
	if *dbName == "" {
		// Run a test cockroach instance to represent the bank.
		security.SetReadFileFn(securitytest.Asset)
		serv := server.StartTestServer(nil)
		defer serv.Stop()
		*dbName = "https://root@" + serv.ServingAddr() + "?certs=test_certs"
	}
	// Create a database handle.
	db, err := client.Open(*dbName)
	if err != nil {
		log.Fatal(err)
	}
	bank.db = db
	// Initialize all the bank accounts.
	initCash := bank.initBankAccounts(1000)

	// Start all the money transfer routines.
	for i := 0; i < *numParallelTransfers; i++ {
		// Keep transferring upto 10% of initCash between accounts.
		go bank.continuouslyTransferMoney(initCash / 10)
	}

	bank.periodicallyCheckBalances(initCash)
}
<|endoftext|>


<commit_before>//
//  Clone server Model Five
//

package main

import (
	zmq "github.com/pebbe/zmq3"
	"github.com/pebbe/zmq3/examples/kvmsg"

	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"
)

//  Our server is defined by these properties
type clonesrv_t struct {
	kvmap     map[string]*kvmsg.Kvmsg //  Key-value store
	port      int                     //  Main port we're working on
	sequence  int64                   //  How many updates we're at
	snapshot  *zmq.Socket             //  Handle snapshot requests
	publisher *zmq.Socket             //  Publish updates to clients
	collector *zmq.Socket             //  Collect updates from clients
}

func main() {

	srv := &clonesrv_t{
		port:  5556,
		kvmap: make(map[string]*kvmsg.Kvmsg),
	}

	//  Set up our clone server sockets
	srv.snapshot, _ = zmq.NewSocket(zmq.ROUTER)
	srv.snapshot.Bind(fmt.Sprint("tcp://*:", srv.port))
	srv.publisher, _ = zmq.NewSocket(zmq.PUB)
	srv.publisher.Bind(fmt.Sprint("tcp://*:", srv.port+1))
	srv.collector, _ = zmq.NewSocket(zmq.PULL)
	srv.collector.Bind(fmt.Sprint("tcp://*:", srv.port+2))

	//  Register our handlers with reactor
	reactor := zmq.NewReactor()
	reactor.AddSocket(srv.snapshot, zmq.POLLIN,
		func(e zmq.State) error { return snapshots(srv) })
	reactor.AddSocket(srv.collector, zmq.POLLIN,
		func(e zmq.State) error { return collector(srv) })
	reactor.AddChannelTime(time.Tick(1000*time.Millisecond), 1,
		func(v interface{}) error { return flush_ttl(srv) })

	log.Println(reactor.Run(1000 * time.Millisecond))
}

//  This is the reactor handler for the snapshot socket; it accepts
//  just the ICANHAZ? request and replies with a state snapshot ending
//  with a KTHXBAI message:

func snapshots(srv *clonesrv_t) (err error) {

	msg, err := srv.snapshot.RecvMessage(0)
	if err != nil {
		return
	}
	identity := msg[0]

	//  Request is in second frame of message
	request := msg[1]
	if request != "ICANHAZ?" {
		err = errors.New("E: bad request, aborting")
		return
	}
	subtree := msg[2]

	//  Send state socket to client
	for _, kvmsg := range srv.kvmap {
		if key, _ := kvmsg.GetKey(); strings.HasPrefix(key, subtree) {
			srv.snapshot.Send(identity, zmq.SNDMORE)
			kvmsg.Send(srv.snapshot)
		}
	}

	//  Now send END message with sequence number
	log.Println("I: sending shapshot =", srv.sequence)
	srv.snapshot.Send(identity, zmq.SNDMORE)
	kvmsg := kvmsg.NewKvmsg(srv.sequence)
	kvmsg.SetKey("KTHXBAI")
	kvmsg.SetBody(subtree)
	kvmsg.Send(srv.snapshot)

	return
}

//  We store each update with a new sequence number, and if necessary, a
//  time-to-live. We publish updates immediately on our publisher socket:

func collector(srv *clonesrv_t) (err error) {
	kvmsg, err := kvmsg.RecvKvmsg(srv.collector)
	if err != nil {
		return
	}

	srv.sequence++
	kvmsg.SetSequence(srv.sequence)
	kvmsg.Send(srv.publisher)
	if ttls, e := kvmsg.GetProp("ttl"); e == nil {
		// change duration into specific time, using the same property: ugly!
		ttl, e := strconv.ParseInt(ttls, 10, 64)
		if e != nil {
			err = e
			return
		}
		kvmsg.SetProp("ttl", fmt.Sprint(time.Now().Add(time.Duration(ttl)*time.Second).Unix()))
	}
	kvmsg.Store(srv.kvmap)
	log.Println("I: publishing update =", srv.sequence)

	return
}

//  At regular intervals we flush ephemeral values that have expired. This
//  could be slow on very large data sets:

func flush_ttl(srv *clonesrv_t) (err error) {

	for _, kvmsg := range srv.kvmap {

		//  If key-value pair has expired, delete it and publish the
		//  fact to listening clients.

		if ttls, e := kvmsg.GetProp("ttl"); e == nil {
			ttl, e := strconv.ParseInt(ttls, 10, 64)
			if e != nil {
				err = e
				continue
			}
			if time.Now().After(time.Unix(ttl, 0)) {
				srv.sequence++
				kvmsg.SetSequence(srv.sequence)
				kvmsg.SetBody("")
				e = kvmsg.Send(srv.publisher)
				if e != nil {
					err = e
				}
				kvmsg.Store(srv.kvmap)
				log.Println("I: publishing delete =", srv.sequence)
			}
		}
	}
	return
}
<commit_msg>Example clonesrv5: minor edit<commit_after>//
//  Clone server Model Five
//

package main

import (
	zmq "github.com/pebbe/zmq3"
	"github.com/pebbe/zmq3/examples/kvmsg"

	"errors"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"
)

//  Our server is defined by these properties
type clonesrv_t struct {
	kvmap     map[string]*kvmsg.Kvmsg //  Key-value store
	port      int                     //  Main port we're working on
	sequence  int64                   //  How many updates we're at
	snapshot  *zmq.Socket             //  Handle snapshot requests
	publisher *zmq.Socket             //  Publish updates to clients
	collector *zmq.Socket             //  Collect updates from clients
}

func main() {

	srv := &clonesrv_t{
		port:  5556,
		kvmap: make(map[string]*kvmsg.Kvmsg),
	}

	//  Set up our clone server sockets
	srv.snapshot, _ = zmq.NewSocket(zmq.ROUTER)
	srv.snapshot.Bind(fmt.Sprint("tcp://*:", srv.port))
	srv.publisher, _ = zmq.NewSocket(zmq.PUB)
	srv.publisher.Bind(fmt.Sprint("tcp://*:", srv.port+1))
	srv.collector, _ = zmq.NewSocket(zmq.PULL)
	srv.collector.Bind(fmt.Sprint("tcp://*:", srv.port+2))

	//  Register our handlers with reactor
	reactor := zmq.NewReactor()
	reactor.AddSocket(srv.snapshot, zmq.POLLIN,
		func(e zmq.State) error { return snapshots(srv) })
	reactor.AddSocket(srv.collector, zmq.POLLIN,
		func(e zmq.State) error { return collector(srv) })
	reactor.AddChannelTime(time.Tick(1000*time.Millisecond), 1,
		func(v interface{}) error { return flush_ttl(srv) })

	log.Println(reactor.Run(100 * time.Millisecond)) // precision: .1 seconds
}

//  This is the reactor handler for the snapshot socket; it accepts
//  just the ICANHAZ? request and replies with a state snapshot ending
//  with a KTHXBAI message:

func snapshots(srv *clonesrv_t) (err error) {

	msg, err := srv.snapshot.RecvMessage(0)
	if err != nil {
		return
	}
	identity := msg[0]

	//  Request is in second frame of message
	request := msg[1]
	if request != "ICANHAZ?" {
		err = errors.New("E: bad request, aborting")
		return
	}
	subtree := msg[2]

	//  Send state socket to client
	for _, kvmsg := range srv.kvmap {
		if key, _ := kvmsg.GetKey(); strings.HasPrefix(key, subtree) {
			srv.snapshot.Send(identity, zmq.SNDMORE)
			kvmsg.Send(srv.snapshot)
		}
	}

	//  Now send END message with sequence number
	log.Println("I: sending shapshot =", srv.sequence)
	srv.snapshot.Send(identity, zmq.SNDMORE)
	kvmsg := kvmsg.NewKvmsg(srv.sequence)
	kvmsg.SetKey("KTHXBAI")
	kvmsg.SetBody(subtree)
	kvmsg.Send(srv.snapshot)

	return
}

//  We store each update with a new sequence number, and if necessary, a
//  time-to-live. We publish updates immediately on our publisher socket:

func collector(srv *clonesrv_t) (err error) {
	kvmsg, err := kvmsg.RecvKvmsg(srv.collector)
	if err != nil {
		return
	}

	srv.sequence++
	kvmsg.SetSequence(srv.sequence)
	kvmsg.Send(srv.publisher)
	if ttls, e := kvmsg.GetProp("ttl"); e == nil {
		// change duration into specific time, using the same property: ugly!
		ttl, e := strconv.ParseInt(ttls, 10, 64)
		if e != nil {
			err = e
			return
		}
		kvmsg.SetProp("ttl", fmt.Sprint(time.Now().Add(time.Duration(ttl)*time.Second).Unix()))
	}
	kvmsg.Store(srv.kvmap)
	log.Println("I: publishing update =", srv.sequence)

	return
}

//  At regular intervals we flush ephemeral values that have expired. This
//  could be slow on very large data sets:

func flush_ttl(srv *clonesrv_t) (err error) {

	for _, kvmsg := range srv.kvmap {

		//  If key-value pair has expired, delete it and publish the
		//  fact to listening clients.

		if ttls, e := kvmsg.GetProp("ttl"); e == nil {
			ttl, e := strconv.ParseInt(ttls, 10, 64)
			if e != nil {
				err = e
				continue
			}
			if time.Now().After(time.Unix(ttl, 0)) {
				srv.sequence++
				kvmsg.SetSequence(srv.sequence)
				kvmsg.SetBody("")
				e = kvmsg.Send(srv.publisher)
				if e != nil {
					err = e
				}
				kvmsg.Store(srv.kvmap)
				log.Println("I: publishing delete =", srv.sequence)
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

import "testing"
import "runtime"
import "time"

func TestMessageInsertAndRetreive(t *testing.T) {
	say := "'Ello, Mister Polly Parrot!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	store.Add <- &Message{at, say}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.Fail()
	}
	if messages[0].Time != at {
		t.Fail()
	}
	if messages[0].Text != say {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}

func TestFetchBlocksUntilSpeak(t *testing.T) {
	start_fetch_wait_count := fetch_wait_count.String()
	say := "I've got a lovely fresh cuttle fish for you"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	for start_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}
	store.Add <- &Message{at, say}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.Fail()
	}
	if messages[0].Time != at {
		t.Fail()
	}
	if messages[0].Text != say {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}

func TestMultipleListeners(t *testing.T) {
	say := "This is your nine o'clock alarm call!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	const num_clients = 13
	var messages_from_store [num_clients]chan []Message
	for i := 0; i < num_clients; i++ {
		messages_from_store[i] = make(chan []Message, 1)
		store.Get <- &StoreRequest{zero_time, messages_from_store[i]}
	}
	store.Add <- &Message{at, say}
	for i := 0; i < num_clients; i++ {
		messages := <-messages_from_store[i]
		if len(messages) != 1 {
			t.Fail()
		}
		if messages[0].Time != at {
			t.Fail()
		}
		if messages[0].Text != say {
			t.Fail()
		}
	}
	close(store.Get)
	close(store.Add)
}
<commit_msg>Test partial fetch<commit_after>package main

import "testing"
import "runtime"
import "strconv"
import "time"

func TestMessageInsertAndRetreive(t *testing.T) {
	say := "'Ello, Mister Polly Parrot!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	store.Add <- &Message{at, say}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	if messages[0].Time != at {
		t.Fail()
	}
	if messages[0].Text != say {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}

func TestFetchBlocksUntilSpeak(t *testing.T) {
	start_fetch_wait_count := fetch_wait_count.String()
	say := "I've got a lovely fresh cuttle fish for you"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	for start_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}
	store.Add <- &Message{at, say}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	if messages[0].Time != at {
		t.Fail()
	}
	if messages[0].Text != say {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}

func TestMultipleListeners(t *testing.T) {
	say := "This is your nine o'clock alarm call!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	const num_clients = 13
	var messages_from_store [num_clients]chan []Message
	for i := 0; i < num_clients; i++ {
		messages_from_store[i] = make(chan []Message, 1)
		store.Get <- &StoreRequest{zero_time, messages_from_store[i]}
	}
	store.Add <- &Message{at, say}
	for i := 0; i < num_clients; i++ {
		messages := <-messages_from_store[i]
		if len(messages) != 1 {
			t.FailNow()
		}
		if messages[0].Time != at {
			t.Fail()
		}
		if messages[0].Text != say {
			t.Fail()
		}
	}
	close(store.Get)
	close(store.Add)
}

func parseDuration(s string) time.Duration {
	d, err := time.ParseDuration(s)
	if err != nil {
		panic(err)
	}
	return d
}

func atoi(s string) int {
	i, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return i
}

func TestPartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	say1 := "No, no.....No, 'e's stunned!"
	say2 := "You stunned him, just as he was wakin' up!"
	say3 := "Norwegian Blues stun easily, major."
	base := time.Now()
	at1 := base.Add(parseDuration("-4m"))
	since := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	store := start_store()
	store.Add <- &Message{at1, say1}
	store.Add <- &Message{at2, say2}
	store.Add <- &Message{at3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 2 {
		t.FailNow()
	}
	if messages[0].Time != at2 {
		t.Fail()
	}
	if messages[0].Text != say2 {
		t.Fail()
	}
	if messages[1].Time != at3 {
		t.Fail()
	}
	if messages[1].Text != say3 {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}
<|endoftext|>


<commit_before>/*  reliable-chat - multipath chat
 *  Copyright (C) 2012  Scott Worley <sworley@chkno.net>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as
 *  published by the Free Software Foundation, either version 3 of the
 *  License, or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package main

import "testing"
import "runtime"
import "strconv"
import "time"

func expectMessage(t *testing.T, m *Message, at time.Time, id, say string) {
	if m.Time != at {
		t.Fail()
	}
	if m.ID != id {
		t.Fail()
	}
	if m.Text != say {
		t.Fail()
	}
}

func TestMessageInsertAndRetreive(t *testing.T) {
	id := "1"
	say := "'Ello, Mister Polly Parrot!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	store.Add <- &Message{at, id, say}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at, id, say)
}

func TestFetchBlocksUntilSpeak(t *testing.T) {
	start_fetch_wait_count := fetch_wait_count.String()
	id := "2"
	say := "I've got a lovely fresh cuttle fish for you"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	for start_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}
	store.Add <- &Message{at, id, say}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at, id, say)
}

func TestMultipleListeners(t *testing.T) {
	id := "3"
	say := "This is your nine o'clock alarm call!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	const num_clients = 13
	var messages_from_store [num_clients]chan []Message
	for i := 0; i < num_clients; i++ {
		messages_from_store[i] = make(chan []Message, 1)
		store.Get <- &StoreRequest{zero_time, messages_from_store[i]}
	}
	store.Add <- &Message{at, id, say}
	for i := 0; i < num_clients; i++ {
		messages := <-messages_from_store[i]
		if len(messages) != 1 {
			t.FailNow()
		}
		expectMessage(t, &messages[0], at, id, say)
	}
}

func parseDuration(s string) time.Duration {
	d, err := time.ParseDuration(s)
	if err != nil {
		panic(err)
	}
	return d
}

func atoi(s string) int {
	i, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return i
}

func TestPartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	id1 := "4"
	id2 := "5"
	id3 := "6"
	say1 := "No, no.....No, 'e's stunned!"
	say2 := "You stunned him, just as he was wakin' up!"
	say3 := "Norwegian Blues stun easily, major."
	base := time.Now()
	at1 := base.Add(parseDuration("-4m"))
	since := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 2 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at2, id2, say2)
	expectMessage(t, &messages[1], at3, id3, say3)
}

func TestPrecisePartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	id1 := "7"
	id2 := "8"
	id3 := "9"
	say1 := "Well, he's...he's, ah...probably pining for the fjords."
	say2 := "PININ' for the FJORDS?!?!?!?"
	say3 := "look, why did he fall flat on his back the moment I got 'im home?"
	base := time.Now()
	at1 := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	since := at2
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at3, id3, say3)
}

func TestTypicalFlow(t *testing.T) {
	id1 := "10"
	id2 := "11"
	say1 := "The Norwegian Blue prefers kippin' on it's back!"
	say2 := "Remarkable bird, innit, squire?  Lovely plumage!"
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	// A waiting zero-time fetch.
	var zero_time time.Time
	prev_fetch_wait_count := fetch_wait_count.String()
	fetch1 := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, fetch1}
	for prev_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}

	// Someone speaks.  This triggers delivery.
	at1 := time.Now()
	store.Add <- &Message{at1, id1, say1}
	messages1 := <-fetch1
	if len(messages1) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages1[0], at1, id1, say1)

	// Upon recipt, client blocks on fetch with since=at1
	prev_fetch_wait_count = fetch_wait_count.String()
	fetch2 := make(chan []Message, 1)
	store.Get <- &StoreRequest{at1, fetch2}
	for prev_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}

	// Someone speaks again.  This triggers another delivery.
	at2 := time.Now()
	if !at2.After(at1) {
		t.Fail()
	}
	store.Add <- &Message{at2, id2, say2}
	messages2 := <-fetch2
	if len(messages2) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages2[0], at2, id2, say2)
}
<commit_msg>Test message drops due to hitting the limit<commit_after>/*  reliable-chat - multipath chat
 *  Copyright (C) 2012  Scott Worley <sworley@chkno.net>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as
 *  published by the Free Software Foundation, either version 3 of the
 *  License, or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package main

import "testing"
import "runtime"
import "strconv"
import "time"

func expectMessage(t *testing.T, m *Message, at time.Time, id, say string) {
	if m.Time != at {
		t.Fail()
	}
	if m.ID != id {
		t.Fail()
	}
	if m.Text != say {
		t.Fail()
	}
}

func TestMessageInsertAndRetreive(t *testing.T) {
	id := "1"
	say := "'Ello, Mister Polly Parrot!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	store.Add <- &Message{at, id, say}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at, id, say)
}

func TestFetchBlocksUntilSpeak(t *testing.T) {
	start_fetch_wait_count := fetch_wait_count.String()
	id := "2"
	say := "I've got a lovely fresh cuttle fish for you"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	for start_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}
	store.Add <- &Message{at, id, say}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at, id, say)
}

func TestMultipleListeners(t *testing.T) {
	id := "3"
	say := "This is your nine o'clock alarm call!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	const num_clients = 13
	var messages_from_store [num_clients]chan []Message
	for i := 0; i < num_clients; i++ {
		messages_from_store[i] = make(chan []Message, 1)
		store.Get <- &StoreRequest{zero_time, messages_from_store[i]}
	}
	store.Add <- &Message{at, id, say}
	for i := 0; i < num_clients; i++ {
		messages := <-messages_from_store[i]
		if len(messages) != 1 {
			t.FailNow()
		}
		expectMessage(t, &messages[0], at, id, say)
	}
}

func parseDuration(s string) time.Duration {
	d, err := time.ParseDuration(s)
	if err != nil {
		panic(err)
	}
	return d
}

func atoi(s string) int {
	i, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return i
}

func TestPartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	id1 := "4"
	id2 := "5"
	id3 := "6"
	say1 := "No, no.....No, 'e's stunned!"
	say2 := "You stunned him, just as he was wakin' up!"
	say3 := "Norwegian Blues stun easily, major."
	base := time.Now()
	at1 := base.Add(parseDuration("-4m"))
	since := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 2 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at2, id2, say2)
	expectMessage(t, &messages[1], at3, id3, say3)
}

func TestPrecisePartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	id1 := "7"
	id2 := "8"
	id3 := "9"
	say1 := "Well, he's...he's, ah...probably pining for the fjords."
	say2 := "PININ' for the FJORDS?!?!?!?"
	say3 := "look, why did he fall flat on his back the moment I got 'im home?"
	base := time.Now()
	at1 := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	since := at2
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at3, id3, say3)
}

func TestTypicalFlow(t *testing.T) {
	id1 := "10"
	id2 := "11"
	say1 := "The Norwegian Blue prefers kippin' on it's back!"
	say2 := "Remarkable bird, innit, squire?  Lovely plumage!"
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	// A waiting zero-time fetch.
	var zero_time time.Time
	prev_fetch_wait_count := fetch_wait_count.String()
	fetch1 := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, fetch1}
	for prev_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}

	// Someone speaks.  This triggers delivery.
	at1 := time.Now()
	store.Add <- &Message{at1, id1, say1}
	messages1 := <-fetch1
	if len(messages1) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages1[0], at1, id1, say1)

	// Upon recipt, client blocks on fetch with since=at1
	prev_fetch_wait_count = fetch_wait_count.String()
	fetch2 := make(chan []Message, 1)
	store.Get <- &StoreRequest{at1, fetch2}
	for prev_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}

	// Someone speaks again.  This triggers another delivery.
	at2 := time.Now()
	if !at2.After(at1) {
		t.Fail()
	}
	store.Add <- &Message{at2, id2, say2}
	messages2 := <-fetch2
	if len(messages2) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages2[0], at2, id2, say2)
}

func TestExpiryDueToLimit(t *testing.T) {
	previous_limit := *max_messages
	defer func() { *max_messages = previous_limit }()
	*max_messages = 2

	start_speak_count := atoi(speak_count.String())
	start_drop_count := atoi(drop_due_to_limit_count.String())
	id1 := "12"
	id2 := "13"
	id3 := "14"
	say1 := "'E's passed on!"
	say2 := "This parrot is no more!"
	say3 := "He has ceased to be!"
	base := time.Now()
	at1 := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	store := start_store()
	defer func() {
		close(store.Get)
		close(store.Add)
	}()

	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	if atoi(drop_due_to_limit_count.String()) != start_drop_count+1 {
		t.Fail()
	}
	messages_from_store := make(chan []Message, 1)
	var zero_time time.Time
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 2 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at2, id2, say2)
	expectMessage(t, &messages[1], at3, id3, say3)
}
<|endoftext|>


<commit_before>package service

import (
	"strconv"
	"strings"

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

var holder *PrometheusHolder

// PrometheusHolder is holds an instances of globally and internally used Prometheus metrics
type PrometheusHolder struct {
	httpLatency *prometheus.SummaryVec
}

// NewPrometheus returns a new instance of a Prometheus holder to bind to a service
func NewPrometheus(serviceName string) *PrometheusHolder {
	if holder != nil {
		return holder
	}

	prefix := strings.Replace(serviceName, "-", "_", -1)

	httpLatency := prometheus.NewSummaryVec(prometheus.SummaryOpts{
		Name: prefix + "_http_handler_latency",
		Help: "HTTP Handler Latency by endpoint",
	},
		[]string{"method", "path", "status_code"},
	)

	prometheus.MustRegister(httpLatency)

	holder = &PrometheusHolder{
		httpLatency: httpLatency,
	}
	return holder
}

// OnRequestFinished will measure the http latency for the respective call to a missy service endpoint
func (p *PrometheusHolder) OnRequestFinished(method string, path string, statusCode int, processTimeMillis float64) {
	p.httpLatency.WithLabelValues(method, path, strconv.Itoa(statusCode)).Observe(processTimeMillis)
}
<commit_msg>refactor(metrics): Use Histogram metric instead of Summary (#26)<commit_after>package service

import (
	"strconv"
	"strings"

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

var holder *PrometheusHolder

// PrometheusHolder is holds an instances of globally and internally used Prometheus metrics
type PrometheusHolder struct {
	httpLatency *prometheus.HistogramVec
}

// NewPrometheus returns a new instance of a Prometheus holder to bind to a service
func NewPrometheus(serviceName string) *PrometheusHolder {
	if holder != nil {
		return holder
	}

	prefix := strings.Replace(serviceName, "-", "_", -1)

	httpLatency := prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Name: prefix + "_http_handler_latency",
		Help: "HTTP Handler Latency by endpoint",
	},
		[]string{"method", "path", "status_code"},
	)

	prometheus.MustRegister(httpLatency)

	holder = &PrometheusHolder{
		httpLatency: httpLatency,
	}
	return holder
}

// OnRequestFinished will measure the http latency for the respective call to a missy service endpoint
func (p *PrometheusHolder) OnRequestFinished(method string, path string, statusCode int, processTimeMillis float64) {
	p.httpLatency.WithLabelValues(method, path, strconv.Itoa(statusCode)).Observe(processTimeMillis)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
	"net/url"
	"strings"
	"time"
)

type AboutNode struct {
	Addr     string `json:"addr"`
	Type     string `json:"type"`
	Time     string `json:"time"`
	BindAddr string `json:"bindaddr"`
	Hash     string `json:"hash"`
}

func getNodeAddress(sid string) (string, error) {
	sidkey := "/" + sid
	aboutSid := AboutNode{}
	err := couchbase.Get(sidkey, &aboutSid)
	if err != nil {
        return "", err
	}
	if strings.HasPrefix(aboutSid.BindAddr, ":") {
		return aboutSid.Addr + aboutSid.BindAddr, nil
	}
	return aboutSid.BindAddr, nil
}

func heartbeat() {
	for {

		u, err := url.Parse(*couchbaseServer)
		c, err := net.Dial("tcp", u.Host)
		localAddr := ""
		if err == nil {
			localAddr = strings.Split(c.LocalAddr().String(), ":")[0]
			c.Close()
		}

		aboutMe := AboutNode{}
		aboutMe.Addr = localAddr
		aboutMe.Type = "storage"
		aboutMe.Time = time.Now().UTC().String()
		aboutMe.BindAddr = *bindAddr
		aboutMe.Hash = *hashType

		err = couchbase.Set("/"+serverIdentifier(), aboutMe)
		if err != nil {
			log.Printf("Failed to record a heartbeat: %v", err)
		}
		time.Sleep(5 * time.Second)
	}
}
<commit_msg>Configurable heartbeat frequency.<commit_after>package main

import (
	"flag"
	"log"
	"net"
	"net/url"
	"strings"
	"time"
)

var heartFreq = flag.Duration("heartbeat", 10*time.Second,
	"Heartbeat frequency")

type AboutNode struct {
	Addr     string `json:"addr"`
	Type     string `json:"type"`
	Time     string `json:"time"`
	BindAddr string `json:"bindaddr"`
	Hash     string `json:"hash"`
}

func getNodeAddress(sid string) (string, error) {
	sidkey := "/" + sid
	aboutSid := AboutNode{}
	err := couchbase.Get(sidkey, &aboutSid)
	if err != nil {
		return "", err
	}
	if strings.HasPrefix(aboutSid.BindAddr, ":") {
		return aboutSid.Addr + aboutSid.BindAddr, nil
	}
	return aboutSid.BindAddr, nil
}

func heartbeat() {
	for {

		u, err := url.Parse(*couchbaseServer)
		c, err := net.Dial("tcp", u.Host)
		localAddr := ""
		if err == nil {
			localAddr = strings.Split(c.LocalAddr().String(), ":")[0]
			c.Close()
		}

		aboutMe := AboutNode{}
		aboutMe.Addr = localAddr
		aboutMe.Type = "storage"
		aboutMe.Time = time.Now().UTC().String()
		aboutMe.BindAddr = *bindAddr
		aboutMe.Hash = *hashType

		err = couchbase.Set("/"+serverIdentifier(), aboutMe)
		if err != nil {
			log.Printf("Failed to record a heartbeat: %v", err)
		}
		time.Sleep(*heartFreq)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/catkins/heartbeat/config"
	"gopkg.in/redis.v2"
)

var appConfig config.Configuration

func main() {
	appConfig = config.Load()

	interval := time.Duration(appConfig.HeartbeatInterval) * time.Second
	ticker := time.NewTicker(interval)

	fmt.Printf("Connecting to redis at redis://%s/%d\n",
		appConfig.RedisAddress, appConfig.RedisDatabase)

	options := appConfig.RedisOptions()
	client := redis.NewTCPClient(&options)
	defer client.Close()

	fmt.Printf("Starting heartbeat on channel \"%s\" every %d seconds\n",
		appConfig.HeartbeatChannel, appConfig.HeartbeatInterval)

	for {
		tick := <-ticker.C

		go func() {
			var message string
			if len(appConfig.HeartbeatMessage) > 0 {
				message = appConfig.HeartbeatMessage
			} else {
				message = fmt.Sprintf("%d", tick.Unix())
			}

			_, err := client.Publish(appConfig.HeartbeatChannel, message).Result()

			if err != nil {
				fmt.Println(time.Now().String(), err.Error())
			}
		}()

	}
}
<commit_msg>refactored main package<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/catkins/heartbeat/config"
	"gopkg.in/redis.v2"
)

var appConfig config.Configuration

func init() {
	appConfig = config.Load()
}

func main() {
	client := connectToRedis()
	defer client.Close()

	startInterval(func(tick time.Time) {
		message := buildMessage(tick)
		publish(client, message)
	})
}

func connectToRedis() *redis.Client {
	fmt.Printf("Connecting to redis at redis://%s/%d\n",
		appConfig.RedisAddress,
		appConfig.RedisDatabase)

	options := appConfig.RedisOptions()
	return redis.NewTCPClient(&options)
}

func startInterval(callback func(time.Time)) {
	fmt.Printf("Starting heartbeat on channel \"%s\" every %d seconds\n",
		appConfig.HeartbeatChannel,
		appConfig.HeartbeatInterval)

	interval := time.Duration(appConfig.HeartbeatInterval) * time.Second
	ticker := time.NewTicker(interval)

	for {
		tick := <-ticker.C
		go callback(tick)
	}
}

func publish(client *redis.Client, message string) {
	_, err := client.Publish(appConfig.HeartbeatChannel, message).Result()

	if err != nil {
		fmt.Println(time.Now().String(), err.Error())
	}
}

func buildMessage(tick time.Time) string {
	if len(appConfig.HeartbeatMessage) > 0 {
		return appConfig.HeartbeatMessage
	}

	return fmt.Sprintf("%d", tick.Unix())
}
<|endoftext|>


<commit_before>package color

import (
	"bytes"
	"os"
	"strconv"
	"strings"
	"sync"
	"unicode"

	"github.com/gdamore/tcell"
)

const (
	errMissing = "%%!h(MISSING)" // no attributes in the highlight verb
	errInvalid = "%%!h(INVALID)" // invalid character in the highlight verb
	errBadAttr = "%%!h(BADATTR)" // unknown attribute in the highlight verb
)

var colors = map[string]tcell.Color{
	"Black":   tcell.ColorBlack,
	"Maroon":  tcell.ColorMaroon,
	"Green":   tcell.ColorGreen,
	"Olive":   tcell.ColorOlive,
	"Navy":    tcell.ColorNavy,
	"Purple":  tcell.ColorPurple,
	"Teal":    tcell.ColorTeal,
	"Silver":  tcell.ColorSilver,
	"Gray":    tcell.ColorGray,
	"Red":     tcell.ColorRed,
	"Lime":    tcell.ColorLime,
	"Yellow":  tcell.ColorYellow,
	"Blue":    tcell.ColorBlue,
	"Fuchsia": tcell.ColorFuchsia,
	"Aqua":    tcell.ColorAqua,
	"White":   tcell.ColorWhite,
}

// highlighter holds the state of the scanner.
type highlighter struct {
	s       string        // string being scanned
	pos     int           // position in s
	buf     *bytes.Buffer // where result is built
	color   bool          // color or strip the highlight verbs
	fg      bool          // foreground or background color attribute
	noAttrs bool          // not written attrs to buf
}

var ti, tiErr = tcell.LookupTerminfo(os.Getenv("TERM"))

// Highlight replaces the highlight verbs in s with the appropriate control sequences and
// then returns the resulting string.
// It is a thin wrapper around Run.
func Highlight(s string) string {
	return Run(s, true)
}

// Strip removes all highlight verbs in s and then returns the resulting string.
// It is a thin wrapper around Run.
func Strip(s string) string {
	return Run(s, false)
}

// Run runs a highlighter with s as the input and then returns the output. The strip argument
// determines whether the highlight verbs will be stripped or instead replaced with
// their appropriate control sequences.
// Do not use this directly unless you know what you are doing.
func Run(s string, color bool) string {
	hl := getHighlighter(s, color)
	defer hl.free()
	return hl.run()
}

// highlighterPool allows the reuse of highlighters to avoid allocations.
var highlighterPool = sync.Pool{
	New: func() interface{} {
		hl := new(highlighter)
		// The initial capacity avoids constant reallocation during growth.
		hl.buf = bytes.NewBuffer(make([]byte, 0, 45))
		return hl
	},
}

// getHighlighter returns a new initialized highlighter from the pool.
func getHighlighter(s string, color bool) (hl *highlighter) {
	hl = highlighterPool.Get().(*highlighter)
	hl.s = s
	if tiErr == nil {
		hl.color = color
	} else {
		hl.color = false
	}
	return
}

// free resets the highlighter.
func (hl *highlighter) free() {
	hl.buf.Reset()
	hl.pos = 0
	highlighterPool.Put(hl)
}

// stateFn represents the state of the scanner as a function that returns the next state.
type stateFn func(*highlighter) stateFn

// run runs the state machine for the highlighter.
func (hl *highlighter) run() string {
	for state := scanText; state != nil; {
		state = state(hl)
	}
	return hl.buf.String()
}

const eof = -1

// get returns the current rune.
func (hl *highlighter) get() rune {
	if hl.pos >= len(hl.s) {
		return eof
	}
	return rune(hl.s[hl.pos])
}

// writePrev writes n previous characters to the buffer
func (hl *highlighter) writePrev(n int) {
	hl.pos++
	hl.buf.WriteString(hl.s[hl.pos-n : hl.pos])
}

// writeFrom writes the characters from ppos to pos to the buffer.
func (hl *highlighter) writeFrom(ppos int) {
	if hl.pos > ppos {
		// Append remaining characters.
		hl.buf.WriteString(hl.s[ppos:hl.pos])
	}
}

// scanText scans until the next verb.
func scanText(hl *highlighter) stateFn {
	ppos := hl.pos
	// Find next verb.
	for {
		switch hl.get() {
		case eof:
			hl.writeFrom(ppos)
			return nil
		case '%':
			hl.writeFrom(ppos)
			hl.pos++
			if hl.color {
				return scanVerb
			}
			return stripVerb
		}
		hl.pos++
	}
}

// stripVerb skips the current verb.
func stripVerb(hl *highlighter) stateFn {
	switch hl.get() {
	case 'r':
		// Strip the reset verb.
		hl.pos++
	case 'h':
		// Strip inside the highlight verb.
		hl.pos++
		j := strings.IndexByte(hl.s[hl.pos:], ']')
		if j == -1 {
			hl.buf.WriteString(errInvalid)
			return nil
		}
		hl.pos += j + 1
	case eof:
		// Let fmt handle "%!h(NOVERB)".
		hl.buf.WriteByte('%')
		return nil
	default:
		// Include the verb.
		hl.writePrev(2)
	}
	return scanText
}

// scanVerb scans the current verb.
func scanVerb(hl *highlighter) stateFn {
	switch hl.get() {
	case 'r':
		hl.pos++
		return verbReset
	case 'h':
		hl.pos += 2
		hl.noAttrs = true
		return scanHighlight
	case eof:
		// Let fmt handle "%!h(NOVERB)".
		hl.buf.WriteByte('%')
		return nil
	}
	hl.writePrev(2)
	return scanText
}

// verbReset writes the reset verb with the reset control sequence.
func verbReset(hl *highlighter) stateFn {
	hl.buf.WriteString(ti.AttrOff)
	return scanText
}

// scanHighlight scans the highlight verb for attributes,
// then writes a control sequence derived from said attributes to the buffer.
func scanHighlight(hl *highlighter) stateFn {
	for {
		r := hl.get()
		switch {
		case r == 'f':
			hl.fg = true
			return scanColor
		case r == 'b':
			hl.fg = false
			return scanColor
		case unicode.IsLetter(r):
			return scanAttribute
		case r == '+':
			hl.pos++
			continue
		case r == ']':
			if hl.noAttrs {
				hl.buf.WriteString(errMissing)
				return nil
			}
			hl.pos++
			return scanText
		default:
			hl.buf.WriteString(errInvalid)
			return nil
		}
	}
}

// scanAttribute scans a named attribute.
func scanAttribute(hl *highlighter) stateFn {
	start := hl.pos
	for unicode.IsLetter(hl.get()) {
		hl.pos++
	}
	a := hl.s[start:hl.pos]
	switch a {
	case "bold":
		a = ti.Bold
	case "underline":
		a = ti.Underline
	case "reverse":
		a = ti.Reverse
	case "blink":
		a = ti.Blink
	case "dim":
		a = ti.Dim
	case "attrOff":
		a = ti.AttrOff
	default:
		hl.buf.WriteString(errBadAttr)
		return nil
	}
	hl.buf.WriteString(a)
	hl.noAttrs = false
	return scanHighlight
}

// scanColor scans a color attribute.
func scanColor(hl *highlighter) stateFn {
	hl.pos++
	if hl.get() != 'g' {
		hl.pos--
		return scanAttribute
	}
	hl.pos++
	r := hl.get()
	switch {
	case unicode.IsNumber(r):
		return scanColor256
	case unicode.IsLetter(r):
		// continue
	default:
		hl.buf.WriteString(errBadAttr)
		return nil
	}
	start := hl.pos
	hl.pos++
	for unicode.IsLetter(hl.get()) {
		hl.pos++
	}
	if c, ok := colors[hl.s[start:hl.pos]]; ok {
		if hl.fg {
			hl.buf.WriteString(ti.TColor(c, -1))
		} else {
			hl.buf.WriteString(ti.TColor(-1, c))
		}
		hl.noAttrs = false
		return scanHighlight
	}
	hl.buf.WriteString(errBadAttr)
	return nil
}

// scanColor256 scans a 256 color attribute.
func scanColor256(hl *highlighter) stateFn {
	start := hl.pos
	hl.pos++
	for unicode.IsNumber(hl.get()) {
		hl.pos++
	}
	t, _ := strconv.Atoi(hl.s[start:hl.pos])
	n := tcell.Color(t)
	if hl.fg {
		hl.buf.WriteString(ti.TColor(n, -1))
	} else {
		hl.buf.WriteString(ti.TColor(-1, n))
	}
	hl.noAttrs = false
	return scanHighlight
}
<commit_msg>moved eof to last<commit_after>package color

import (
	"bytes"
	"os"
	"strconv"
	"strings"
	"sync"
	"unicode"

	"github.com/gdamore/tcell"
)

const (
	errMissing = "%%!h(MISSING)" // no attributes in the highlight verb
	errInvalid = "%%!h(INVALID)" // invalid character in the highlight verb
	errBadAttr = "%%!h(BADATTR)" // unknown attribute in the highlight verb
)

var colors = map[string]tcell.Color{
	"Black":   tcell.ColorBlack,
	"Maroon":  tcell.ColorMaroon,
	"Green":   tcell.ColorGreen,
	"Olive":   tcell.ColorOlive,
	"Navy":    tcell.ColorNavy,
	"Purple":  tcell.ColorPurple,
	"Teal":    tcell.ColorTeal,
	"Silver":  tcell.ColorSilver,
	"Gray":    tcell.ColorGray,
	"Red":     tcell.ColorRed,
	"Lime":    tcell.ColorLime,
	"Yellow":  tcell.ColorYellow,
	"Blue":    tcell.ColorBlue,
	"Fuchsia": tcell.ColorFuchsia,
	"Aqua":    tcell.ColorAqua,
	"White":   tcell.ColorWhite,
}

// highlighter holds the state of the scanner.
type highlighter struct {
	s       string        // string being scanned
	pos     int           // position in s
	buf     *bytes.Buffer // where result is built
	color   bool          // color or strip the highlight verbs
	fg      bool          // foreground or background color attribute
	noAttrs bool          // not written attrs to buf
}

var ti, tiErr = tcell.LookupTerminfo(os.Getenv("TERM"))

// Highlight replaces the highlight verbs in s with the appropriate control sequences and
// then returns the resulting string.
// It is a thin wrapper around Run.
func Highlight(s string) string {
	return Run(s, true)
}

// Strip removes all highlight verbs in s and then returns the resulting string.
// It is a thin wrapper around Run.
func Strip(s string) string {
	return Run(s, false)
}

// Run runs a highlighter with s as the input and then returns the output. The strip argument
// determines whether the highlight verbs will be stripped or instead replaced with
// their appropriate control sequences.
// Do not use this directly unless you know what you are doing.
func Run(s string, color bool) string {
	hl := getHighlighter(s, color)
	defer hl.free()
	return hl.run()
}

// highlighterPool allows the reuse of highlighters to avoid allocations.
var highlighterPool = sync.Pool{
	New: func() interface{} {
		hl := new(highlighter)
		// The initial capacity avoids constant reallocation during growth.
		hl.buf = bytes.NewBuffer(make([]byte, 0, 45))
		return hl
	},
}

// getHighlighter returns a new initialized highlighter from the pool.
func getHighlighter(s string, color bool) (hl *highlighter) {
	hl = highlighterPool.Get().(*highlighter)
	hl.s = s
	if tiErr == nil {
		hl.color = color
	} else {
		hl.color = false
	}
	return
}

// free resets the highlighter.
func (hl *highlighter) free() {
	hl.buf.Reset()
	hl.pos = 0
	highlighterPool.Put(hl)
}

// stateFn represents the state of the scanner as a function that returns the next state.
type stateFn func(*highlighter) stateFn

// run runs the state machine for the highlighter.
func (hl *highlighter) run() string {
	for state := scanText; state != nil; {
		state = state(hl)
	}
	return hl.buf.String()
}

const eof = -1

// get returns the current rune.
func (hl *highlighter) get() rune {
	if hl.pos >= len(hl.s) {
		return eof
	}
	return rune(hl.s[hl.pos])
}

// writePrev writes n previous characters to the buffer
func (hl *highlighter) writePrev(n int) {
	hl.pos++
	hl.buf.WriteString(hl.s[hl.pos-n : hl.pos])
}

// writeFrom writes the characters from ppos to pos to the buffer.
func (hl *highlighter) writeFrom(ppos int) {
	if hl.pos > ppos {
		// Append remaining characters.
		hl.buf.WriteString(hl.s[ppos:hl.pos])
	}
}

// scanText scans until the next verb.
func scanText(hl *highlighter) stateFn {
	ppos := hl.pos
	// Find next verb.
	for {
		switch hl.get() {
		case '%':
			hl.writeFrom(ppos)
			hl.pos++
			if hl.color {
				return scanVerb
			}
			return stripVerb
		case eof:
			hl.writeFrom(ppos)
			return nil
		}
		hl.pos++
	}
}

// stripVerb skips the current verb.
func stripVerb(hl *highlighter) stateFn {
	switch hl.get() {
	case 'r':
		// Strip the reset verb.
		hl.pos++
	case 'h':
		// Strip inside the highlight verb.
		hl.pos++
		j := strings.IndexByte(hl.s[hl.pos:], ']')
		if j == -1 {
			hl.buf.WriteString(errInvalid)
			return nil
		}
		hl.pos += j + 1
	case eof:
		// Let fmt handle "%!h(NOVERB)".
		hl.buf.WriteByte('%')
		return nil
	default:
		// Include the verb.
		hl.writePrev(2)
	}
	return scanText
}

// scanVerb scans the current verb.
func scanVerb(hl *highlighter) stateFn {
	switch hl.get() {
	case 'r':
		hl.pos++
		return verbReset
	case 'h':
		hl.pos += 2
		hl.noAttrs = true
		return scanHighlight
	case eof:
		// Let fmt handle "%!h(NOVERB)".
		hl.buf.WriteByte('%')
		return nil
	}
	hl.writePrev(2)
	return scanText
}

// verbReset writes the reset verb with the reset control sequence.
func verbReset(hl *highlighter) stateFn {
	hl.buf.WriteString(ti.AttrOff)
	return scanText
}

// scanHighlight scans the highlight verb for attributes,
// then writes a control sequence derived from said attributes to the buffer.
func scanHighlight(hl *highlighter) stateFn {
	for {
		r := hl.get()
		switch {
		case r == 'f':
			hl.fg = true
			return scanColor
		case r == 'b':
			hl.fg = false
			return scanColor
		case unicode.IsLetter(r):
			return scanAttribute
		case r == '+':
			hl.pos++
			continue
		case r == ']':
			if hl.noAttrs {
				hl.buf.WriteString(errMissing)
				return nil
			}
			hl.pos++
			return scanText
		default:
			hl.buf.WriteString(errInvalid)
			return nil
		}
	}
}

// scanAttribute scans a named attribute.
func scanAttribute(hl *highlighter) stateFn {
	start := hl.pos
	for unicode.IsLetter(hl.get()) {
		hl.pos++
	}
	a := hl.s[start:hl.pos]
	switch a {
	case "bold":
		a = ti.Bold
	case "underline":
		a = ti.Underline
	case "reverse":
		a = ti.Reverse
	case "blink":
		a = ti.Blink
	case "dim":
		a = ti.Dim
	case "attrOff":
		a = ti.AttrOff
	default:
		hl.buf.WriteString(errBadAttr)
		return nil
	}
	hl.buf.WriteString(a)
	hl.noAttrs = false
	return scanHighlight
}

// scanColor scans a color attribute.
func scanColor(hl *highlighter) stateFn {
	hl.pos++
	if hl.get() != 'g' {
		hl.pos--
		return scanAttribute
	}
	hl.pos++
	r := hl.get()
	switch {
	case unicode.IsNumber(r):
		return scanColor256
	case unicode.IsLetter(r):
		// continue
	default:
		hl.buf.WriteString(errBadAttr)
		return nil
	}
	start := hl.pos
	hl.pos++
	for unicode.IsLetter(hl.get()) {
		hl.pos++
	}
	if c, ok := colors[hl.s[start:hl.pos]]; ok {
		if hl.fg {
			hl.buf.WriteString(ti.TColor(c, -1))
		} else {
			hl.buf.WriteString(ti.TColor(-1, c))
		}
		hl.noAttrs = false
		return scanHighlight
	}
	hl.buf.WriteString(errBadAttr)
	return nil
}

// scanColor256 scans a 256 color attribute.
func scanColor256(hl *highlighter) stateFn {
	start := hl.pos
	hl.pos++
	for unicode.IsNumber(hl.get()) {
		hl.pos++
	}
	t, _ := strconv.Atoi(hl.s[start:hl.pos])
	n := tcell.Color(t)
	if hl.fg {
		hl.buf.WriteString(ti.TColor(n, -1))
	} else {
		hl.buf.WriteString(ti.TColor(-1, n))
	}
	hl.noAttrs = false
	return scanHighlight
}
<|endoftext|>