<commit_before>package effio

/*
 * effio provides ways to build a suite of fio tests & devices, combine them
 * then generate reports / graphs.
 *
 * The fio config file name is stripped of its path and .fio extension to be
 * used in composite test names. A test against the above device with
 * "rand_512b_write_iops.fio" will be named
 * "rand_512b_write_iops-samsung_840_pro_256".
 */

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strings"
	"time"
)

// test data generated on the fly based on info above
type Test struct {
	Name        string      // name to be used in tests, files, etc.
	Dir         string      // directory for writing configs, logs, etc.
	FioArgs     []string    // the arguments to the fio command for the test
	FioFile     string      // generated fio config file name
	FioJson     string      // generated fio json output file name
	FioBWLog    string      // filename for the bandwidth log
	FioLatLog   string      // filename for the latency log
	FioIopsLog  string      // filename for the iops log
	TestJson    string      // dump the test data (this struct) to this file
	CmdFile     string      // write the exact fio command used to this file
	FioConfTmpl FioConfTmpl // template info struct
	Device      Device      // device info struct
}

// a test suite has a global id that is also used as a directory name
type Suite struct {
	Id        string
	Created   time.Time // time the test was generated / run
	EffioCmd  []string  // os.Args() of the effio command used to generate the suite
	SuiteJson string
	Tests     []Test
}

// NewSuite returns an initialized Suite with the given
// id and the Created field set to the current time.
func NewSuite(id string) Suite {
	now := time.Now()
	fname := path.Join(id, "suite.json")
	return Suite{id, now, os.Args, fname, []Test{}}
}

// LoadSuiteJson loads a suite from JSON. Argument is a path to a
// JSON file that has a complete suite's information in it.
func LoadSuiteJson(spath string) (suite Suite) {
	data, err := ioutil.ReadFile(spath)
	if err != nil {
		log.Fatalf("Could not read suite JSON file '%s': %s", spath, err)
	}

	err = json.Unmarshal(data, &suite)
	if err != nil {
		log.Fatalf("Could not parse suite JSON in file '%s': %s", spath, err)
	}

	return suite
}

// Run the whole suite one at a time letting fio write its output into
// the suite directories. Repeated runs will overwrite files; behavior
// is dependent on what fio does with existing files for now.
func (suite *Suite) Run(spath string) {
	for _, test := range suite.Tests {
		test.Run(spath)
	}
}

// Populate the test suite with the (cartesian) product of
// Devices x FioConfTmpls to get all combinations.
// This does not modify the filesystem.
func (suite *Suite) Populate(dl Devices, ftl FioConfTmpls) {
	for _, tp := range ftl {
		for _, dev := range dl {
			if dev.Ignore {
				continue
			}

			// I suppose these conventions could be defined higher up in the call stack
			// but this makes things a little easier to modify down the road.
			testName := fmt.Sprintf("%s-%s", dev.Name, tp.Name)
			testDir := path.Join(suite.Id, testName)
			args := []string{"--output-format=json", "--output=output.json", "config.fio"}

			// fio adds _$type.log to log file names so only provide the base name
			test := Test{
				Name:        testName,
				Dir:         testDir,
				FioArgs:     args,
				FioFile:     "config.fio",
				FioJson:     "output.json",
				FioBWLog:    "bw",
				FioLatLog:   "lat",
				FioIopsLog:  "iops",
				TestJson:    "test.json",
				CmdFile:     "run.sh",
				FioConfTmpl: tp,
				Device:      dev,
			}

			suite.Tests = append(suite.Tests, test)
		}
	}
}

// WriteAll(path) writes a suite out to a set of directories and files.
func (suite *Suite) WriteAll(basePath string) {
	suite.mkdirAll(basePath)

	suite.WriteSuiteJson(basePath)

	for _, test := range suite.Tests {
		test.WriteFioFile(basePath)
		test.WriteTestJson(basePath)
		test.WriteCmdFile(basePath)
	}
}

// WriteSuiteJson dumps the suite data structure to a JSON file. This
// file is used by some effio subcommands, such as run_suite and various
// reports.
// <basePath>/<suite id>/suite.json
func (suite *Suite) WriteSuiteJson(basePath string) {
	outfile := path.Join(basePath, suite.SuiteJson)

	js, err := json.MarshalIndent(suite, "", "  ")
	if err != nil {
		log.Fatalf("Failed to encode suite data as JSON: %s\n", err)
	}

	// MarshalIndent does not follow the final brace with a newline
	js = append(js, byte('\n'))

	err = ioutil.WriteFile(outfile, js, 0644)
	if err != nil {
		log.Fatalf("Failed to write suite JSON data file '%s': %s\n", outfile, err)
	}
}

// mkdirAll(path) creates the directory structure of a test suite
// under directory 'path'. This must be called before the Write*()
// methods or they will fail. It only makes sense to call this after
// Populate().
func (suite *Suite) mkdirAll(basePath string) {
	sdir := path.Join(basePath, suite.Id)

	for _, t := range suite.Tests {
		tdir := path.Join(sdir, t.Name)
		err := os.MkdirAll(tdir, 0755)
		if err != nil {
			log.Fatalf("Failed to create test directory '%s': %s\n", tdir, err)
		}
	}
}

func (test *Test) Run(spath string) {
	tpath := path.Join(spath, test.Dir)
	err := os.Chdir(tpath)
	if err != nil {
		log.Fatalf("Could not chdir to '%s': %s\n", tpath, err)
	}

	fioPath, err := exec.LookPath("fio")
	if err != nil {
		log.Fatalf("Could not locate an fio command in PATH: %s\n", err)
	}

	cmd := exec.Command(fioPath, test.FioArgs...)
	before := time.Now()

	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Fatal(err)
	}

	if err := cmd.Start(); err != nil {
		log.Fatalf("Could not run '%s %s': %s\n", fioPath, strings.Join(test.FioArgs, " "))
	}

	// grab stderr in case something goes wrong
	// TODO: switch this to io.Copy?
	errors, err := ioutil.ReadAll(stderr)
	if err != nil {
		log.Fatal(err)
	}

	err = cmd.Wait()

	// TODO: figure out if this is worth recording and record it
	elapsed := time.Since(before)

	// it might be OK to let 1 fio command out of a suite fail?
	if err != nil {
		log.Printf(string(errors))
		log.Fatalf("Command '%s %s' failed: %s\n", fioPath, strings.Join(test.FioArgs, " "), err)
	}
	log.Printf("Elapsed: %s\n", elapsed)
}

// DumpFioFile writes the fio configuration file.
// <basePath>/<suite id>/<generated test name>/config.fio
func (test *Test) WriteFioFile(basePath string) {
	outfile := path.Join(basePath, test.Dir, test.FioFile)

	fd, err := os.OpenFile(outfile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("Failed to create fio config file '%s': %s\n", outfile, err)
	}
	defer fd.Close()

	err = test.FioConfTmpl.tmpl.Execute(fd, test)
	if err != nil {
		log.Fatalf("Template execution failed for '%s': %s\n", outfile, err)
	}
}

// WriteTestJson dumps the suite data structure to a JSON file for posterity (and debugging).
// <basePath>/<suite id>/<generated test name>/test.json
func (test *Test) WriteTestJson(basePath string) {
	outfile := path.Join(basePath, test.Dir, test.TestJson)

	js, err := json.MarshalIndent(test, "", "  ")
	if err != nil {
		log.Fatalf("Failed to encode test data as JSON: %s\n", err)
	}

	// MarshalIndent does not follow the final brace with a newline
	js = append(js, byte('\n'))

	err = ioutil.WriteFile(outfile, js, 0644)
	if err != nil {
		log.Fatalf("Failed to write test JSON data file '%s': %s\n", outfile, err)
	}
}

// WriteCmdFile writes the command to a file as a mini shell script.
// <basePath>/<suite id>/<test name>/run.sh
func (test *Test) WriteCmdFile(basePath string) {
	outfile := path.Join(basePath, test.Dir, test.CmdFile)

	fd, err := os.OpenFile(outfile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		log.Fatalf("Failed to create command file '%s': %s\n", outfile, err)
	}
	defer fd.Close()

	// just use 'fio' if it isn't found on the path
	fioPath, err := exec.LookPath("fio")
	if err != nil {
		fioPath = "fio"
	}
	fmt.Fprintf(fd, "#!/bin/bash -x\n%s %s\n", fioPath, strings.Join(test.FioArgs, " "))
}
<commit_msg>Use fully-qualified path to launch tests.<commit_after>package effio

/*
 * effio provides ways to build a suite of fio tests & devices, combine them
 * then generate reports / graphs.
 *
 * The fio config file name is stripped of its path and .fio extension to be
 * used in composite test names. A test against the above device with
 * "rand_512b_write_iops.fio" will be named
 * "rand_512b_write_iops-samsung_840_pro_256".
 */

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"strings"
	"time"
)

// test data generated on the fly based on info above
type Test struct {
	Name        string      // name to be used in tests, files, etc.
	Dir         string      // directory for writing configs, logs, etc.
	FioArgs     []string    // the arguments to the fio command for the test
	FioFile     string      // generated fio config file name
	FioJson     string      // generated fio json output file name
	FioBWLog    string      // filename for the bandwidth log
	FioLatLog   string      // filename for the latency log
	FioIopsLog  string      // filename for the iops log
	TestJson    string      // dump the test data (this struct) to this file
	CmdFile     string      // write the exact fio command used to this file
	FioConfTmpl FioConfTmpl // template info struct
	Device      Device      // device info struct
}

// a test suite has a global id that is also used as a directory name
type Suite struct {
	Id        string
	Created   time.Time // time the test was generated / run
	EffioCmd  []string  // os.Args() of the effio command used to generate the suite
	SuiteJson string
	Tests     []Test
}

// NewSuite returns an initialized Suite with the given
// id and the Created field set to the current time.
func NewSuite(id string) Suite {
	now := time.Now()
	fname := path.Join(id, "suite.json")
	return Suite{id, now, os.Args, fname, []Test{}}
}

// LoadSuiteJson loads a suite from JSON. Argument is a path to a
// JSON file that has a complete suite's information in it.
func LoadSuiteJson(spath string) (suite Suite) {
	data, err := ioutil.ReadFile(spath)
	if err != nil {
		log.Fatalf("Could not read suite JSON file '%s': %s", spath, err)
	}

	err = json.Unmarshal(data, &suite)
	if err != nil {
		log.Fatalf("Could not parse suite JSON in file '%s': %s", spath, err)
	}

	return suite
}

// Run the whole suite one at a time letting fio write its output into
// the suite directories. Repeated runs will overwrite files; behavior
// is dependent on what fio does with existing files for now.
func (suite *Suite) Run(spath string) {
	wd, err := os.Getwd()
	if err != nil {
		log.Fatalf("Could not get working directory: %s\n", err)
	}

	for _, test := range suite.Tests {
		test.Run(path.Join(wd, spath))
	}
}

// Populate the test suite with the (cartesian) product of
// Devices x FioConfTmpls to get all combinations.
// This does not modify the filesystem.
func (suite *Suite) Populate(dl Devices, ftl FioConfTmpls) {
	for _, tp := range ftl {
		for _, dev := range dl {
			if dev.Ignore {
				continue
			}

			// I suppose these conventions could be defined higher up in the call stack
			// but this makes things a little easier to modify down the road.
			testName := fmt.Sprintf("%s-%s", dev.Name, tp.Name)
			testDir := path.Join(suite.Id, testName)
			args := []string{"--output-format=json", "--output=output.json", "config.fio"}

			// fio adds _$type.log to log file names so only provide the base name
			test := Test{
				Name:        testName,
				Dir:         testDir,
				FioArgs:     args,
				FioFile:     "config.fio",
				FioJson:     "output.json",
				FioBWLog:    "bw",
				FioLatLog:   "lat",
				FioIopsLog:  "iops",
				TestJson:    "test.json",
				CmdFile:     "run.sh",
				FioConfTmpl: tp,
				Device:      dev,
			}

			suite.Tests = append(suite.Tests, test)
		}
	}
}

// WriteAll(path) writes a suite out to a set of directories and files.
func (suite *Suite) WriteAll(basePath string) {
	suite.mkdirAll(basePath)

	suite.WriteSuiteJson(basePath)

	for _, test := range suite.Tests {
		test.WriteFioFile(basePath)
		test.WriteTestJson(basePath)
		test.WriteCmdFile(basePath)
	}
}

// WriteSuiteJson dumps the suite data structure to a JSON file. This
// file is used by some effio subcommands, such as run_suite and various
// reports.
// <basePath>/<suite id>/suite.json
func (suite *Suite) WriteSuiteJson(basePath string) {
	outfile := path.Join(basePath, suite.SuiteJson)

	js, err := json.MarshalIndent(suite, "", "  ")
	if err != nil {
		log.Fatalf("Failed to encode suite data as JSON: %s\n", err)
	}

	// MarshalIndent does not follow the final brace with a newline
	js = append(js, byte('\n'))

	err = ioutil.WriteFile(outfile, js, 0644)
	if err != nil {
		log.Fatalf("Failed to write suite JSON data file '%s': %s\n", outfile, err)
	}
}

// mkdirAll(path) creates the directory structure of a test suite
// under directory 'path'. This must be called before the Write*()
// methods or they will fail. It only makes sense to call this after
// Populate().
func (suite *Suite) mkdirAll(basePath string) {
	sdir := path.Join(basePath, suite.Id)

	for _, t := range suite.Tests {
		tdir := path.Join(sdir, t.Name)
		err := os.MkdirAll(tdir, 0755)
		if err != nil {
			log.Fatalf("Failed to create test directory '%s': %s\n", tdir, err)
		}
	}
}

func (test *Test) Run(spath string) {
	tpath := path.Join(spath, test.Dir)

	err := os.Chdir(tpath)
	if err != nil {
		log.Fatalf("Could not chdir to '%s': %s\n", tpath, err)
	}

	fioPath, err := exec.LookPath("fio")
	if err != nil {
		log.Fatalf("Could not locate an fio command in PATH: %s\n", err)
	}

	cmd := exec.Command(fioPath, test.FioArgs...)
	before := time.Now()

	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Fatal(err)
	}

	if err := cmd.Start(); err != nil {
		log.Fatalf("Could not run '%s %s': %s\n", fioPath, strings.Join(test.FioArgs, " "))
	}

	// grab stderr in case something goes wrong
	// TODO: switch this to io.Copy?
	errors, err := ioutil.ReadAll(stderr)
	if err != nil {
		log.Fatal(err)
	}

	err = cmd.Wait()

	// TODO: figure out if this is worth recording and record it
	elapsed := time.Since(before)

	// it might be OK to let 1 fio command out of a suite fail?
	if err != nil {
		log.Printf(string(errors))
		log.Fatalf("Command '%s %s' failed: %s\n", fioPath, strings.Join(test.FioArgs, " "), err)
	}
	log.Printf("Elapsed: %s\n", elapsed)
}

// DumpFioFile writes the fio configuration file.
// <basePath>/<suite id>/<generated test name>/config.fio
func (test *Test) WriteFioFile(basePath string) {
	outfile := path.Join(basePath, test.Dir, test.FioFile)

	fd, err := os.OpenFile(outfile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("Failed to create fio config file '%s': %s\n", outfile, err)
	}
	defer fd.Close()

	err = test.FioConfTmpl.tmpl.Execute(fd, test)
	if err != nil {
		log.Fatalf("Template execution failed for '%s': %s\n", outfile, err)
	}
}

// WriteTestJson dumps the suite data structure to a JSON file for posterity (and debugging).
// <basePath>/<suite id>/<generated test name>/test.json
func (test *Test) WriteTestJson(basePath string) {
	outfile := path.Join(basePath, test.Dir, test.TestJson)

	js, err := json.MarshalIndent(test, "", "  ")
	if err != nil {
		log.Fatalf("Failed to encode test data as JSON: %s\n", err)
	}

	// MarshalIndent does not follow the final brace with a newline
	js = append(js, byte('\n'))

	err = ioutil.WriteFile(outfile, js, 0644)
	if err != nil {
		log.Fatalf("Failed to write test JSON data file '%s': %s\n", outfile, err)
	}
}

// WriteCmdFile writes the command to a file as a mini shell script.
// <basePath>/<suite id>/<test name>/run.sh
func (test *Test) WriteCmdFile(basePath string) {
	outfile := path.Join(basePath, test.Dir, test.CmdFile)

	fd, err := os.OpenFile(outfile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		log.Fatalf("Failed to create command file '%s': %s\n", outfile, err)
	}
	defer fd.Close()

	// just use 'fio' if it isn't found on the path
	fioPath, err := exec.LookPath("fio")
	if err != nil {
		fioPath = "fio"
	}
	fmt.Fprintf(fd, "#!/bin/bash -x\n%s %s\n", fioPath, strings.Join(test.FioArgs, " "))
}
<|endoftext|>


<commit_before>package format

import (
	"bytes"
	"io"
	"os/exec"
	"sort"
	"time"

	"github.com/AlexanderThaller/lablog/src/data"
	"github.com/AlexanderThaller/lablog/src/helper"
	"github.com/juju/errgo"
)

const (
	Name = "format"
)

const AsciiDocSettings = `:toc: right
:toclevels: 2
:sectanchors:
:sectlink:
:icons: font
:linkattrs:
:numbered:
:idprefix:
:idseparator: -
:doctype: book
:source-highlighter: pygments
:listing-caption: Listing`

func ProjectsEntries(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Entries \n\n")

	for _, project := range projects {
		notes, err := helper.FilteredNotesByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}

		todos, err := helper.FilteredTodosByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}
		todos = data.FilterTodosLatest(todos)
		todos = data.FilterTodosAreNotDone(todos)

		project.Format(writer, 1)
		if len(todos) != 0 {
			Todos(writer, todos)
			io.WriteString(writer, "\n")
		}

		if len(notes) != 0 {
			Notes(writer, notes)
		}
	}

	return nil
}

func ProjectsNotes(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Notes \n\n")

	for _, project := range projects {
		notes, err := helper.FilteredNotesByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}

		project.Format(writer, 1)
		Notes(writer, notes)
	}

	return nil
}

func ProjectsTodos(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Todos \n\n")

	for _, project := range projects {
		todos, err := helper.FilteredTodosByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}
		todos = data.FilterTodosLatest(todos)
		todos = data.FilterTodosAreNotDone(todos)

		if len(todos) == 0 {
			continue
		}

		project.Format(writer, 1)
		Todos(writer, todos)
		io.WriteString(writer, "\n")
	}

	return nil
}

func ProjectsTracks(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Tracks \n\n")

	for _, project := range projects {
		tracks, err := helper.FilteredTracksByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered tracks")
		}

		if len(tracks) == 0 {
			continue
		}

		project.Format(writer, 1)
		Tracks(writer, tracks)
		io.WriteString(writer, "\n")
	}

	return nil
}

func ProjectsDates(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Dates \n\n")

	dates, err := helper.ProjectsDates(projects, start, end)
	if err != nil {
		return errgo.Notef(err, "can not get dates for projects")
	}

	sort.Strings(dates)

	for _, date := range dates {
		io.WriteString(writer, "* "+date+"\n")
	}

	return nil
}

func Todos(writer io.Writer, todos []data.Todo) {
	io.WriteString(writer, "=== Todos\n\n")

	sort.Sort(data.TodosByName(todos))
	for _, todo := range todos {
		todo.Format(writer, 2)
	}
}

func Notes(writer io.Writer, notes []data.Note) {
	io.WriteString(writer, "=== Notes\n\n")

	notes = data.FilterNotesNotEmpty(notes)

	sort.Sort(data.NotesByTimeStamp(notes))
	for _, note := range notes {
		note.Format(writer, 2)
	}
}

func Tracks(writer io.Writer, tracks []data.Track) {
	io.WriteString(writer, "=== Tracks\n\n")

	sort.Sort(data.TracksByTimeStamp(tracks))
	for _, track := range tracks {
		track.Format(writer, 2)
	}
}

func AsciiDoctor(reader io.Reader, writer io.Writer) error {
	stderr := new(bytes.Buffer)

	command := exec.Command("asciidoctor", "-")
	command.Stdin = reader
	command.Stdout = writer
	command.Stderr = stderr

	err := command.Run()
	if err != nil {
		return errgo.Notef(errgo.Notef(err, "can not run asciidoctor"),
			stderr.String())
	}

	return nil
}

func Log(writer io.Writer, projects []data.Project, start, end time.Time) error {
	var allnotes []data.Note

	for _, project := range projects {
		notes, err := helper.FilteredNotesByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}

		for _, note := range notes {
			allnotes = append(allnotes, note)
		}
	}

	allnotes = data.FilterNotesNotEmpty(allnotes)
	sort.Sort(data.NotesByTimeStamp(allnotes))

	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Log \n\n")
	for _, note := range allnotes {
		note.FormatTimeStamp(writer, 2)
		note.Project.Format(writer, 2)
		note.FormatText(writer, 4)
	}

	return nil
}
<commit_msg>will not print notes header if there are no notes for the project.<commit_after>package format

import (
	"bytes"
	"io"
	"os/exec"
	"sort"
	"time"

	"github.com/AlexanderThaller/lablog/src/data"
	"github.com/AlexanderThaller/lablog/src/helper"
	"github.com/juju/errgo"
)

const (
	Name = "format"
)

const AsciiDocSettings = `:toc: right
:toclevels: 2
:sectanchors:
:sectlink:
:icons: font
:linkattrs:
:numbered:
:idprefix:
:idseparator: -
:doctype: book
:source-highlighter: pygments
:listing-caption: Listing`

func ProjectsEntries(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Entries \n\n")

	for _, project := range projects {
		notes, err := helper.FilteredNotesByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}

		todos, err := helper.FilteredTodosByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}
		todos = data.FilterTodosLatest(todos)
		todos = data.FilterTodosAreNotDone(todos)

		project.Format(writer, 1)
		if len(todos) != 0 {
			Todos(writer, todos)
			io.WriteString(writer, "\n")
		}

		if len(notes) != 0 {
			Notes(writer, notes)
		}
	}

	return nil
}

func ProjectsNotes(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Notes \n\n")

	for _, project := range projects {
		notes, err := helper.FilteredNotesByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}

		if len(notes) == 0 {
			continue
		}

		project.Format(writer, 1)
		Notes(writer, notes)
	}

	return nil
}

func ProjectsTodos(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Todos \n\n")

	for _, project := range projects {
		todos, err := helper.FilteredTodosByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}
		todos = data.FilterTodosLatest(todos)
		todos = data.FilterTodosAreNotDone(todos)

		if len(todos) == 0 {
			continue
		}

		project.Format(writer, 1)
		Todos(writer, todos)
		io.WriteString(writer, "\n")
	}

	return nil
}

func ProjectsTracks(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Tracks \n\n")

	for _, project := range projects {
		tracks, err := helper.FilteredTracksByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered tracks")
		}

		if len(tracks) == 0 {
			continue
		}

		project.Format(writer, 1)
		Tracks(writer, tracks)
		io.WriteString(writer, "\n")
	}

	return nil
}

func ProjectsDates(writer io.Writer, projects []data.Project, start, end time.Time) error {
	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Dates \n\n")

	dates, err := helper.ProjectsDates(projects, start, end)
	if err != nil {
		return errgo.Notef(err, "can not get dates for projects")
	}

	sort.Strings(dates)

	for _, date := range dates {
		io.WriteString(writer, "* "+date+"\n")
	}

	return nil
}

func Todos(writer io.Writer, todos []data.Todo) {
	io.WriteString(writer, "=== Todos\n\n")

	sort.Sort(data.TodosByName(todos))
	for _, todo := range todos {
		todo.Format(writer, 2)
	}
}

func Notes(writer io.Writer, notes []data.Note) {
	io.WriteString(writer, "=== Notes\n\n")

	notes = data.FilterNotesNotEmpty(notes)

	sort.Sort(data.NotesByTimeStamp(notes))
	for _, note := range notes {
		note.Format(writer, 2)
	}
}

func Tracks(writer io.Writer, tracks []data.Track) {
	io.WriteString(writer, "=== Tracks\n\n")

	sort.Sort(data.TracksByTimeStamp(tracks))
	for _, track := range tracks {
		track.Format(writer, 2)
	}
}

func AsciiDoctor(reader io.Reader, writer io.Writer) error {
	stderr := new(bytes.Buffer)

	command := exec.Command("asciidoctor", "-")
	command.Stdin = reader
	command.Stdout = writer
	command.Stderr = stderr

	err := command.Run()
	if err != nil {
		return errgo.Notef(errgo.Notef(err, "can not run asciidoctor"),
			stderr.String())
	}

	return nil
}

func Log(writer io.Writer, projects []data.Project, start, end time.Time) error {
	var allnotes []data.Note

	for _, project := range projects {
		notes, err := helper.FilteredNotesByStartEnd(project, start, end)
		if err != nil {
			return errgo.Notef(err, "can not get filtered notes")
		}

		for _, note := range notes {
			allnotes = append(allnotes, note)
		}
	}

	allnotes = data.FilterNotesNotEmpty(allnotes)
	sort.Sort(data.NotesByTimeStamp(allnotes))

	io.WriteString(writer, AsciiDocSettings+"\n\n")
	io.WriteString(writer, "= Log \n\n")
	for _, note := range allnotes {
		note.FormatTimeStamp(writer, 2)
		note.Project.Format(writer, 2)
		note.FormatText(writer, 4)
	}

	return nil
}
<|endoftext|>


<commit_before>package gloom

import (
  "fmt"
  "hash"
  "hash/fnv"
)

type BloomFilter struct {
   m      int //number of bits in filter
   k      int //number of hash functions
   bits   []uint64
   hash   hash.Hash64
}

func NewFilter(m int, k int) *BloomFilter {
  size := m / 64
  if m % 64 > 0 {
    size++
  }
  bits := make([]uint64, size)
  return &BloomFilter{m,k,bits,fnv.New64a()}
}

func (bf * BloomFilter) GetHashes(key []byte) []uint64 {
  bf.hash.Reset()
  bf.hash.Write(key)
  h := bf.hash.Sum64()
  h1 := uint32(h)
  h2 := uint32(h >> 32)
  hs := make([]uint64,bf.k)
  for i := 0; i < bf.k; i++ {
   //This distributes values so that one hash function
   //can effectively act as k hash functions
   hs[i] = uint64(h1 + uint32(i)*h2)
  }
  return hs
}

func (bf *BloomFilter) HashLoc(h uint64) (chnk uint64,shft uint64) {
    bitInd := h % uint64(bf.m)
    chnkInd := bitInd / 64
    shftInd := bitInd - chnkInd*64
    return chnkInd, shftInd
}

func (bf *BloomFilter) Has(data []byte) bool {
  hs := bf.GetHashes(data)
  for i := 0; i < bf.k; i++ {
    hk := hs[i]
    chnk, shft := bf.HashLoc(hk)
    if bf.bits[chnk] ^ (1 << shft) == bf.bits[chnk] {
      return false
    }
  }
  return true
}

func (bf *BloomFilter) HasString(key string) bool {
  return bf.Has([]byte(key))
}

func (bf *BloomFilter) Put(data []byte) *BloomFilter {
  hs := bf.GetHashes(data)
  for i := 0; i < bf.k; i++ {
    hk := hs[i]
    chnk, shft := bf.HashLoc(hk)
    bf.bits[chnk] |= (1 << shft)
  }
  return bf
}

func (bf *BloomFilter) PutString(key string) *BloomFilter {
  return bf.Put([]byte(key))
}

func (bf *BloomFilter) Print() {
  for _, b := range bf.bits {
    fmt.Printf("%64b\n",b)
  }
}

<commit_msg>Update gloom.go<commit_after>package gloom

import (
  "fmt"
  "hash"
  "hash/fnv"
)

type BloomFilter struct {
   m      int //number of bits in filter
   k      int //number of hash functions
   bits   []uint64
   hash   hash.Hash64
}

func NewFilter(m int, k int) *BloomFilter {
  size := m / 64
  if m % 64 > 0 {
    size++
  }
  bits := make([]uint64, size)
  return &BloomFilter{m,k,bits,fnv.New64a()}
}

func (bf * BloomFilter) GetHashes(key []byte) []uint64 {
  bf.hash.Reset()
  bf.hash.Write(key)
  h := bf.hash.Sum64()
  h1 := uint32(h)
  h2 := uint32(h >> 32)
  hs := make([]uint64,bf.k)
  for i := 0; i < bf.k; i++ {
   //This distributes values so that one hash function
   //can effectively act as k hash functions
   //https://www.eecs.harvard.edu/~michaelm/postscripts/rsa2008.pdf
   hs[i] = uint64(h1 + uint32(i)*h2)
  }
  return hs
}

func (bf *BloomFilter) HashLoc(h uint64) (chnk uint64,shft uint64) {
    bitInd := h % uint64(bf.m)
    chnkInd := bitInd / 64
    shftInd := bitInd - chnkInd*64
    return chnkInd, shftInd
}

func (bf *BloomFilter) Has(data []byte) bool {
  hs := bf.GetHashes(data)
  for i := 0; i < bf.k; i++ {
    hk := hs[i]
    chnk, shft := bf.HashLoc(hk)
    if bf.bits[chnk] ^ (1 << shft) == bf.bits[chnk] {
      return false
    }
  }
  return true
}

func (bf *BloomFilter) HasString(key string) bool {
  return bf.Has([]byte(key))
}

func (bf *BloomFilter) Put(data []byte) *BloomFilter {
  hs := bf.GetHashes(data)
  for i := 0; i < bf.k; i++ {
    hk := hs[i]
    chnk, shft := bf.HashLoc(hk)
    bf.bits[chnk] |= (1 << shft)
  }
  return bf
}

func (bf *BloomFilter) PutString(key string) *BloomFilter {
  return bf.Put([]byte(key))
}

func (bf *BloomFilter) Print() {
  for _, b := range bf.bits {
    fmt.Printf("%64b\n",b)
  }
}

<|endoftext|>


<commit_before>package raft

import (
	"testing"
)

const (
	defaultHeartbeat = 1
	defaultElection  = 5
)

func TestTickMsgHup(t *testing.T) {
	n := New(0, defaultHeartbeat, defaultElection)
	n.sm = newStateMachine(0, []int{0, 1, 2})

	for i := 0; i < defaultElection+1; i++ {
		n.Tick()
	}

	called := false
	for _, m := range n.Msgs() {
		if m.Type == msgVote {
			called = true
		}
	}

	if !called {
		t.Errorf("called = %v, want true", called)
	}
}

func TestTickMsgBeat(t *testing.T) {
	k := 3
	n := New(0, defaultHeartbeat, defaultElection)
	n.sm = newStateMachine(0, []int{0, 1, 2})

	n.Step(Message{Type: msgHup}) // become leader please
	for _, m := range n.Msgs() {
		if m.Type == msgVote {
			n.Step(Message{From: 1, Type: msgVoteResp, Index: 1, Term: 1})
		}
	}

	for i := 0; i < defaultHeartbeat+1; i++ {
		n.Tick()
	}

	called := 0
	for _, m := range n.Msgs() {
		if m.Type == msgApp {
			called++
		}
	}

	// becomeLeader -> k-1 append
	// msgBeat -> k-1 append
	w := (k - 1) * 2
	if called != w {
		t.Errorf("called = %v, want %v", called, w)
	}
}

func TestResetElapse(t *testing.T) {
	tests := []struct {
		msg      Message
		welapsed tick
	}{
		{Message{From: 0, To: 1, Type: msgApp, Term: 2, Entries: []Entry{{Term: 1}}}, 0},
		{Message{From: 0, To: 1, Type: msgApp, Term: 1, Entries: []Entry{{Term: 1}}}, 1},
		{Message{From: 0, To: 1, Type: msgVote, Term: 2}, 0},
		{Message{From: 0, To: 1, Type: msgVote, Term: 1}, 1},
	}

	for i, tt := range tests {
		n := New(0, defaultHeartbeat, defaultElection)
		n.sm = newStateMachine(0, []int{0, 1, 2})
		n.sm.term = 2

		n.Tick()
		if n.elapsed != 1 {
			t.Errorf("%d: elpased = %d, want %d", i, n.elapsed, 1)
		}

		n.Step(tt.msg)
		if n.elapsed != tt.welapsed {
			t.Errorf("%d: elpased = %d, want %d", i, n.elapsed, tt.welapsed)
		}
	}
}

func TestStartCluster(t *testing.T) {
	n := Dictate(New(0, defaultHeartbeat, defaultElection))
	n.Next()

	if len(n.sm.ins) != 1 {
		t.Errorf("k = %d, want 1", len(n.sm.ins))
	}
	if n.sm.id != 0 {
		t.Errorf("addr = %d, want 0", n.sm.id)
	}
	if n.sm.state != stateLeader {
		t.Errorf("state = %s, want %s", n.sm.state, stateLeader)
	}
}

func TestAdd(t *testing.T) {
	n := Dictate(New(0, defaultHeartbeat, defaultElection))
	n.Next()
	n.Add(1)
	n.Next()

	if len(n.sm.ins) != 2 {
		t.Errorf("k = %d, want 2", len(n.sm.ins))
	}
	if n.sm.id != 0 {
		t.Errorf("addr = %d, want 0", n.sm.id)
	}
}

func TestRemove(t *testing.T) {
	n := Dictate(New(0, defaultHeartbeat, defaultElection))
	n.Next()
	n.Add(1)
	n.Next()
	n.Remove(0)
	n.Step(Message{Type: msgAppResp, From: 1, Term: 1, Index: 3})
	n.Next()

	if len(n.sm.ins) != 1 {
		t.Errorf("k = %d, want 1", len(n.sm.ins))
	}
	if n.sm.id != 0 {
		t.Errorf("addr = %d, want 0", n.sm.id)
	}
}
<commit_msg>raft: addr -> id<commit_after>package raft

import (
	"testing"
)

const (
	defaultHeartbeat = 1
	defaultElection  = 5
)

func TestTickMsgHup(t *testing.T) {
	n := New(0, defaultHeartbeat, defaultElection)
	n.sm = newStateMachine(0, []int{0, 1, 2})

	for i := 0; i < defaultElection+1; i++ {
		n.Tick()
	}

	called := false
	for _, m := range n.Msgs() {
		if m.Type == msgVote {
			called = true
		}
	}

	if !called {
		t.Errorf("called = %v, want true", called)
	}
}

func TestTickMsgBeat(t *testing.T) {
	k := 3
	n := New(0, defaultHeartbeat, defaultElection)
	n.sm = newStateMachine(0, []int{0, 1, 2})

	n.Step(Message{Type: msgHup}) // become leader please
	for _, m := range n.Msgs() {
		if m.Type == msgVote {
			n.Step(Message{From: 1, Type: msgVoteResp, Index: 1, Term: 1})
		}
	}

	for i := 0; i < defaultHeartbeat+1; i++ {
		n.Tick()
	}

	called := 0
	for _, m := range n.Msgs() {
		if m.Type == msgApp {
			called++
		}
	}

	// becomeLeader -> k-1 append
	// msgBeat -> k-1 append
	w := (k - 1) * 2
	if called != w {
		t.Errorf("called = %v, want %v", called, w)
	}
}

func TestResetElapse(t *testing.T) {
	tests := []struct {
		msg      Message
		welapsed tick
	}{
		{Message{From: 0, To: 1, Type: msgApp, Term: 2, Entries: []Entry{{Term: 1}}}, 0},
		{Message{From: 0, To: 1, Type: msgApp, Term: 1, Entries: []Entry{{Term: 1}}}, 1},
		{Message{From: 0, To: 1, Type: msgVote, Term: 2}, 0},
		{Message{From: 0, To: 1, Type: msgVote, Term: 1}, 1},
	}

	for i, tt := range tests {
		n := New(0, defaultHeartbeat, defaultElection)
		n.sm = newStateMachine(0, []int{0, 1, 2})
		n.sm.term = 2

		n.Tick()
		if n.elapsed != 1 {
			t.Errorf("%d: elpased = %d, want %d", i, n.elapsed, 1)
		}

		n.Step(tt.msg)
		if n.elapsed != tt.welapsed {
			t.Errorf("%d: elpased = %d, want %d", i, n.elapsed, tt.welapsed)
		}
	}
}

func TestStartCluster(t *testing.T) {
	n := Dictate(New(0, defaultHeartbeat, defaultElection))
	n.Next()

	if len(n.sm.ins) != 1 {
		t.Errorf("k = %d, want 1", len(n.sm.ins))
	}
	if n.sm.id != 0 {
		t.Errorf("id = %d, want 0", n.sm.id)
	}
	if n.sm.state != stateLeader {
		t.Errorf("state = %s, want %s", n.sm.state, stateLeader)
	}
}

func TestAdd(t *testing.T) {
	n := Dictate(New(0, defaultHeartbeat, defaultElection))
	n.Next()
	n.Add(1)
	n.Next()

	if len(n.sm.ins) != 2 {
		t.Errorf("k = %d, want 2", len(n.sm.ins))
	}
	if n.sm.id != 0 {
		t.Errorf("id = %d, want 0", n.sm.id)
	}
}

func TestRemove(t *testing.T) {
	n := Dictate(New(0, defaultHeartbeat, defaultElection))
	n.Next()
	n.Add(1)
	n.Next()
	n.Remove(0)
	n.Step(Message{Type: msgAppResp, From: 1, Term: 1, Index: 3})
	n.Next()

	if len(n.sm.ins) != 1 {
		t.Errorf("k = %d, want 1", len(n.sm.ins))
	}
	if n.sm.id != 0 {
		t.Errorf("id = %d, want 0", n.sm.id)
	}
}
<|endoftext|>


<commit_before>package rainrpc

import (
	"encoding/base64"
	"errors"
	"net"
	"net/rpc"
	"net/rpc/jsonrpc"
	"strconv"
	"strings"

	"github.com/cenkalti/rain/client"
	"github.com/cenkalti/rain/internal/logger"
)

type Server struct {
	config    ServerConfig
	rpcServer *rpc.Server
	log       logger.Logger
}

type ServerConfig struct {
	Host   string
	Port   int
	Client client.Config
}

var DefaultServerConfig = ServerConfig{
	Host:   "127.0.0.1",
	Port:   7246,
	Client: client.DefaultConfig,
}

func NewServer(cfg ServerConfig) (*Server, error) {
	clt, err := client.New(cfg.Client)
	if err != nil {
		return nil, err
	}
	h := &handler{client: clt}
	srv := rpc.NewServer()
	srv.RegisterName("Client", h)
	srv.HandleHTTP(rpc.DefaultRPCPath, rpc.DefaultDebugPath)
	return &Server{
		config:    cfg,
		rpcServer: srv,
		log:       logger.New("rpc server"),
	}, nil
}

func (s *Server) ListenAndServe() error {
	addr := net.JoinHostPort(s.config.Host, strconv.Itoa(s.config.Port))
	s.log.Infoln("RPC server is listening on", addr)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	for {
		conn, err := listener.Accept()
		if err != nil {
			return err
		}
		go s.rpcServer.ServeCodec(jsonrpc.NewServerCodec(conn))
	}
}

type handler struct {
	client *client.Client
}

func (h *handler) ListTorrents(args *ListTorrentsRequest, reply *ListTorrentsResponse) error {
	torrents := h.client.ListTorrents()
	reply.Torrents = make([]Torrent, 0, len(torrents))
	for _, t := range torrents {
		reply.Torrents = append(reply.Torrents, newTorrent(t))
	}
	return nil
}

func (h *handler) AddTorrent(args *AddTorrentRequest, reply *AddTorrentResponse) error {
	r := base64.NewDecoder(base64.StdEncoding, strings.NewReader(args.Torrent))
	t, err := h.client.AddTorrent(r)
	if err != nil {
		return err
	}
	reply.Torrent = newTorrent(t)
	return nil
}

func (h *handler) AddMagnet(args *AddMagnetRequest, reply *AddMagnetResponse) error {
	t, err := h.client.AddMagnet(args.Magnet)
	if err != nil {
		return err
	}
	reply.Torrent = newTorrent(t)
	return nil
}

func newTorrent(t *client.Torrent) Torrent {
	return Torrent{
		ID:       t.ID(),
		Name:     t.Name(),
		InfoHash: t.InfoHash(),
		Port:     t.Port(),
	}
}

func (h *handler) RemoveTorrent(args *RemoveTorrentRequest, reply *RemoveTorrentResponse) error {
	h.client.RemoveTorrent(args.ID)
	return nil
}

func (h *handler) GetTorrentStats(args *GetTorrentStatsRequest, reply *GetTorrentStatsResponse) error {
	t := h.client.GetTorrent(args.ID)
	if t == nil {
		return errors.New("torrent not found")
	}
	reply.Stats = t.Stats()
	return nil
}

func (h *handler) StartTorrent(args *StartTorrentRequest, reply *StartTorrentResponse) error {
	t := h.client.GetTorrent(args.ID)
	if t == nil {
		return errors.New("torrent not found")
	}
	t.Start()
	return nil
}

func (h *handler) StopTorrent(args *StopTorrentRequest, reply *StopTorrentResponse) error {
	t := h.client.GetTorrent(args.ID)
	if t == nil {
		return errors.New("torrent not found")
	}
	t.Stop()
	return nil
}
<commit_msg>rpc server does not exit on accept error<commit_after>package rainrpc

import (
	"encoding/base64"
	"errors"
	"net"
	"net/rpc"
	"net/rpc/jsonrpc"
	"strconv"
	"strings"

	"github.com/cenkalti/rain/client"
	"github.com/cenkalti/rain/internal/logger"
)

type Server struct {
	config    ServerConfig
	rpcServer *rpc.Server
	log       logger.Logger
}

type ServerConfig struct {
	Host   string
	Port   int
	Client client.Config
}

var DefaultServerConfig = ServerConfig{
	Host:   "127.0.0.1",
	Port:   7246,
	Client: client.DefaultConfig,
}

func NewServer(cfg ServerConfig) (*Server, error) {
	clt, err := client.New(cfg.Client)
	if err != nil {
		return nil, err
	}
	h := &handler{client: clt}
	srv := rpc.NewServer()
	srv.RegisterName("Client", h)
	srv.HandleHTTP(rpc.DefaultRPCPath, rpc.DefaultDebugPath)
	return &Server{
		config:    cfg,
		rpcServer: srv,
		log:       logger.New("rpc server"),
	}, nil
}

func (s *Server) ListenAndServe() error {
	addr := net.JoinHostPort(s.config.Host, strconv.Itoa(s.config.Port))
	s.log.Infoln("RPC server is listening on", addr)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	for {
		conn, err := listener.Accept()
		if err != nil {
			s.log.Errorln("accept error:", err)
			continue
		}
		go s.rpcServer.ServeCodec(jsonrpc.NewServerCodec(conn))
	}
}

type handler struct {
	client *client.Client
}

func (h *handler) ListTorrents(args *ListTorrentsRequest, reply *ListTorrentsResponse) error {
	torrents := h.client.ListTorrents()
	reply.Torrents = make([]Torrent, 0, len(torrents))
	for _, t := range torrents {
		reply.Torrents = append(reply.Torrents, newTorrent(t))
	}
	return nil
}

func (h *handler) AddTorrent(args *AddTorrentRequest, reply *AddTorrentResponse) error {
	r := base64.NewDecoder(base64.StdEncoding, strings.NewReader(args.Torrent))
	t, err := h.client.AddTorrent(r)
	if err != nil {
		return err
	}
	reply.Torrent = newTorrent(t)
	return nil
}

func (h *handler) AddMagnet(args *AddMagnetRequest, reply *AddMagnetResponse) error {
	t, err := h.client.AddMagnet(args.Magnet)
	if err != nil {
		return err
	}
	reply.Torrent = newTorrent(t)
	return nil
}

func newTorrent(t *client.Torrent) Torrent {
	return Torrent{
		ID:       t.ID(),
		Name:     t.Name(),
		InfoHash: t.InfoHash(),
		Port:     t.Port(),
	}
}

func (h *handler) RemoveTorrent(args *RemoveTorrentRequest, reply *RemoveTorrentResponse) error {
	h.client.RemoveTorrent(args.ID)
	return nil
}

func (h *handler) GetTorrentStats(args *GetTorrentStatsRequest, reply *GetTorrentStatsResponse) error {
	t := h.client.GetTorrent(args.ID)
	if t == nil {
		return errors.New("torrent not found")
	}
	reply.Stats = t.Stats()
	return nil
}

func (h *handler) StartTorrent(args *StartTorrentRequest, reply *StartTorrentResponse) error {
	t := h.client.GetTorrent(args.ID)
	if t == nil {
		return errors.New("torrent not found")
	}
	t.Start()
	return nil
}

func (h *handler) StopTorrent(args *StopTorrentRequest, reply *StopTorrentResponse) error {
	t := h.client.GetTorrent(args.ID)
	if t == nil {
		return errors.New("torrent not found")
	}
	t.Stop()
	return nil
}
<|endoftext|>


<commit_before>package db

import (
	"time"

	"github.com/couchbase/sync_gateway/base"
	"golang.org/x/net/context"
)

type BackgroundTaskFunc func(ctx context.Context) error

// backgroundTask runs task at the specified time interval in its own goroutine until stopped or an error is thrown by
// the BackgroundTaskFunc
func NewBackgroundTask(taskName string, dbName string, task BackgroundTaskFunc, interval time.Duration, c chan bool) {
	base.Infof(base.KeyAll, "Created background task: %q with interval %v", taskName, interval)
	go func() {
		for {
			select {
			case <-time.After(interval):
				ctx := context.WithValue(context.Background(), base.LogContextKey{}, base.LogContext{CorrelationID: base.NewTaskID(dbName, taskName)})
				base.DebugfCtx(ctx, base.KeyAll, "Running background task")
				if err := task(ctx); err != nil {
					base.ErrorfCtx(ctx, base.KeyAll, "Background task returned error: %v", err)
					return
				}
			case <-c:
				base.Debugf(base.KeyAll, "Terminating background task: %q", taskName)
				return
			}
		}
	}()
}
<commit_msg>Context package changed (#4102)<commit_after>package db

import (
	"context"
	"time"

	"github.com/couchbase/sync_gateway/base"
)

type BackgroundTaskFunc func(ctx context.Context) error

// backgroundTask runs task at the specified time interval in its own goroutine until stopped or an error is thrown by
// the BackgroundTaskFunc
func NewBackgroundTask(taskName string, dbName string, task BackgroundTaskFunc, interval time.Duration, c chan bool) {
	base.Infof(base.KeyAll, "Created background task: %q with interval %v", taskName, interval)
	go func() {
		for {
			select {
			case <-time.After(interval):
				ctx := context.WithValue(context.Background(), base.LogContextKey{}, base.LogContext{CorrelationID: base.NewTaskID(dbName, taskName)})
				base.DebugfCtx(ctx, base.KeyAll, "Running background task")
				if err := task(ctx); err != nil {
					base.ErrorfCtx(ctx, base.KeyAll, "Background task returned error: %v", err)
					return
				}
			case <-c:
				base.Debugf(base.KeyAll, "Terminating background task: %q", taskName)
				return
			}
		}
	}()
}
<|endoftext|>


<commit_before>// Package router provides a network routing control plane
package router

import (
	"time"
)

var (
	// DefaultAddress is default router address
	DefaultAddress = ":9093"
	// DefaultName is default router service name
	DefaultName = "go.micro.router"
	// DefaultNetwork is default micro network
	DefaultNetwork = "go.micro"
	// DefaultRouter is default network router
	DefaultRouter = NewRouter()
)

// Router is an interface for a routing control plane
type Router interface {
	// Init initializes the router with options
	Init(...Option) error
	// Options returns the router options
	Options() Options
	// The routing table
	Table() Table
	// Advertise advertises routes to the network
	Advertise() (<-chan *Advert, error)
	// Process processes incoming adverts
	Process(*Advert) error
	// Solicit advertises the whole routing table to the network
	Solicit() error
	// Lookup queries routes in the routing table
	Lookup(...QueryOption) ([]Route, error)
	// Watch returns a watcher which tracks updates to the routing table
	Watch(opts ...WatchOption) (Watcher, error)
	// Start starts the router
	Start() error
	// Status returns router status
	Status() Status
	// Stop stops the router
	Stop() error
	// Returns the router implementation
	String() string
}

// Table is an interface for routing table
type Table interface {
	// Create new route in the routing table
	Create(Route) error
	// Delete existing route from the routing table
	Delete(Route) error
	// Update route in the routing table
	Update(Route) error
	// List all routes in the table
	List() ([]Route, error)
	// Query routes in the routing table
	Query(...QueryOption) ([]Route, error)
}

// Option used by the router
type Option func(*Options)

// StatusCode defines router status
type StatusCode int

const (
	// Running means the router is up and running
	Running StatusCode = iota
	// Advertising means the router is advertising
	Advertising
	// Stopped means the router has been stopped
	Stopped
	// Error means the router has encountered error
	Error
)

func (s StatusCode) String() string {
	switch s {
	case Running:
		return "running"
	case Advertising:
		return "advertising"
	case Stopped:
		return "stopped"
	case Error:
		return "error"
	default:
		return "unknown"
	}
}

// Status is router status
type Status struct {
	// Code defines router status
	Code StatusCode
	// Error contains error description
	Error error
}

// String returns human readable status
func (s Status) String() string {
	return s.Code.String()
}

// AdvertType is route advertisement type
type AdvertType int

const (
	// Announce is advertised when the router announces itself
	Announce AdvertType = iota
	// RouteUpdate advertises route updates
	RouteUpdate
	// Solicitation indicates routes were solicited
	Solicitation
)

// String returns human readable advertisement type
func (t AdvertType) String() string {
	switch t {
	case Announce:
		return "announce"
	case RouteUpdate:
		return "update"
	case Solicitation:
		return "solicitation"
	default:
		return "unknown"
	}
}

// Advert contains a list of events advertised by the router to the network
type Advert struct {
	// Id is the router Id
	Id string
	// Type is type of advert
	Type AdvertType
	// Timestamp marks the time when the update is sent
	Timestamp time.Time
	// TTL is Advert TTL
	TTL time.Duration
	// Events is a list of routing table events to advertise
	Events []*Event
}

// Strategy is route advertisement strategy
type Strategy int

const (
	// AdvertiseAll advertises all routes to the network
	AdvertiseAll Strategy = iota
	// AdvertiseBest advertises optimal routes to the network
	AdvertiseBest
	// AdvertiseLocal will only advertise the local routes
	AdvertiseLocal
	// AdvertiseNone will not advertise any routes
	AdvertiseNone
)

// String returns human readable Strategy
func (s Strategy) String() string {
	switch s {
	case AdvertiseAll:
		return "all"
	case AdvertiseBest:
		return "best"
	case AdvertiseLocal:
		return "local"
	case AdvertiseNone:
		return "none"
	default:
		return "unknown"
	}
}

// NewRouter creates new Router and returns it
func NewRouter(opts ...Option) Router {
	return newRouter(opts...)
}
<commit_msg>Update router comments<commit_after>// Package router provides a network routing control plane
package router

import (
	"time"
)

var (
	// DefaultAddress is default router address
	DefaultAddress = ":9093"
	// DefaultName is default router service name
	DefaultName = "go.micro.router"
	// DefaultNetwork is default micro network
	DefaultNetwork = "go.micro"
	// DefaultRouter is default network router
	DefaultRouter = NewRouter()
)

// Router is an interface for a routing control plane
type Router interface {
	// Init initializes the router with options
	Init(...Option) error
	// Options returns the router options
	Options() Options
	// The routing table
	Table() Table
	// Advertise advertises routes
	Advertise() (<-chan *Advert, error)
	// Process processes incoming adverts
	Process(*Advert) error
	// Solicit advertises the whole routing table
	Solicit() error
	// Lookup queries routes in the routing table
	Lookup(...QueryOption) ([]Route, error)
	// Watch returns a watcher which tracks updates to the routing table
	Watch(opts ...WatchOption) (Watcher, error)
	// Start starts the router
	Start() error
	// Status returns router status
	Status() Status
	// Stop stops the router
	Stop() error
	// Returns the router implementation
	String() string
}

// Table is an interface for routing table
type Table interface {
	// Create new route in the routing table
	Create(Route) error
	// Delete existing route from the routing table
	Delete(Route) error
	// Update route in the routing table
	Update(Route) error
	// List all routes in the table
	List() ([]Route, error)
	// Query routes in the routing table
	Query(...QueryOption) ([]Route, error)
}

// Option used by the router
type Option func(*Options)

// StatusCode defines router status
type StatusCode int

const (
	// Running means the router is up and running
	Running StatusCode = iota
	// Advertising means the router is advertising
	Advertising
	// Stopped means the router has been stopped
	Stopped
	// Error means the router has encountered error
	Error
)

func (s StatusCode) String() string {
	switch s {
	case Running:
		return "running"
	case Advertising:
		return "advertising"
	case Stopped:
		return "stopped"
	case Error:
		return "error"
	default:
		return "unknown"
	}
}

// Status is router status
type Status struct {
	// Code defines router status
	Code StatusCode
	// Error contains error description
	Error error
}

// String returns human readable status
func (s Status) String() string {
	return s.Code.String()
}

// AdvertType is route advertisement type
type AdvertType int

const (
	// Announce is advertised when the router announces itself
	Announce AdvertType = iota
	// RouteUpdate advertises route updates
	RouteUpdate
	// Solicitation indicates routes were solicited
	Solicitation
)

// String returns human readable advertisement type
func (t AdvertType) String() string {
	switch t {
	case Announce:
		return "announce"
	case RouteUpdate:
		return "update"
	case Solicitation:
		return "solicitation"
	default:
		return "unknown"
	}
}

// Advert contains a list of events advertised by the router to the network
type Advert struct {
	// Id is the router Id
	Id string
	// Type is type of advert
	Type AdvertType
	// Timestamp marks the time when the update is sent
	Timestamp time.Time
	// TTL is Advert TTL
	TTL time.Duration
	// Events is a list of routing table events to advertise
	Events []*Event
}

// Strategy is route advertisement strategy
type Strategy int

const (
	// AdvertiseAll advertises all routes to the network
	AdvertiseAll Strategy = iota
	// AdvertiseBest advertises optimal routes to the network
	AdvertiseBest
	// AdvertiseLocal will only advertise the local routes
	AdvertiseLocal
	// AdvertiseNone will not advertise any routes
	AdvertiseNone
)

// String returns human readable Strategy
func (s Strategy) String() string {
	switch s {
	case AdvertiseAll:
		return "all"
	case AdvertiseBest:
		return "best"
	case AdvertiseLocal:
		return "local"
	case AdvertiseNone:
		return "none"
	default:
		return "unknown"
	}
}

// NewRouter creates new Router and returns it
func NewRouter(opts ...Option) Router {
	return newRouter(opts...)
}
<|endoftext|>


<commit_before>package timing

import (
	"fmt"
	"sort"
	"sync/atomic"
	"time"
)

var percentiles = []int{90, 95, 99}

func New() *Timing {
	t := &Timing{
		begin:     time.Now(),
		results:   make(chan record, 10),
		collected: make(map[string][]record),
	}
	go t.collector()
	return t
}

type Timing struct {
	begin     time.Time
	results   chan record
	done      int32
	collected map[string][]record
}

func (t *Timing) collector() {
	for result := range t.results {
		t.collected[result.key] = append(t.collected[result.key], result)
	}
	atomic.AddInt32(&t.done, 1)
}

func (t *Timing) Log(elapsed time.Duration, key string) {
	t.results <- record{
		key:     key,
		elapsed: elapsed,
	}
}

func (t *Timing) LogSince(begin time.Time, key string) {
	t.Log(time.Since(begin), key)
}

func (t *Timing) Report() Report {
	dest := make(Report)
	t.ReportInto(dest)
	return dest
}

func (t *Timing) ReportInto(dest Report) {
	dest["elapsed"] = time.Since(t.begin).String()
	if atomic.LoadInt32(&t.done) == 0 {
		close(t.results)
	}
	for k, records := range t.collected {
		sort.Sort(recordsByTime(records))

		// calculate mean
		var total time.Duration
		for _, record := range records {
			total += record.elapsed
		}
		dest[fmt.Sprintf("%s_avg", k)] = (total / time.Duration(len(records))).String()

		// calculate all the percentiles
		for _, p := range percentiles {
			// TODO do we need a ceiling operator?
			target := len(records) * 100 / p
			dest[fmt.Sprintf("%s_%dth", k, p)] = records[target].elapsed.String()
		}
	}
}

type Report map[string]interface{}

type record struct {
	key     string
	elapsed time.Duration
}

type recordsByTime []record

func (r recordsByTime) Len() int           { return len(r) }
func (r recordsByTime) Less(i, j int) bool { return r[i].elapsed < r[j].elapsed }
func (r recordsByTime) Swap(i, j int)      { r[i], r[j] = r[j], r[i] }
<commit_msg>Minor fix<commit_after>package timing

import (
	"fmt"
	"sort"
	"sync/atomic"
	"time"
)

var percentiles = []int{90, 95, 99}

func New() *Timing {
	t := &Timing{
		begin:     time.Now(),
		results:   make(chan record, 10),
		collected: make(map[string][]record),
	}
	go t.collector()
	return t
}

type Timing struct {
	begin     time.Time
	results   chan record
	done      int32
	collected map[string][]record
}

func (t *Timing) collector() {
	for result := range t.results {
		t.collected[result.key] = append(t.collected[result.key], result)
	}
	atomic.AddInt32(&t.done, 1)
}

func (t *Timing) Log(elapsed time.Duration, key string) {
	t.results <- record{
		key:     key,
		elapsed: elapsed,
	}
}

func (t *Timing) LogSince(begin time.Time, key string) {
	t.Log(time.Since(begin), key)
}

func (t *Timing) Report() Report {
	dest := make(Report)
	t.ReportInto(dest)
	return dest
}

func (t *Timing) ReportInto(dest Report) {
	dest["elapsed"] = time.Since(t.begin).String()
	if atomic.LoadInt32(&t.done) == 0 {
		close(t.results)
	}
	for k, records := range t.collected {
		sort.Sort(recordsByTime(records))

		// calculate mean
		var total time.Duration
		for _, record := range records {
			total += record.elapsed
		}
		dest[fmt.Sprintf("%s_avg", k)] = (total / time.Duration(len(records))).String()

		// calculate all the percentiles
		for _, p := range percentiles {
			// TODO do we need a ceiling operator?
			target := len(records) * p / 100
			dest[fmt.Sprintf("%s_%dth", k, p)] = records[target].elapsed.String()
		}
	}
}

type Report map[string]interface{}

type record struct {
	key     string
	elapsed time.Duration
}

type recordsByTime []record

func (r recordsByTime) Len() int           { return len(r) }
func (r recordsByTime) Less(i, j int) bool { return r[i].elapsed < r[j].elapsed }
func (r recordsByTime) Swap(i, j int)      { r[i], r[j] = r[j], r[i] }
<|endoftext|>


<commit_before>package director

import (
	"time"
)

const (
	FOREVER = -1
)

// A Looper is used in place of a direct call to "for {}" and implements some
// controls over how the loop will be run. The Loop() function is the main
// call used by dependant routines. Common patterns like Quit and Done
// channels are easily implemented in a Looper.
type Looper interface {
	Loop(fn func() error)
	Wait() error
	Done(err error)
	Quit()
}

// A TimedLooper is a Looper that runs on a timed schedule, using a Timer
// underneath. It also implements Quit and Done channels to allow external
// routines to more easily control and synchronize the loop.
//
// If you pass in a DoneChan at creation time, it will send a true on the
// channel when the loop has completed. If you pass in a quitChan it can
// be used to tell the loop to exit on completion of the next iteration.
// Using both Done and Quit channels allows full external control of the
// loop.
type TimedLooper struct {
	Count int
	Interval time.Duration
	DoneChan chan error
	quitChan chan bool
}

func NewTimedLooper(count int, interval time.Duration, done chan error) *TimedLooper {
	return &TimedLooper{count, interval, done, nil}
}

func (l *TimedLooper) Wait() error {
	return <-l.DoneChan
}

// Signal a dependant routine that we're done with our work
func (l *TimedLooper) Done(err error) {
	if l.DoneChan != nil {
		l.DoneChan <-err
	}
}

// The main method of the Looper. This call takes a function with a single
// return value, an error. If the error is nil, the Looper will run the
// next iteration. If it's an error, it will not run the next iteration,
// will clean up any internals that need to be, and will invoke done().
func (l *TimedLooper) Loop(fn func() error) {
	if l.quitChan == nil {
		l.quitChan = make(chan bool)
	}

	i := 0
	ticks := time.Tick(l.Interval)
	for range ticks {

		err := fn()
		if err != nil {
			l.Done(err)
			return
		}

		// We have to make sure not to increment if we started
		// at -1 otherwise we quit on maxint rollover.
		if l.Count != FOREVER {
			i = i + 1
			if i >= l.Count {
				l.Done(nil)
				return
			}
		}

		select {
		case <-l.quitChan:
			l.Done(nil)
			return
		default:
		}
	}
}

// Quit() signals to the Looper to not run the next iteration and to
// call done() and return as quickly as possible. It is does not
// intervene between iterations.
func (l *TimedLooper) Quit() {
	go func() {
		l.quitChan <-true
	}()
}

// A FreeLooper is like a TimedLooper but doesn't wait between iterations.
type FreeLooper struct {
	Count int
	DoneChan chan error
	quitChan chan bool
}

func NewFreeLooper(count int, done chan error) *FreeLooper {
	return &FreeLooper{count, done, nil}
}

func (l *FreeLooper) Wait() error {
	return <-l.DoneChan
}

// This is used internally, but can also be used by controlling routines
// to signal that a job is completed. The FreeLooper doesn's support its
// use outside the internals.
func (l *FreeLooper) Done(err error) {
	if l.DoneChan != nil {
		l.DoneChan <-err
	}
}

func (l *FreeLooper) Loop(fn func() error) {
	if l.quitChan == nil {
		l.quitChan = make(chan bool)
	}

	i := 0
	for {
		err := fn()
		if err != nil {
			l.Done(err)
			return
		}

		// We have to make sure not to increment if we started
		// at -1 otherwise we quit on maxint rollover.
		if l.Count != FOREVER {
			i = i + 1
			if i >= l.Count {
				l.Done(nil)
				return
			}
		}

		select {
		case <-l.quitChan:
			l.Done(nil)
			return
		default:
		}
	}
}

// Quit() signals to the Looper to not run the next iteration and to
// call Done() and return as quickly as possible. It is does not
// intervene between iterations.
func (l *FreeLooper) Quit() {
	go func() {
		l.quitChan <-true
	}()
}
<commit_msg>Added ONCE constant.<commit_after>package director

import (
	"time"
)

const (
	FOREVER = -1
	ONCE = 1
)

// A Looper is used in place of a direct call to "for {}" and implements some
// controls over how the loop will be run. The Loop() function is the main
// call used by dependant routines. Common patterns like Quit and Done
// channels are easily implemented in a Looper.
type Looper interface {
	Loop(fn func() error)
	Wait() error
	Done(err error)
	Quit()
}

// A TimedLooper is a Looper that runs on a timed schedule, using a Timer
// underneath. It also implements Quit and Done channels to allow external
// routines to more easily control and synchronize the loop.
//
// If you pass in a DoneChan at creation time, it will send a true on the
// channel when the loop has completed. If you pass in a quitChan it can
// be used to tell the loop to exit on completion of the next iteration.
// Using both Done and Quit channels allows full external control of the
// loop.
type TimedLooper struct {
	Count int
	Interval time.Duration
	DoneChan chan error
	quitChan chan bool
}

func NewTimedLooper(count int, interval time.Duration, done chan error) *TimedLooper {
	return &TimedLooper{count, interval, done, nil}
}

func (l *TimedLooper) Wait() error {
	return <-l.DoneChan
}

// Signal a dependant routine that we're done with our work
func (l *TimedLooper) Done(err error) {
	if l.DoneChan != nil {
		l.DoneChan <-err
	}
}

// The main method of the Looper. This call takes a function with a single
// return value, an error. If the error is nil, the Looper will run the
// next iteration. If it's an error, it will not run the next iteration,
// will clean up any internals that need to be, and will invoke done().
func (l *TimedLooper) Loop(fn func() error) {
	if l.quitChan == nil {
		l.quitChan = make(chan bool)
	}

	i := 0
	ticks := time.Tick(l.Interval)
	for range ticks {

		err := fn()
		if err != nil {
			l.Done(err)
			return
		}

		// We have to make sure not to increment if we started
		// at -1 otherwise we quit on maxint rollover.
		if l.Count != FOREVER {
			i = i + 1
			if i >= l.Count {
				l.Done(nil)
				return
			}
		}

		select {
		case <-l.quitChan:
			l.Done(nil)
			return
		default:
		}
	}
}

// Quit() signals to the Looper to not run the next iteration and to
// call done() and return as quickly as possible. It is does not
// intervene between iterations.
func (l *TimedLooper) Quit() {
	go func() {
		l.quitChan <-true
	}()
}

// A FreeLooper is like a TimedLooper but doesn't wait between iterations.
type FreeLooper struct {
	Count int
	DoneChan chan error
	quitChan chan bool
}

func NewFreeLooper(count int, done chan error) *FreeLooper {
	return &FreeLooper{count, done, nil}
}

func (l *FreeLooper) Wait() error {
	return <-l.DoneChan
}

// This is used internally, but can also be used by controlling routines
// to signal that a job is completed. The FreeLooper doesn's support its
// use outside the internals.
func (l *FreeLooper) Done(err error) {
	if l.DoneChan != nil {
		l.DoneChan <-err
	}
}

func (l *FreeLooper) Loop(fn func() error) {
	if l.quitChan == nil {
		l.quitChan = make(chan bool)
	}

	i := 0
	for {
		err := fn()
		if err != nil {
			l.Done(err)
			return
		}

		// We have to make sure not to increment if we started
		// at -1 otherwise we quit on maxint rollover.
		if l.Count != FOREVER {
			i = i + 1
			if i >= l.Count {
				l.Done(nil)
				return
			}
		}

		select {
		case <-l.quitChan:
			l.Done(nil)
			return
		default:
		}
	}
}

// Quit() signals to the Looper to not run the next iteration and to
// call Done() and return as quickly as possible. It is does not
// intervene between iterations.
func (l *FreeLooper) Quit() {
	go func() {
		l.quitChan <-true
	}()
}
<|endoftext|>


<commit_before>package muta

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

type FileInfo struct {
	Name string
	Path string
	Ctx  *interface{}
}

type SrcOpts struct {
	ReadSize uint
}

type Streamer func(*FileInfo, []byte) (*FileInfo, []byte, error)

// A convenience function to let functions that return Streamers
// "return an error". Ie, the following syntax:
//
// ```golang
// err := doSomething()
// if err != nil {
//   return ErrorStreamer(err)
// }
// ```
//
// ErrorStreamer will simply return a Streamer that will return an
// error when called.
func ErrorStreamer(err error) Streamer {
	return func(fi *FileInfo, chunk []byte) (*FileInfo, []byte, error) {
		return fi, chunk, err
	}
}

func SrcStreamer(ps []string, opts SrcOpts) Streamer {
	if opts.ReadSize == 0 {
		opts.ReadSize = 50
	}

	// Setup our channels
	fi := make(chan *FileInfo)
	chunk := make(chan []byte)
	err := make(chan error)
	read := make(chan bool)

	// This method of reading files needs to be abstracted further
	// to ensure that the file closing is deferred. In this
	// implementation i can't think of a way to test that.
	// Also, moving it out would let us ensure closing of the files
	// in tests
	go func() {
		sendErr := func(_fi *FileInfo, _chunk []byte, _err error) {
			<-read
			fi <- _fi
			chunk <- _chunk
			err <- _err
		}

		loadFile := func(p string) error {
			pchunks := make([]byte, opts.ReadSize)
			pfi := &FileInfo{
				Name: filepath.Base(p),
				Path: filepath.Dir(p),
			}

			f, ferr := os.Open(p)
			defer f.Close()
			if ferr != nil {
				sendErr(pfi, nil, ferr)
				return ferr
			}

			// Wait for a read request
			for <-read {
				// Read
				count, ferr := f.Read(pchunks)
				if ferr != nil && ferr == io.EOF {
					break
				}

				// Send
				fi <- pfi
				chunk <- pchunks[0:count]
				err <- ferr
				if ferr != nil {
					return ferr
				}
			}

			// The for loop stopped, send EOF
			fi <- pfi
			chunk <- nil
			err <- nil
			return nil
		}

		// Go through the paths and globbify any globbed paths
		globbedPaths := []string{}
		for _, p := range ps {
			// If it hs a *, it is a glob
			if strings.Contains(p, "*") {
				expandedGlobs, err := filepath.Glob(p)
				if err != nil {
					sendErr(nil, nil, err)
					return
				}
				globbedPaths = append(globbedPaths, expandedGlobs...)
			} else {
				globbedPaths = append(globbedPaths, p)
			}
		}

		for _, p := range globbedPaths {
			err := loadFile(p)
			if err != nil {
				return
			}
		}

		<-read
		// send EOS
		fi <- nil
		chunk <- nil
		err <- nil
	}()

	return func(inFi *FileInfo, inC []byte) (*FileInfo, []byte, error) {
		// If there is an incoming file pass the data along unmodified. This
		// func doesn't care to modify the data in any way
		if inFi != nil {
			return inFi, inC, nil
		}

		read <- true
		return <-fi, <-chunk, <-err
	}
}

type DestOpts struct {
	// Not implemented
	Clean bool
	// Not implemented
	Overwrite bool
}

func Dest(d string, args ...interface{}) Streamer {
	var opts DestOpts
	if len(args) == 0 {
		opts = DestOpts{
			Clean:     false,
			Overwrite: true,
		}
	} else if len(args) == 1 {
		_opts, ok := args[0].(DestOpts)
		opts = _opts
		if !ok {
			return ErrorStreamer(errors.New(
				"Unrecognized type in Dest(string, ...interface{}). " +
					"Use DestOpts()",
			))
		}
	}

	if opts.Clean {
		err := os.RemoveAll(d)
		if err != nil {
			return ErrorStreamer(err)
		}
	}

	// Make the destination if needed
	if err := os.MkdirAll(d, 0755); err != nil {
		return ErrorStreamer(err)
	}

	// A staging variable for the currently working file.
	var f *os.File
	return func(fi *FileInfo, chunk []byte) (*FileInfo, []byte, error) {
		// If fi is nil, then this func is now the generator. Dest() has no
		// need to generate, so signal EOS
		if fi == nil {
			return nil, chunk, nil
		}

		destPath := filepath.Join(d, fi.Path)
		destFilepath := filepath.Join(destPath, fi.Name)
		// MkdirAll checks if the given path is a dir, and exists. So
		// i believe there is no reason for us to bother checking.
		err := os.MkdirAll(destPath, 0755)
		if err != nil {
			return fi, chunk, err
		}

		if chunk == nil && f != nil {
			// f is open for writing, but chunk is nil, we're at EOF.
			// Close f, and set it to nil
			err = f.Close()
			f = nil
			return nil, nil, err
		}

		if f == nil {
			osFi, err := os.Stat(destFilepath)
			if err == nil && osFi.IsDir() {
				return fi, chunk, errors.New(fmt.Sprintf(
					"Cannot write to '%s', path is directory.",
					destFilepath,
				))
			}

			// This area is a bit of a cluster f*ck. In short:
			//
			// 1. If there is an error, and the error is that the file
			// does not exist, create it.
			// 2. If it's not a file does not exist error, return it.
			// 3. If there is no error, and the filepath is a directory,
			// return an error.
			// 4. If it's not a directory, and we're not allowed to overwrite
			// it, return an error.
			// 5. If we are allowed to overwrite it, open it up.
			//
			// Did i drink too much while writing this? It feels so messy.
			if err != nil {
				if os.IsNotExist(err) {
					f, err = os.Create(destFilepath)
					if err != nil {
						// Failed to create file, return
						return fi, chunk, err
					}
				} else {
					// Stat() error is unknown, return
					return fi, chunk, err
				}
			} else {
				// There was no error Stating path, it exist
				if osFi.IsDir() {
					// The file path is a dir, return error
					return fi, chunk, errors.New(fmt.Sprintf(
						"Cannot write to '%s', path is directory.",
						destFilepath,
					))
				} else if !opts.Overwrite {
					// We're not allowed to overwrite. Return error.
					return fi, chunk, errors.New(fmt.Sprintf(
						"Cannot write to '%s', path exists and Overwrite is set "+
							"to false.",
						destFilepath,
					))
				} else {
					f, err = os.Open(destFilepath)
					if err != nil {
						// Failed to open file for writing.
						return fi, chunk, err
					}
				}
			}
		}

		// lenth written can be ignored, because Write() returns an error
		// if len(chunk) != n
		_, err = f.Write(chunk)

		// Return EOS always. Dest() writes everything, like a boss..?
		return nil, nil, err
	}
}
<commit_msg>Changed Ctx to map[string]interface{}<commit_after>package muta

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

type FileInfo struct {
	Name string
	Path string
	Ctx  map[string]interface{}
}

type SrcOpts struct {
	ReadSize uint
}

type Streamer func(*FileInfo, []byte) (*FileInfo, []byte, error)

// A convenience function to let functions that return Streamers
// "return an error". Ie, the following syntax:
//
// ```golang
// err := doSomething()
// if err != nil {
//   return ErrorStreamer(err)
// }
// ```
//
// ErrorStreamer will simply return a Streamer that will return an
// error when called.
func ErrorStreamer(err error) Streamer {
	return func(fi *FileInfo, chunk []byte) (*FileInfo, []byte, error) {
		return fi, chunk, err
	}
}

func SrcStreamer(ps []string, opts SrcOpts) Streamer {
	if opts.ReadSize == 0 {
		opts.ReadSize = 50
	}

	// Setup our channels
	fi := make(chan *FileInfo)
	chunk := make(chan []byte)
	err := make(chan error)
	read := make(chan bool)

	// This method of reading files needs to be abstracted further
	// to ensure that the file closing is deferred. In this
	// implementation i can't think of a way to test that.
	// Also, moving it out would let us ensure closing of the files
	// in tests
	go func() {
		sendErr := func(_fi *FileInfo, _chunk []byte, _err error) {
			<-read
			fi <- _fi
			chunk <- _chunk
			err <- _err
		}

		loadFile := func(p string) error {
			pchunks := make([]byte, opts.ReadSize)
			pfi := &FileInfo{
				Name: filepath.Base(p),
				Path: filepath.Dir(p),
			}

			f, ferr := os.Open(p)
			defer f.Close()
			if ferr != nil {
				sendErr(pfi, nil, ferr)
				return ferr
			}

			// Wait for a read request
			for <-read {
				// Read
				count, ferr := f.Read(pchunks)
				if ferr != nil && ferr == io.EOF {
					break
				}

				// Send
				fi <- pfi
				chunk <- pchunks[0:count]
				err <- ferr
				if ferr != nil {
					return ferr
				}
			}

			// The for loop stopped, send EOF
			fi <- pfi
			chunk <- nil
			err <- nil
			return nil
		}

		// Go through the paths and globbify any globbed paths
		globbedPaths := []string{}
		for _, p := range ps {
			// If it hs a *, it is a glob
			if strings.Contains(p, "*") {
				expandedGlobs, err := filepath.Glob(p)
				if err != nil {
					sendErr(nil, nil, err)
					return
				}
				globbedPaths = append(globbedPaths, expandedGlobs...)
			} else {
				globbedPaths = append(globbedPaths, p)
			}
		}

		for _, p := range globbedPaths {
			err := loadFile(p)
			if err != nil {
				return
			}
		}

		<-read
		// send EOS
		fi <- nil
		chunk <- nil
		err <- nil
	}()

	return func(inFi *FileInfo, inC []byte) (*FileInfo, []byte, error) {
		// If there is an incoming file pass the data along unmodified. This
		// func doesn't care to modify the data in any way
		if inFi != nil {
			return inFi, inC, nil
		}

		read <- true
		return <-fi, <-chunk, <-err
	}
}

type DestOpts struct {
	// Not implemented
	Clean bool
	// Not implemented
	Overwrite bool
}

func Dest(d string, args ...interface{}) Streamer {
	var opts DestOpts
	if len(args) == 0 {
		opts = DestOpts{
			Clean:     false,
			Overwrite: true,
		}
	} else if len(args) == 1 {
		_opts, ok := args[0].(DestOpts)
		opts = _opts
		if !ok {
			return ErrorStreamer(errors.New(
				"Unrecognized type in Dest(string, ...interface{}). " +
					"Use DestOpts()",
			))
		}
	}

	if opts.Clean {
		err := os.RemoveAll(d)
		if err != nil {
			return ErrorStreamer(err)
		}
	}

	// Make the destination if needed
	if err := os.MkdirAll(d, 0755); err != nil {
		return ErrorStreamer(err)
	}

	// A staging variable for the currently working file.
	var f *os.File
	return func(fi *FileInfo, chunk []byte) (*FileInfo, []byte, error) {
		// If fi is nil, then this func is now the generator. Dest() has no
		// need to generate, so signal EOS
		if fi == nil {
			return nil, chunk, nil
		}

		destPath := filepath.Join(d, fi.Path)
		destFilepath := filepath.Join(destPath, fi.Name)
		// MkdirAll checks if the given path is a dir, and exists. So
		// i believe there is no reason for us to bother checking.
		err := os.MkdirAll(destPath, 0755)
		if err != nil {
			return fi, chunk, err
		}

		if chunk == nil && f != nil {
			// f is open for writing, but chunk is nil, we're at EOF.
			// Close f, and set it to nil
			err = f.Close()
			f = nil
			return nil, nil, err
		}

		if f == nil {
			osFi, err := os.Stat(destFilepath)
			if err == nil && osFi.IsDir() {
				return fi, chunk, errors.New(fmt.Sprintf(
					"Cannot write to '%s', path is directory.",
					destFilepath,
				))
			}

			// This area is a bit of a cluster f*ck. In short:
			//
			// 1. If there is an error, and the error is that the file
			// does not exist, create it.
			// 2. If it's not a file does not exist error, return it.
			// 3. If there is no error, and the filepath is a directory,
			// return an error.
			// 4. If it's not a directory, and we're not allowed to overwrite
			// it, return an error.
			// 5. If we are allowed to overwrite it, open it up.
			//
			// Did i drink too much while writing this? It feels so messy.
			if err != nil {
				if os.IsNotExist(err) {
					f, err = os.Create(destFilepath)
					if err != nil {
						// Failed to create file, return
						return fi, chunk, err
					}
				} else {
					// Stat() error is unknown, return
					return fi, chunk, err
				}
			} else {
				// There was no error Stating path, it exist
				if osFi.IsDir() {
					// The file path is a dir, return error
					return fi, chunk, errors.New(fmt.Sprintf(
						"Cannot write to '%s', path is directory.",
						destFilepath,
					))
				} else if !opts.Overwrite {
					// We're not allowed to overwrite. Return error.
					return fi, chunk, errors.New(fmt.Sprintf(
						"Cannot write to '%s', path exists and Overwrite is set "+
							"to false.",
						destFilepath,
					))
				} else {
					f, err = os.Open(destFilepath)
					if err != nil {
						// Failed to open file for writing.
						return fi, chunk, err
					}
				}
			}
		}

		// lenth written can be ignored, because Write() returns an error
		// if len(chunk) != n
		_, err = f.Write(chunk)

		// Return EOS always. Dest() writes everything, like a boss..?
		return nil, nil, err
	}
}
<|endoftext|>


<commit_before>// Package wkhtmltoimage provides a simple wrapper around wkhtmltoimage (http://wkhtmltopdf.org/) binary.
package wkhtmltoimage

import (
	"bytes"
	"errors"
	"image/jpeg"
	"image/png"
	"os/exec"
	"strconv"
	"strings"
)

// ImageOptions represent the options to generate the image.
type ImageOptions struct {
	// BinaryPrefix any prefix command to wkhtmltoimage
	//
	// Any command to prefex wkhtmltoimage (e.g. xvfb-run)
	BinaryPrefix string
	// BinaryPath the path to your wkhtmltoimage binary. REQUIRED
	//
	// Must be absolute path e.g /usr/local/bin/wkhtmltoimage
	BinaryPath string
	// Input is the content to turn into an image. REQUIRED
	//
	// Can be a url (http://example.com), a local file (/tmp/example.html), or html as a string (send "-" and set the Html value)
	Input string
	// Format is the type of image to generate
	//
	// jpg, png, svg, bmp supported. Defaults to local wkhtmltoimage default
	Format string
	// Height is the height of the screen used to render in pixels.
	//
	// Default is calculated from page content. Default 0 (renders entire page top to bottom)
	Height int
	// Width is the width of the screen used to render in pixels.
	//
	// Note that this is used only as a guide line. Default 1024
	Width int
	// Quality determines the final image quality.
	//
	// Values supported between 1 and 100. Default is 94
	Quality int
	// Html is a string of html to render into and image.
	//
	// Only needed to be set if Input is set to "-"
	Html string
	// Output controls how to save or return the image.
	//
	// Leave nil to return a []byte of the image. Set to a path (/tmp/example.png) to save as a file.
	Output string
}

// GenerateImage creates an image from an input.
// It returns the image ([]byte) and any error encountered.
func GenerateImage(options *ImageOptions) ([]byte, error) {
	arr, err := buildParams(options)
	if err != nil {
		return []byte{}, err
	}

	if options.BinaryPath == "" {
		return []byte{}, errors.New("BinaryPath not set")
	}

	cmd := new(exec.Cmd)
	if options.BinaryPrefix != "" {
		arr = append([]string{options.BinaryPath}, arr...)
		cmd = exec.Command(options.BinaryPrefix, arr...)
	} else {
		cmd = exec.Command(options.BinaryPath, arr...)
	}

	if options.Html != "" {
		cmd.Stdin = strings.NewReader(options.Html)
	}

	output, err := cmd.CombinedOutput()

	trimmed := cleanupOutput(output, options.Format)

	return trimmed, err
}

// buildParams takes the image options set by the user and turns them into command flags for wkhtmltoimage
// It returns an array of command flags.
func buildParams(options *ImageOptions) ([]string, error) {
	a := []string{}

	if options.BinaryPath == "" {
		return []string{}, errors.New("BinaryPath not set")
	}

	// silence extra wkhtmltoimage output
	// might want to add --javascript-delay too?
	a = append(a, "-q")
	a = append(a, "--disable-plugins")

	a = append(a, "--format")
	if options.Format != "" {
		a = append(a, options.Format)
	} else {
		a = append(a, "png")
	}

	if options.Height != 0 {
		a = append(a, "--height")
		a = append(a, strconv.Itoa(options.Height))
	}

	if options.Width != 0 {
		a = append(a, "--width")
		a = append(a, strconv.Itoa(options.Width))
	}

	if options.Quality != 0 {
		a = append(a, "--quality")
		a = append(a, strconv.Itoa(options.Quality))
	}

	// url and output come last
	if options.Input != "-" {
		// make sure we dont pass stdin if we aren't expecting it
		options.Html = ""
	}

	a = append(a, options.Input)

	if options.Output == "" {
		a = append(a, "-")
	} else {
		a = append(a, options.Output)
	}

	return a, nil
}

func cleanupOutput(img []byte, format string) []byte {
	buf := new(bytes.Buffer)
	switch {
	case format == "png":
		decoded, err := png.Decode(bytes.NewReader(img))
		for err != nil {
			img = img[1:]
			if len(img) == 0 {
				break
			}
			decoded, err = png.Decode(bytes.NewReader(img))
		}
		png.Encode(buf, decoded)
		return buf.Bytes()
	case format == "jpg":
		decoded, err := jpeg.Decode(bytes.NewReader(img))
		for err != nil {
			img = img[1:]
			if len(img) == 0 {
				break
			}
			decoded, err = jpeg.Decode(bytes.NewReader(img))
		}
		jpeg.Encode(buf, decoded, nil)
		return buf.Bytes()
		// case format == "svg":
		// 	return img
	}
	return img
}
<commit_msg>add binary prefix option<commit_after>// Package wkhtmltoimage provides a simple wrapper around wkhtmltoimage (http://wkhtmltopdf.org/) binary.
package wkhtmltoimage

import (
	// "bytes"
	"errors"
	// "image/jpeg"
	// "image/png"
	"os/exec"
	"strconv"
	"strings"
)

// ImageOptions represent the options to generate the image.
type ImageOptions struct {
	// BinaryPrefix any prefix command to wkhtmltoimage
	//
	// Any command to prefex wkhtmltoimage (e.g. xvfb-run)
	BinaryPrefix string
	// BinaryPath the path to your wkhtmltoimage binary. REQUIRED
	//
	// Must be absolute path e.g /usr/local/bin/wkhtmltoimage
	BinaryPath string
	// Input is the content to turn into an image. REQUIRED
	//
	// Can be a url (http://example.com), a local file (/tmp/example.html), or html as a string (send "-" and set the Html value)
	Input string
	// Format is the type of image to generate
	//
	// jpg, png, svg, bmp supported. Defaults to local wkhtmltoimage default
	Format string
	// Height is the height of the screen used to render in pixels.
	//
	// Default is calculated from page content. Default 0 (renders entire page top to bottom)
	Height int
	// Width is the width of the screen used to render in pixels.
	//
	// Note that this is used only as a guide line. Default 1024
	Width int
	// Quality determines the final image quality.
	//
	// Values supported between 1 and 100. Default is 94
	Quality int
	// Html is a string of html to render into and image.
	//
	// Only needed to be set if Input is set to "-"
	Html string
	// Output controls how to save or return the image.
	//
	// Leave nil to return a []byte of the image. Set to a path (/tmp/example.png) to save as a file.
	Output string
}

// GenerateImage creates an image from an input.
// It returns the image ([]byte) and any error encountered.
func GenerateImage(options *ImageOptions) ([]byte, error) {
	arr, err := buildParams(options)
	if err != nil {
		return []byte{}, err
	}

	if options.BinaryPath == "" {
		return []byte{}, errors.New("BinaryPath not set")
	}

	cmd := new(exec.Cmd)
	if options.BinaryPrefix != "" {
		arr = append([]string{options.BinaryPath}, arr...)
		cmd = exec.Command(options.BinaryPrefix, arr...)
	} else {
		cmd = exec.Command(options.BinaryPath, arr...)
	}

	if options.Html != "" {
		cmd.Stdin = strings.NewReader(options.Html)
	}

	output, err := cmd.CombinedOutput()

	trimmed := cleanupOutput(output, options.Format)

	return trimmed, err
}

// buildParams takes the image options set by the user and turns them into command flags for wkhtmltoimage
// It returns an array of command flags.
func buildParams(options *ImageOptions) ([]string, error) {
	a := []string{}

	if options.BinaryPath == "" {
		return []string{}, errors.New("BinaryPath not set")
	}

	// silence extra wkhtmltoimage output
	// might want to add --javascript-delay too?
	a = append(a, "-q")
	a = append(a, "--disable-plugins")

	a = append(a, "--format")
	if options.Format != "" {
		a = append(a, options.Format)
	} else {
		a = append(a, "png")
	}

	if options.Height != 0 {
		a = append(a, "--height")
		a = append(a, strconv.Itoa(options.Height))
	}

	if options.Width != 0 {
		a = append(a, "--width")
		a = append(a, strconv.Itoa(options.Width))
	}

	if options.Quality != 0 {
		a = append(a, "--quality")
		a = append(a, strconv.Itoa(options.Quality))
	}

	// url and output come last
	if options.Input != "-" {
		// make sure we dont pass stdin if we aren't expecting it
		options.Html = ""
	}

	a = append(a, options.Input)

	if options.Output == "" {
		a = append(a, "-")
	} else {
		a = append(a, options.Output)
	}

	return a, nil
}

func cleanupOutput(img []byte, format string) []byte {
	// buf := new(bytes.Buffer)
	// switch {
	// case format == "png":
	// 	decoded, err := png.Decode(bytes.NewReader(img))
	// 	for err != nil {
	// 		img = img[1:]
	// 		if len(img) == 0 {
	// 			break
	// 		}
	// 		decoded, err = png.Decode(bytes.NewReader(img))
	// 	}
	// 	png.Encode(buf, decoded)
	// 	return buf.Bytes()
	// case format == "jpg":
	// 	decoded, err := jpeg.Decode(bytes.NewReader(img))
	// 	for err != nil {
	// 		img = img[1:]
	// 		if len(img) == 0 {
	// 			break
	// 		}
	// 		decoded, err = jpeg.Decode(bytes.NewReader(img))
	// 	}
	// 	jpeg.Encode(buf, decoded, nil)
	// 	return buf.Bytes()
	// 	// case format == "svg":
	// 	// 	return img
	// }
	return img
}
<|endoftext|>


<commit_before>package workers

import (
	"fmt"
	"time"

	"github.com/Odinman/omq/utils"
	zmq "github.com/pebbe/zmq4"
)

type Node struct {
	identity  string    //  Identity of worker
	id_string string    //  Printable identity
	expire    time.Time //  Expires at this time
}

/* {{{ func newNode(identity string) node
 *  Construct new worker
 */
func newNode(identity string) Node {
	return Node{
		identity:  identity,
		id_string: fmt.Sprintf("%q", identity),
		expire:    time.Now().Add(HEARTBEAT_INTERVAL * HEARTBEAT_LIVENESS),
	}
}

/* }}} */

/* {{{ func nodeReady(self Node, nodes []Node) []Node
 * 把一个ready的节点放到队列最后
 */
func nodeReady(self Node, nodes []Node) []Node {
	for i, worker := range nodes {
		if self.identity == worker.identity {
			if i == 0 { //第一个
				nodes = nodes[1:]
			} else if i == len(nodes)-1 { //最后一个
				nodes = nodes[:i-1]
			} else { //中间
				nodes = append(nodes[:i], nodes[i+1:]...)
			}
			break
		}
	}
	return append(nodes, self)
}

/* }}} */

/* {{{ func purgeNodes(nodes []Node) []Node
 *  The purge method looks for and kills expired nodes. We hold nodes
 *  from oldest to most recent, so we stop at the first alive worker:
 */
func purgeNodes(nodes []Node) []Node {
	now := time.Now()
	for i, worker := range nodes {
		if now.Before(worker.expire) {
			return nodes[i:] //  Worker is alive, we're done here
		}
	}
	return nodes[0:0]
}

/* }}} */

/* {{{ func (w *OmqWorker) serve() {
 *
 */
func (w *OmqWorker) serve() {
	//  接受请求的socket
	frontend, _ := zmq.NewSocket(zmq.ROUTER)
	defer frontend.Close()
	frontend.Bind(fmt.Sprint("tcp://*:", basePort))
	w.Debug("frontend bind port: %v", basePort)

	// backend
	backend, _ := zmq.NewSocket(zmq.ROUTER)
	defer backend.Close()
	backend.Bind("inproc://backend")

	//  可用节点列表,LRU算法,最空的节点保持在队列最前
	nodes := make([]Node, 0)

	// 心跳
	heartbeat_at := time.Tick(HEARTBEAT_INTERVAL)

	poller1 := zmq.NewPoller()
	poller1.Add(backend, zmq.POLLIN)

	poller2 := zmq.NewPoller()
	poller2.Add(backend, zmq.POLLIN)
	poller2.Add(frontend, zmq.POLLIN)

	// spawn responser
	w.Info("create %d responsers", responseNodes)
	for i := 1; i <= responseNodes; i++ {
		go w.newResponser(i)
	}
	// loop
	for {
		//  Poll frontend only if we have available nodes
		var sockets []zmq.Polled
		var err error
		if nl := len(nodes); nl > 0 {
			//w.Info("nodes len: %d", nl)
			sockets, err = poller2.Poll(HEARTBEAT_INTERVAL)
		} else {
			w.Info("nodes empty")
			sockets, err = poller1.Poll(HEARTBEAT_INTERVAL)
		}
		if err != nil {
			w.Critical("big wrong: %s", err)
			break //  Interrupted
		}

		for _, socket := range sockets {
			switch socket.Socket {
			case backend:
				//  Handle worker activity on backend
				//  Use worker identity for load-balancing
				msg, err := backend.RecvMessage(0)
				if err != nil {
					w.Error("backend wrong: %s", err)
					break //  Interrupted
				}

				//  Any sign of life from worker means it's ready
				identity, msg := utils.Unwrap(msg)
				nodes = nodeReady(newNode(identity), nodes)

				//  Validate control message, or return reply to client
				if len(msg) == 1 {
					// 控制信息, ready or heartbeat
					if msg[0] != PPP_READY && msg[0] != PPP_HEARTBEAT {
						w.Info("Error: invalid message from worker, %s", msg)
					}
				} else {
					// 任务处理完毕的回复(带信封), 直接返回前台
					w.Trace("backend recv: %s", msg)
					frontend.SendMessage(msg)
				}
			case frontend:
				//  Now get next client request, route to next worker
				msg, err := frontend.RecvMessage(0)
				if err != nil {
					w.Error("frontend wrong: %s", err)
					break //  Interrupted
				}
				w.Trace("frontend recv: %s", msg)

				//定向发送到后台(带信封), 将来可以有多组后台, 分别处理不同的任务
				backend.SendMessage(nodes[0].identity, msg)
				nodes = nodes[1:]
			}
		}

		//  We handle heartbeating after any socket activity. First we send
		//  heartbeats to any idle nodes if it's time. Then we purge any
		//  dead nodes:

		select {
		case <-heartbeat_at:
			//向节点发送心跳
			for _, worker := range nodes {
				backend.SendMessage(worker.identity, PPP_HEARTBEAT)
			}
			//向订阅者发送心跳
			publisher.SendMessage(PPP_HEARTBEAT)
		default:
		}
		nodes = purgeNodes(nodes)

	}
}

/* }}} */
<commit_msg>bugfix<commit_after>package workers

import (
	"fmt"
	"time"

	"github.com/Odinman/omq/utils"
	zmq "github.com/pebbe/zmq4"
)

type Node struct {
	identity  string    //  Identity of worker
	id_string string    //  Printable identity
	expire    time.Time //  Expires at this time
}

/* {{{ func newNode(identity string) node
 *  Construct new worker
 */
func newNode(identity string) Node {
	return Node{
		identity:  identity,
		id_string: fmt.Sprintf("%q", identity),
		expire:    time.Now().Add(HEARTBEAT_INTERVAL * HEARTBEAT_LIVENESS),
	}
}

/* }}} */

/* {{{ func nodeReady(self Node, nodes []Node) []Node
 * 把一个ready的节点放到队列最后
 */
func nodeReady(self Node, nodes []Node) []Node {
	for i, worker := range nodes {
		if self.identity == worker.identity {
			if i == 0 { //第一个
				nodes = nodes[1:]
			} else if i == len(nodes)-1 { //最后一个
				nodes = nodes[:i]
			} else { //中间
				nodes = append(nodes[:i], nodes[i+1:]...)
			}
			break
		}
	}
	return append(nodes, self)
}

/* }}} */

/* {{{ func purgeNodes(nodes []Node) []Node
 *  The purge method looks for and kills expired nodes. We hold nodes
 *  from oldest to most recent, so we stop at the first alive worker:
 */
func purgeNodes(nodes []Node) []Node {
	now := time.Now()
	for i, worker := range nodes {
		if now.Before(worker.expire) {
			return nodes[i:] //  Worker is alive, we're done here
		}
	}
	return nodes[0:0]
}

/* }}} */

/* {{{ func (w *OmqWorker) serve() {
 *
 */
func (w *OmqWorker) serve() {
	//  接受请求的socket
	frontend, _ := zmq.NewSocket(zmq.ROUTER)
	defer frontend.Close()
	frontend.Bind(fmt.Sprint("tcp://*:", basePort))
	w.Debug("frontend bind port: %v", basePort)

	// backend
	backend, _ := zmq.NewSocket(zmq.ROUTER)
	defer backend.Close()
	backend.Bind("inproc://backend")

	//  可用节点列表,LRU算法,最空的节点保持在队列最前
	nodes := make([]Node, 0)

	// 心跳
	heartbeat_at := time.Tick(HEARTBEAT_INTERVAL)

	poller1 := zmq.NewPoller()
	poller1.Add(backend, zmq.POLLIN)

	poller2 := zmq.NewPoller()
	poller2.Add(backend, zmq.POLLIN)
	poller2.Add(frontend, zmq.POLLIN)

	// spawn responser
	w.Info("create %d responsers", responseNodes)
	for i := 1; i <= responseNodes; i++ {
		go w.newResponser(i)
	}
	// loop
	for {
		//  Poll frontend only if we have available nodes
		var sockets []zmq.Polled
		var err error
		if nl := len(nodes); nl > 0 {
			//w.Info("nodes len: %d", nl)
			sockets, err = poller2.Poll(HEARTBEAT_INTERVAL)
		} else {
			w.Info("nodes empty")
			sockets, err = poller1.Poll(HEARTBEAT_INTERVAL)
		}
		if err != nil {
			w.Critical("big wrong: %s", err)
			break //  Interrupted
		}

		for _, socket := range sockets {
			switch socket.Socket {
			case backend:
				//  Handle worker activity on backend
				//  Use worker identity for load-balancing
				msg, err := backend.RecvMessage(0)
				if err != nil {
					w.Error("backend wrong: %s", err)
					break //  Interrupted
				}

				//  Any sign of life from worker means it's ready
				identity, msg := utils.Unwrap(msg)
				nodes = nodeReady(newNode(identity), nodes)

				//  Validate control message, or return reply to client
				if len(msg) == 1 {
					// 控制信息, ready or heartbeat
					if msg[0] != PPP_READY && msg[0] != PPP_HEARTBEAT {
						w.Info("Error: invalid message from worker, %s", msg)
					}
				} else {
					// 任务处理完毕的回复(带信封), 直接返回前台
					w.Trace("backend recv: %s", msg)
					frontend.SendMessage(msg)
				}
			case frontend:
				//  Now get next client request, route to next worker
				msg, err := frontend.RecvMessage(0)
				if err != nil {
					w.Error("frontend wrong: %s", err)
					break //  Interrupted
				}
				w.Trace("frontend recv: %s", msg)

				//定向发送到后台(带信封), 将来可以有多组后台, 分别处理不同的任务
				backend.SendMessage(nodes[0].identity, msg)
				nodes = nodes[1:]
			}
		}

		//  We handle heartbeating after any socket activity. First we send
		//  heartbeats to any idle nodes if it's time. Then we purge any
		//  dead nodes:

		select {
		case <-heartbeat_at:
			//向节点发送心跳
			for _, worker := range nodes {
				backend.SendMessage(worker.identity, PPP_HEARTBEAT)
			}
			//向订阅者发送心跳
			publisher.SendMessage(PPP_HEARTBEAT)
		default:
		}
		nodes = purgeNodes(nodes)

	}
}

/* }}} */
<|endoftext|>


<commit_before>package main

import (
	"../util"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
)

const VERSION = "0.1"

var (
	showVersion     = flag.Bool("version", false, "print version string")
	bindAddress     = flag.String("address", "0.0.0.0", "address to bind to")
	webPort         = flag.String("web-port", "5151", "port to listen on for HTTP connections")
	tcpPort         = flag.String("tcp-port", "5150", "port to listen on for TCP connections")
	debugMode       = flag.Bool("debug", false, "enable debug mode")
	memQueueSize    = flag.Int("mem-queue-size", 10000, "number of messages to keep in memory (per topic)")
	cpuProfile      = flag.String("cpu-profile", "", "write cpu profile to file")
	goMaxProcs      = flag.Int("go-max-procs", 4, "runtime configuration for GOMAXPROCS")
	dataPath        = flag.String("data-path", "", "path to store disk-backed messages")
	lookupAddresses = util.StringArray{}
)

func init() {
	flag.Var(&lookupAddresses, "lookup-address", "lookup address (may be given multiple times)")
}

func main() {
	flag.Parse()

	if *showVersion {
		fmt.Printf("nsqd v%s\n", VERSION)
		return
	}

	if *goMaxProcs > 0 {
		runtime.GOMAXPROCS(*goMaxProcs)
	}

	endChan := make(chan int)
	signalChan := make(chan os.Signal, 1)

	if *cpuProfile != "" {
		log.Printf("CPU Profiling Enabled")
		f, err := os.Create(*cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	go func() {
		<-signalChan
		endChan <- 1
	}()
	signal.Notify(signalChan, os.Interrupt)

	tcpAddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(*bindAddress, *tcpPort))
	if err != nil {
		log.Fatal(err)
	}

	webAddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(*bindAddress, *webPort))
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("nsqd v%s", VERSION)

	// prepare the lookup peer instances
	lookupHosts := make([]string, 0)
	lookupHosts = append(lookupHosts, "127.0.0.1:5160")
	go LookupRouter(lookupHosts)

	go TopicFactory(*memQueueSize, *dataPath)
	go UuidFactory()

	tcpListener, err := net.Listen("tcp", tcpAddr.String())
	if err != nil {
		log.Fatalf("FATAL: listen (%s) failed - %s", tcpAddr.String(), err.Error())
	}
	go util.TcpServer(tcpListener, tcpClientHandler)

	webListener, err := net.Listen("tcp", webAddr.String())
	if err != nil {
		log.Fatalf("FATAL: listen (%s) failed - %s", webAddr.String(), err.Error())
	}
	go HttpServer(webListener)

	<-endChan

	for _, topic := range TopicMap {
		topic.Close()
	}
}
<commit_msg>default goMaxProcs to 0 so that it uses environment<commit_after>package main

import (
	"../util"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
)

const VERSION = "0.1"

var (
	showVersion     = flag.Bool("version", false, "print version string")
	bindAddress     = flag.String("address", "0.0.0.0", "address to bind to")
	webPort         = flag.String("web-port", "5151", "port to listen on for HTTP connections")
	tcpPort         = flag.String("tcp-port", "5150", "port to listen on for TCP connections")
	debugMode       = flag.Bool("debug", false, "enable debug mode")
	memQueueSize    = flag.Int("mem-queue-size", 10000, "number of messages to keep in memory (per topic)")
	cpuProfile      = flag.String("cpu-profile", "", "write cpu profile to file")
	goMaxProcs      = flag.Int("go-max-procs", 0, "runtime configuration for GOMAXPROCS")
	dataPath        = flag.String("data-path", "", "path to store disk-backed messages")
	lookupAddresses = util.StringArray{}
)

func init() {
	flag.Var(&lookupAddresses, "lookup-address", "lookup address (may be given multiple times)")
}

func main() {
	flag.Parse()

	if *showVersion {
		fmt.Printf("nsqd v%s\n", VERSION)
		return
	}

	if *goMaxProcs > 0 {
		runtime.GOMAXPROCS(*goMaxProcs)
	}

	endChan := make(chan int)
	signalChan := make(chan os.Signal, 1)

	if *cpuProfile != "" {
		log.Printf("CPU Profiling Enabled")
		f, err := os.Create(*cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	go func() {
		<-signalChan
		endChan <- 1
	}()
	signal.Notify(signalChan, os.Interrupt)

	tcpAddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(*bindAddress, *tcpPort))
	if err != nil {
		log.Fatal(err)
	}

	webAddr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(*bindAddress, *webPort))
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("nsqd v%s", VERSION)

	// prepare the lookup peer instances
	lookupHosts := make([]string, 0)
	lookupHosts = append(lookupHosts, "127.0.0.1:5160")
	go LookupRouter(lookupHosts)

	go TopicFactory(*memQueueSize, *dataPath)
	go UuidFactory()

	tcpListener, err := net.Listen("tcp", tcpAddr.String())
	if err != nil {
		log.Fatalf("FATAL: listen (%s) failed - %s", tcpAddr.String(), err.Error())
	}
	go util.TcpServer(tcpListener, tcpClientHandler)

	webListener, err := net.Listen("tcp", webAddr.String())
	if err != nil {
		log.Fatalf("FATAL: listen (%s) failed - %s", webAddr.String(), err.Error())
	}
	go HttpServer(webListener)

	<-endChan

	for _, topic := range TopicMap {
		topic.Close()
	}
}
<|endoftext|>