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

package kernel

import (
	"fmt"
	"io"
	"sync"

	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/context"
	"gvisor.dev/gvisor/pkg/coverage"
	"gvisor.dev/gvisor/pkg/safemem"
	"gvisor.dev/gvisor/pkg/sentry/memmap"
	"gvisor.dev/gvisor/pkg/sentry/mm"
	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
	"gvisor.dev/gvisor/pkg/sentry/usage"
	"gvisor.dev/gvisor/pkg/syserror"
	"gvisor.dev/gvisor/pkg/usermem"
)

// kcovAreaSizeMax is the maximum number of uint64 entries allowed in the kcov
// area. On Linux, the maximum is INT_MAX / 8.
const kcovAreaSizeMax = 10 * 1024 * 1024

// Kcov provides kernel coverage data to userspace through a memory-mapped
// region, as kcov does in Linux.
//
// To give the illusion that the data is always up to date, we update the shared
// memory every time before we return to userspace.
type Kcov struct {
	// mfp provides application memory. It is immutable after creation.
	mfp pgalloc.MemoryFileProvider

	// mu protects all of the fields below.
	mu sync.RWMutex

	// mode is the current kcov mode.
	mode uint8

	// size is the size of the mapping through which the kernel conveys coverage
	// information to userspace.
	size uint64

	// owningTask is the task that currently owns coverage data on the system. The
	// interface for kcov essentially requires that coverage is only going to a
	// single task. Note that kcov should only generate coverage data for the
	// owning task, but we currently generate global coverage.
	owningTask *Task

	// count is a locally cached version of the first uint64 in the kcov data,
	// which is the number of subsequent entries representing PCs.
	//
	// It is used with kcovInode.countBlock(), to copy in/out the first element of
	// the actual data in an efficient manner, avoid boilerplate, and prevent
	// accidental garbage escapes by the temporary counts.
	count uint64

	mappable *mm.SpecialMappable
}

// NewKcov creates and returns a Kcov instance.
func (k *Kernel) NewKcov() *Kcov {
	return &Kcov{
		mfp: k,
	}
}

var coveragePool = sync.Pool{
	New: func() interface{} {
		return make([]byte, 0)
	},
}

// TaskWork implements TaskWorker.TaskWork.
func (kcov *Kcov) TaskWork(t *Task) {
	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	if kcov.mode != linux.KCOV_MODE_TRACE_PC {
		return
	}

	rw := &kcovReadWriter{
		mf: kcov.mfp.MemoryFile(),
		fr: kcov.mappable.FileRange(),
	}

	// Read in the PC count.
	if _, err := safemem.ReadFullToBlocks(rw, kcov.countBlock()); err != nil {
		panic(fmt.Sprintf("Internal error reading count from kcov area: %v", err))
	}

	rw.off = 8 * (1 + kcov.count)
	n := coverage.ConsumeCoverageData(&kcovIOWriter{rw})

	// Update the pc count, based on the number of entries written. Note that if
	// we reached the end of the kcov area, we may not have written everything in
	// output.
	kcov.count += uint64(n / 8)
	rw.off = 0
	if _, err := safemem.WriteFullFromBlocks(rw, kcov.countBlock()); err != nil {
		panic(fmt.Sprintf("Internal error writing count to kcov area: %v", err))
	}

	// Re-register for future work.
	t.RegisterWork(kcov)
}

// InitTrace performs the KCOV_INIT_TRACE ioctl.
func (kcov *Kcov) InitTrace(size uint64) error {
	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	if kcov.mode != linux.KCOV_MODE_DISABLED {
		return syserror.EBUSY
	}

	// To simplify all the logic around mapping, we require that the length of the
	// shared region is a multiple of the system page size.
	if (8*size)&(usermem.PageSize-1) != 0 {
		return syserror.EINVAL
	}

	// We need space for at least two uint64s to hold current position and a
	// single PC.
	if size < 2 || size > kcovAreaSizeMax {
		return syserror.EINVAL
	}

	kcov.size = size
	kcov.mode = linux.KCOV_MODE_INIT
	return nil
}

// EnableTrace performs the KCOV_ENABLE_TRACE ioctl.
func (kcov *Kcov) EnableTrace(ctx context.Context, traceKind uint8) error {
	t := TaskFromContext(ctx)
	if t == nil {
		panic("kcovInode.EnableTrace() cannot be used outside of a task goroutine")
	}

	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	// KCOV_ENABLE must be preceded by KCOV_INIT_TRACE and an mmap call.
	if kcov.mode != linux.KCOV_MODE_INIT || kcov.mappable == nil {
		return syserror.EINVAL
	}

	switch traceKind {
	case linux.KCOV_TRACE_PC:
		kcov.mode = linux.KCOV_MODE_TRACE_PC
	case linux.KCOV_TRACE_CMP:
		// We do not support KCOV_MODE_TRACE_CMP.
		return syserror.ENOTSUP
	default:
		return syserror.EINVAL
	}

	if kcov.owningTask != nil && kcov.owningTask != t {
		return syserror.EBUSY
	}

	kcov.owningTask = t
	t.SetKcov(kcov)
	t.RegisterWork(kcov)

	// Clear existing coverage data; the task expects to read only coverage data
	// from the time it is activated.
	coverage.ClearCoverageData()
	return nil
}

// DisableTrace performs the KCOV_DISABLE_TRACE ioctl.
func (kcov *Kcov) DisableTrace(ctx context.Context) error {
	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	t := TaskFromContext(ctx)
	if t == nil {
		panic("kcovInode.EnableTrace() cannot be used outside of a task goroutine")
	}

	if t != kcov.owningTask {
		return syserror.EINVAL
	}
	kcov.mode = linux.KCOV_MODE_INIT
	kcov.owningTask = nil
	kcov.mappable = nil
	return nil
}

// Clear resets the mode and clears the owning task and memory mapping for kcov.
// It is called when the fd corresponding to kcov is closed. Note that the mode
// needs to be set so that the next call to kcov.TaskWork() will exit early.
func (kcov *Kcov) Clear() {
	kcov.mu.Lock()
	kcov.clearLocked()
	kcov.mu.Unlock()
}

func (kcov *Kcov) clearLocked() {
	kcov.mode = linux.KCOV_MODE_INIT
	kcov.owningTask = nil
	if kcov.mappable != nil {
		kcov.mappable = nil
	}
}

// OnTaskExit is called when the owning task exits. It is similar to
// kcov.Clear(), except the memory mapping is not cleared, so that the same
// mapping can be used in the future if kcov is enabled again by another task.
func (kcov *Kcov) OnTaskExit() {
	kcov.mu.Lock()
	kcov.mode = linux.KCOV_MODE_INIT
	kcov.owningTask = nil
	kcov.mu.Unlock()
}

// ConfigureMMap is called by the vfs.FileDescription for this kcov instance to
// implement vfs.FileDescription.ConfigureMMap.
func (kcov *Kcov) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	if kcov.mode != linux.KCOV_MODE_INIT {
		return syserror.EINVAL
	}

	if kcov.mappable == nil {
		// Set up the kcov area.
		fr, err := kcov.mfp.MemoryFile().Allocate(kcov.size*8, usage.Anonymous)
		if err != nil {
			return err
		}

		// Get the thread id for the mmap name.
		t := TaskFromContext(ctx)
		if t == nil {
			panic("ThreadFromContext returned nil")
		}
		// For convenience, a special mappable is used here. Note that these mappings
		// will look different under /proc/[pid]/maps than they do on Linux.
		kcov.mappable = mm.NewSpecialMappable(fmt.Sprintf("[kcov:%d]", t.ThreadID()), kcov.mfp, fr)
	}
	opts.Mappable = kcov.mappable
	opts.MappingIdentity = kcov.mappable
	return nil
}

// kcovReadWriter implements safemem.Reader and safemem.Writer.
type kcovReadWriter struct {
	off uint64
	mf  *pgalloc.MemoryFile
	fr  memmap.FileRange
}

// ReadToBlocks implements safemem.Reader.ReadToBlocks.
func (rw *kcovReadWriter) ReadToBlocks(dsts safemem.BlockSeq) (uint64, error) {
	if dsts.IsEmpty() {
		return 0, nil
	}

	// Limit the read to the kcov range and check for overflow.
	if rw.fr.Length() <= rw.off {
		return 0, io.EOF
	}
	start := rw.fr.Start + rw.off
	end := rw.fr.Start + rw.fr.Length()
	if rend := start + dsts.NumBytes(); rend < end {
		end = rend
	}

	// Get internal mappings.
	bs, err := rw.mf.MapInternal(memmap.FileRange{start, end}, usermem.Read)
	if err != nil {
		return 0, err
	}

	// Copy from internal mappings.
	n, err := safemem.CopySeq(dsts, bs)
	rw.off += n
	return n, err
}

// WriteFromBlocks implements safemem.Writer.WriteFromBlocks.
func (rw *kcovReadWriter) WriteFromBlocks(srcs safemem.BlockSeq) (uint64, error) {
	if srcs.IsEmpty() {
		return 0, nil
	}

	// Limit the write to the kcov area and check for overflow.
	if rw.fr.Length() <= rw.off {
		return 0, io.EOF
	}
	start := rw.fr.Start + rw.off
	end := rw.fr.Start + rw.fr.Length()
	if wend := start + srcs.NumBytes(); wend < end {
		end = wend
	}

	// Get internal mapping.
	bs, err := rw.mf.MapInternal(memmap.FileRange{start, end}, usermem.Write)
	if err != nil {
		return 0, err
	}

	// Copy to internal mapping.
	n, err := safemem.CopySeq(bs, srcs)
	rw.off += n
	return n, err
}

// kcovIOWriter implements io.Writer as a basic wrapper over kcovReadWriter.
type kcovIOWriter struct {
	rw *kcovReadWriter
}

// Write implements io.Writer.Write.
func (w *kcovIOWriter) Write(p []byte) (int, error) {
	bs := safemem.BlockSeqOf(safemem.BlockFromSafeSlice(p))
	n, err := safemem.WriteFullFromBlocks(w.rw, bs)
	return int(n), err
}
<commit_msg>Simplify nil assignment in kcov.<commit_after>// Copyright 2020 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package kernel

import (
	"fmt"
	"io"
	"sync"

	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/context"
	"gvisor.dev/gvisor/pkg/coverage"
	"gvisor.dev/gvisor/pkg/safemem"
	"gvisor.dev/gvisor/pkg/sentry/memmap"
	"gvisor.dev/gvisor/pkg/sentry/mm"
	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
	"gvisor.dev/gvisor/pkg/sentry/usage"
	"gvisor.dev/gvisor/pkg/syserror"
	"gvisor.dev/gvisor/pkg/usermem"
)

// kcovAreaSizeMax is the maximum number of uint64 entries allowed in the kcov
// area. On Linux, the maximum is INT_MAX / 8.
const kcovAreaSizeMax = 10 * 1024 * 1024

// Kcov provides kernel coverage data to userspace through a memory-mapped
// region, as kcov does in Linux.
//
// To give the illusion that the data is always up to date, we update the shared
// memory every time before we return to userspace.
type Kcov struct {
	// mfp provides application memory. It is immutable after creation.
	mfp pgalloc.MemoryFileProvider

	// mu protects all of the fields below.
	mu sync.RWMutex

	// mode is the current kcov mode.
	mode uint8

	// size is the size of the mapping through which the kernel conveys coverage
	// information to userspace.
	size uint64

	// owningTask is the task that currently owns coverage data on the system. The
	// interface for kcov essentially requires that coverage is only going to a
	// single task. Note that kcov should only generate coverage data for the
	// owning task, but we currently generate global coverage.
	owningTask *Task

	// count is a locally cached version of the first uint64 in the kcov data,
	// which is the number of subsequent entries representing PCs.
	//
	// It is used with kcovInode.countBlock(), to copy in/out the first element of
	// the actual data in an efficient manner, avoid boilerplate, and prevent
	// accidental garbage escapes by the temporary counts.
	count uint64

	mappable *mm.SpecialMappable
}

// NewKcov creates and returns a Kcov instance.
func (k *Kernel) NewKcov() *Kcov {
	return &Kcov{
		mfp: k,
	}
}

var coveragePool = sync.Pool{
	New: func() interface{} {
		return make([]byte, 0)
	},
}

// TaskWork implements TaskWorker.TaskWork.
func (kcov *Kcov) TaskWork(t *Task) {
	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	if kcov.mode != linux.KCOV_MODE_TRACE_PC {
		return
	}

	rw := &kcovReadWriter{
		mf: kcov.mfp.MemoryFile(),
		fr: kcov.mappable.FileRange(),
	}

	// Read in the PC count.
	if _, err := safemem.ReadFullToBlocks(rw, kcov.countBlock()); err != nil {
		panic(fmt.Sprintf("Internal error reading count from kcov area: %v", err))
	}

	rw.off = 8 * (1 + kcov.count)
	n := coverage.ConsumeCoverageData(&kcovIOWriter{rw})

	// Update the pc count, based on the number of entries written. Note that if
	// we reached the end of the kcov area, we may not have written everything in
	// output.
	kcov.count += uint64(n / 8)
	rw.off = 0
	if _, err := safemem.WriteFullFromBlocks(rw, kcov.countBlock()); err != nil {
		panic(fmt.Sprintf("Internal error writing count to kcov area: %v", err))
	}

	// Re-register for future work.
	t.RegisterWork(kcov)
}

// InitTrace performs the KCOV_INIT_TRACE ioctl.
func (kcov *Kcov) InitTrace(size uint64) error {
	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	if kcov.mode != linux.KCOV_MODE_DISABLED {
		return syserror.EBUSY
	}

	// To simplify all the logic around mapping, we require that the length of the
	// shared region is a multiple of the system page size.
	if (8*size)&(usermem.PageSize-1) != 0 {
		return syserror.EINVAL
	}

	// We need space for at least two uint64s to hold current position and a
	// single PC.
	if size < 2 || size > kcovAreaSizeMax {
		return syserror.EINVAL
	}

	kcov.size = size
	kcov.mode = linux.KCOV_MODE_INIT
	return nil
}

// EnableTrace performs the KCOV_ENABLE_TRACE ioctl.
func (kcov *Kcov) EnableTrace(ctx context.Context, traceKind uint8) error {
	t := TaskFromContext(ctx)
	if t == nil {
		panic("kcovInode.EnableTrace() cannot be used outside of a task goroutine")
	}

	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	// KCOV_ENABLE must be preceded by KCOV_INIT_TRACE and an mmap call.
	if kcov.mode != linux.KCOV_MODE_INIT || kcov.mappable == nil {
		return syserror.EINVAL
	}

	switch traceKind {
	case linux.KCOV_TRACE_PC:
		kcov.mode = linux.KCOV_MODE_TRACE_PC
	case linux.KCOV_TRACE_CMP:
		// We do not support KCOV_MODE_TRACE_CMP.
		return syserror.ENOTSUP
	default:
		return syserror.EINVAL
	}

	if kcov.owningTask != nil && kcov.owningTask != t {
		return syserror.EBUSY
	}

	kcov.owningTask = t
	t.SetKcov(kcov)
	t.RegisterWork(kcov)

	// Clear existing coverage data; the task expects to read only coverage data
	// from the time it is activated.
	coverage.ClearCoverageData()
	return nil
}

// DisableTrace performs the KCOV_DISABLE_TRACE ioctl.
func (kcov *Kcov) DisableTrace(ctx context.Context) error {
	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	t := TaskFromContext(ctx)
	if t == nil {
		panic("kcovInode.EnableTrace() cannot be used outside of a task goroutine")
	}

	if t != kcov.owningTask {
		return syserror.EINVAL
	}
	kcov.mode = linux.KCOV_MODE_INIT
	kcov.owningTask = nil
	kcov.mappable = nil
	return nil
}

// Clear resets the mode and clears the owning task and memory mapping for kcov.
// It is called when the fd corresponding to kcov is closed. Note that the mode
// needs to be set so that the next call to kcov.TaskWork() will exit early.
func (kcov *Kcov) Clear() {
	kcov.mu.Lock()
	kcov.clearLocked()
	kcov.mu.Unlock()
}

func (kcov *Kcov) clearLocked() {
	kcov.mode = linux.KCOV_MODE_INIT
	kcov.owningTask = nil
	kcov.mappable = nil
}

// OnTaskExit is called when the owning task exits. It is similar to
// kcov.Clear(), except the memory mapping is not cleared, so that the same
// mapping can be used in the future if kcov is enabled again by another task.
func (kcov *Kcov) OnTaskExit() {
	kcov.mu.Lock()
	kcov.mode = linux.KCOV_MODE_INIT
	kcov.owningTask = nil
	kcov.mu.Unlock()
}

// ConfigureMMap is called by the vfs.FileDescription for this kcov instance to
// implement vfs.FileDescription.ConfigureMMap.
func (kcov *Kcov) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) error {
	kcov.mu.Lock()
	defer kcov.mu.Unlock()

	if kcov.mode != linux.KCOV_MODE_INIT {
		return syserror.EINVAL
	}

	if kcov.mappable == nil {
		// Set up the kcov area.
		fr, err := kcov.mfp.MemoryFile().Allocate(kcov.size*8, usage.Anonymous)
		if err != nil {
			return err
		}

		// Get the thread id for the mmap name.
		t := TaskFromContext(ctx)
		if t == nil {
			panic("ThreadFromContext returned nil")
		}
		// For convenience, a special mappable is used here. Note that these mappings
		// will look different under /proc/[pid]/maps than they do on Linux.
		kcov.mappable = mm.NewSpecialMappable(fmt.Sprintf("[kcov:%d]", t.ThreadID()), kcov.mfp, fr)
	}
	opts.Mappable = kcov.mappable
	opts.MappingIdentity = kcov.mappable
	return nil
}

// kcovReadWriter implements safemem.Reader and safemem.Writer.
type kcovReadWriter struct {
	off uint64
	mf  *pgalloc.MemoryFile
	fr  memmap.FileRange
}

// ReadToBlocks implements safemem.Reader.ReadToBlocks.
func (rw *kcovReadWriter) ReadToBlocks(dsts safemem.BlockSeq) (uint64, error) {
	if dsts.IsEmpty() {
		return 0, nil
	}

	// Limit the read to the kcov range and check for overflow.
	if rw.fr.Length() <= rw.off {
		return 0, io.EOF
	}
	start := rw.fr.Start + rw.off
	end := rw.fr.Start + rw.fr.Length()
	if rend := start + dsts.NumBytes(); rend < end {
		end = rend
	}

	// Get internal mappings.
	bs, err := rw.mf.MapInternal(memmap.FileRange{start, end}, usermem.Read)
	if err != nil {
		return 0, err
	}

	// Copy from internal mappings.
	n, err := safemem.CopySeq(dsts, bs)
	rw.off += n
	return n, err
}

// WriteFromBlocks implements safemem.Writer.WriteFromBlocks.
func (rw *kcovReadWriter) WriteFromBlocks(srcs safemem.BlockSeq) (uint64, error) {
	if srcs.IsEmpty() {
		return 0, nil
	}

	// Limit the write to the kcov area and check for overflow.
	if rw.fr.Length() <= rw.off {
		return 0, io.EOF
	}
	start := rw.fr.Start + rw.off
	end := rw.fr.Start + rw.fr.Length()
	if wend := start + srcs.NumBytes(); wend < end {
		end = wend
	}

	// Get internal mapping.
	bs, err := rw.mf.MapInternal(memmap.FileRange{start, end}, usermem.Write)
	if err != nil {
		return 0, err
	}

	// Copy to internal mapping.
	n, err := safemem.CopySeq(bs, srcs)
	rw.off += n
	return n, err
}

// kcovIOWriter implements io.Writer as a basic wrapper over kcovReadWriter.
type kcovIOWriter struct {
	rw *kcovReadWriter
}

// Write implements io.Writer.Write.
func (w *kcovIOWriter) Write(p []byte) (int, error) {
	bs := safemem.BlockSeqOf(safemem.BlockFromSafeSlice(p))
	n, err := safemem.WriteFullFromBlocks(w.rw, bs)
	return int(n), err
}
<|endoftext|>


<commit_before>package forge

import (
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
)

// Section struct holds a map of values
type Section struct {
	comments []string
	includes []string
	parent   *Section
	values   map[string]Value
}

// NewSection will create and initialize a new Section
func NewSection() *Section {
	return &Section{
		comments: make([]string, 0),
		includes: make([]string, 0),
		values:   make(map[string]Value),
	}
}

func newChildSection(parent *Section) *Section {
	return &Section{
		comments: make([]string, 0),
		includes: make([]string, 0),
		parent:   parent,
		values:   make(map[string]Value),
	}
}

// AddComment will append a new comment into the section
func (section *Section) AddComment(comment string) {
	section.comments = append(section.comments, comment)
}

// AddInclude will append a new filename into the section
func (section *Section) AddInclude(filename string) {
	section.includes = append(section.includes, filename)
}

// GetComments will return all the comments were defined for this Section
func (section *Section) GetComments() []string {
	return section.comments
}

// GetIncludes will return the filenames of all the includes were parsed for this Section
func (section *Section) GetIncludes() []string {
	return section.includes
}

// GetType will respond with the ValueType of this Section (hint, always SECTION)
func (section *Section) GetType() ValueType {
	return SECTION
}

// GetValue retrieves the raw underlying value stored in this Section
func (section *Section) GetValue() interface{} {
	return section.values
}

// UpdateValue updates the raw underlying value stored in this Section
func (section *Section) UpdateValue(value interface{}) error {
	switch value.(type) {
	case map[string]Value:
		section.values = value.(map[string]Value)
		return nil
	}

	msg := fmt.Sprintf("unsupported type, %s must be of type `map[string]Value`", value)
	return errors.New(msg)
}

// AddSection adds a new child section to this Section with the provided name
func (section *Section) AddSection(name string) *Section {
	childSection := newChildSection(section)
	section.values[name] = childSection
	return childSection
}

// Exists returns true when a value stored under the key exists
func (section *Section) Exists(name string) bool {
	_, err := section.Get(name)
	return err == nil
}

// Get the value (Primative or Section) stored under the name
// will respond with an error if the value does not exist
func (section *Section) Get(name string) (Value, error) {
	value, ok := section.values[name]
	var err error
	if ok == false {
		err = errors.New("value does not exist")
	}
	return value, err
}

// GetBoolean will try to get the value stored under name as a bool
// will respond with an error if the value does not exist or cannot be converted to a bool
func (section *Section) GetBoolean(name string) (bool, error) {
	value, err := section.Get(name)
	if err != nil {
		return false, err
	}

	switch value.(type) {
	case *Primative:
		return value.(*Primative).AsBoolean()
	case *Section:
		return true, nil
	}

	return false, errors.New("could not convert unknown value to boolean")
}

// GetFloat will try to get the value stored under name as a float64
// will respond with an error if the value does not exist or cannot be converted to a float64
func (section *Section) GetFloat(name string) (float64, error) {
	value, err := section.Get(name)
	if err != nil {
		return float64(0), err
	}

	switch value.(type) {
	case *Primative:
		return value.(*Primative).AsFloat()
	}

	return float64(0), errors.New("could not convert non-primative value to float")
}

// GetInteger will try to get the value stored under name as a int64
// will respond with an error if the value does not exist or cannot be converted to a int64
func (section *Section) GetInteger(name string) (int64, error) {
	value, err := section.Get(name)
	if err != nil {
		return int64(0), err
	}

	switch value.(type) {
	case *Primative:
		return value.(*Primative).AsInteger()
	}

	return int64(0), errors.New("could not convert non-primative value to integer")
}

// GetList will try to get the value stored under name as a List
// will respond with an error if the value does not exist or is not a List
func (section *Section) GetList(name string) (*List, error) {
	value, err := section.Get(name)
	if err != nil {
		return nil, err
	}

	if value.GetType() == LIST {
		return value.(*List), nil
	}

	return nil, errors.New("could not fetch value as list")
}

// GetSection will try to get the value stored under name as a Section
// will respond with an error if the value does not exist or is not a Section
func (section *Section) GetSection(name string) (*Section, error) {
	value, err := section.Get(name)
	if err != nil {
		return nil, err
	}

	if value.GetType() == SECTION {
		return value.(*Section), nil
	}
	return nil, errors.New("could not fetch value as section")
}

// GetString will try to get the value stored under name as a string
// will respond with an error if the value does not exist or cannot be converted to a string
func (section *Section) GetString(name string) (string, error) {
	value, err := section.Get(name)
	if err != nil {
		return "", err
	}

	switch value.(type) {
	case *Primative:
		return value.(*Primative).AsString()
	}

	return "", errors.New("could not convert non-primative value to string")
}

// GetParent will get the parent section associated with this Section or nil
// if it does not have one
func (section *Section) GetParent() *Section {
	return section.parent
}

// HasParent will return true if this Section has a parent
func (section *Section) HasParent() bool {
	return section.parent != nil
}

// Keys will return back a list of all setting names in this Section
func (section *Section) Keys() []string {
	var keys []string
	for key := range section.values {
		keys = append(keys, key)
	}

	sort.Strings(keys)
	return keys
}

// Set will set a value (Primative or Section) to the provided name
func (section *Section) Set(name string, value Value) {
	section.values[name] = value
}

// SetBoolean will set the value for name as a bool
func (section *Section) SetBoolean(name string, value bool) {
	current, err := section.Get(name)

	// Exists just update the value/type
	if err == nil {
		current.UpdateValue(value)
	} else {
		section.values[name] = NewBoolean(value)
	}
}

// SetFloat will set the value for name as a float64
func (section *Section) SetFloat(name string, value float64) {
	current, err := section.Get(name)

	// Exists just update the value/type
	if err == nil {
		current.UpdateValue(value)
	} else {
		section.values[name] = NewFloat(value)
	}
}

// SetInteger will set the value for name as a int64
func (section *Section) SetInteger(name string, value int64) {
	current, err := section.Get(name)

	// Exists just update the value/type
	if err == nil {
		current.UpdateValue(value)
	} else {
		section.values[name] = NewInteger(value)
	}
}

// SetNull will set the value for name as nil
func (section *Section) SetNull(name string) {
	current, err := section.Get(name)

	// Already is a Null, nothing to do
	if err == nil && current.GetType() == NULL {
		return
	}
	section.Set(name, NewNull())
}

// SetString will set the value for name as a string
func (section *Section) SetString(name string, value string) {
	current, err := section.Get(name)

	// Exists just update the value/type
	if err == nil {
		current.UpdateValue(value)
	} else {
		section.Set(name, NewString(value))
	}
}

// Resolve will recursively try to fetch the provided value and will respond
// with an error if the name does not exist or tries to be resolved through
// a non-section value
func (section *Section) Resolve(name string) (Value, error) {
	// Used only in error state return value
	var value Value

	parts := strings.Split(name, ".")
	if len(parts) == 0 {
		return value, errors.New("no name provided")
	}

	var current Value
	current = section
	for _, part := range parts {
		if current.GetType() != SECTION {
			return value, errors.New("trying to resolve value from non-section")
		}

		nextCurrent, err := current.(*Section).Get(part)
		if err != nil {
			return value, errors.New("could not find value in section")
		}
		current = nextCurrent
	}
	return current, nil
}

// Merge merges the given section to current section. Settings from source
// section overwites the values in the current section
func (section *Section) Merge(source *Section) error {
	for _, key := range source.Keys() {
		sourceValue, _ := source.Get(key)
		targetValue, err := section.Get(key)

		// not found, so add it
		if err != nil {
			section.Set(key, sourceValue)
			continue
		}

		// found existing one and it's type SECTION, merge it
		if targetValue.GetType() == SECTION {
			// Source value have to be SECTION type here
			if sourceValue.GetType() != SECTION {
				return fmt.Errorf("source (%v) and target (%v) type doesn't match: %v",
					sourceValue.GetType(),
					targetValue.GetType(),
					key)
			}

			if err = targetValue.(*Section).Merge(sourceValue.(*Section)); err != nil {
				return err
			}

			continue
		}

		// found existing one, update it
		if err = targetValue.UpdateValue(sourceValue.GetValue()); err != nil {
			return fmt.Errorf("%v: %v", err, key)
		}
	}
	return nil
}

// ToJSON will convert this Section and all it's underlying values and Sections
// into JSON as a []byte
func (section *Section) ToJSON() ([]byte, error) {
	data := section.ToMap()
	return json.Marshal(data)
}

// ToMap will convert this Section and all it's underlying values and Sections into
// a map[string]interface{}
func (section *Section) ToMap() map[string]interface{} {
	output := make(map[string]interface{})

	for key, value := range section.values {
		if value.GetType() == SECTION {
			output[key] = value.(*Section).ToMap()
		} else {
			output[key] = value.GetValue()
		}
	}
	return output
}
<commit_msg>Expose 'value does not exists' error<commit_after>package forge

import (
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
)

var (
	// ErrNotExists represents a nonexistent value error
	ErrNotExists = errors.New("value does not exist")
)

// Section struct holds a map of values
type Section struct {
	comments []string
	includes []string
	parent   *Section
	values   map[string]Value
}

// NewSection will create and initialize a new Section
func NewSection() *Section {
	return &Section{
		comments: make([]string, 0),
		includes: make([]string, 0),
		values:   make(map[string]Value),
	}
}

func newChildSection(parent *Section) *Section {
	return &Section{
		comments: make([]string, 0),
		includes: make([]string, 0),
		parent:   parent,
		values:   make(map[string]Value),
	}
}

// AddComment will append a new comment into the section
func (section *Section) AddComment(comment string) {
	section.comments = append(section.comments, comment)
}

// AddInclude will append a new filename into the section
func (section *Section) AddInclude(filename string) {
	section.includes = append(section.includes, filename)
}

// GetComments will return all the comments were defined for this Section
func (section *Section) GetComments() []string {
	return section.comments
}

// GetIncludes will return the filenames of all the includes were parsed for this Section
func (section *Section) GetIncludes() []string {
	return section.includes
}

// GetType will respond with the ValueType of this Section (hint, always SECTION)
func (section *Section) GetType() ValueType {
	return SECTION
}

// GetValue retrieves the raw underlying value stored in this Section
func (section *Section) GetValue() interface{} {
	return section.values
}

// UpdateValue updates the raw underlying value stored in this Section
func (section *Section) UpdateValue(value interface{}) error {
	switch value.(type) {
	case map[string]Value:
		section.values = value.(map[string]Value)
		return nil
	}

	msg := fmt.Sprintf("unsupported type, %s must be of type `map[string]Value`", value)
	return errors.New(msg)
}

// AddSection adds a new child section to this Section with the provided name
func (section *Section) AddSection(name string) *Section {
	childSection := newChildSection(section)
	section.values[name] = childSection
	return childSection
}

// Exists returns true when a value stored under the key exists
func (section *Section) Exists(name string) bool {
	_, err := section.Get(name)
	return err == nil
}

// Get the value (Primative or Section) stored under the name
// will respond with an error if the value does not exist
func (section *Section) Get(name string) (Value, error) {
	value, ok := section.values[name]
	var err error
	if ok == false {
		err = ErrNotExists
	}
	return value, err
}

// GetBoolean will try to get the value stored under name as a bool
// will respond with an error if the value does not exist or cannot be converted to a bool
func (section *Section) GetBoolean(name string) (bool, error) {
	value, err := section.Get(name)
	if err != nil {
		return false, err
	}

	switch value.(type) {
	case *Primative:
		return value.(*Primative).AsBoolean()
	case *Section:
		return true, nil
	}

	return false, errors.New("could not convert unknown value to boolean")
}

// GetFloat will try to get the value stored under name as a float64
// will respond with an error if the value does not exist or cannot be converted to a float64
func (section *Section) GetFloat(name string) (float64, error) {
	value, err := section.Get(name)
	if err != nil {
		return float64(0), err
	}

	switch value.(type) {
	case *Primative:
		return value.(*Primative).AsFloat()
	}

	return float64(0), errors.New("could not convert non-primative value to float")
}

// GetInteger will try to get the value stored under name as a int64
// will respond with an error if the value does not exist or cannot be converted to a int64
func (section *Section) GetInteger(name string) (int64, error) {
	value, err := section.Get(name)
	if err != nil {
		return int64(0), err
	}

	switch value.(type) {
	case *Primative:
		return value.(*Primative).AsInteger()
	}

	return int64(0), errors.New("could not convert non-primative value to integer")
}

// GetList will try to get the value stored under name as a List
// will respond with an error if the value does not exist or is not a List
func (section *Section) GetList(name string) (*List, error) {
	value, err := section.Get(name)
	if err != nil {
		return nil, err
	}

	if value.GetType() == LIST {
		return value.(*List), nil
	}

	return nil, errors.New("could not fetch value as list")
}

// GetSection will try to get the value stored under name as a Section
// will respond with an error if the value does not exist or is not a Section
func (section *Section) GetSection(name string) (*Section, error) {
	value, err := section.Get(name)
	if err != nil {
		return nil, err
	}

	if value.GetType() == SECTION {
		return value.(*Section), nil
	}
	return nil, errors.New("could not fetch value as section")
}

// GetString will try to get the value stored under name as a string
// will respond with an error if the value does not exist or cannot be converted to a string
func (section *Section) GetString(name string) (string, error) {
	value, err := section.Get(name)
	if err != nil {
		return "", err
	}

	switch value.(type) {
	case *Primative:
		return value.(*Primative).AsString()
	}

	return "", errors.New("could not convert non-primative value to string")
}

// GetParent will get the parent section associated with this Section or nil
// if it does not have one
func (section *Section) GetParent() *Section {
	return section.parent
}

// HasParent will return true if this Section has a parent
func (section *Section) HasParent() bool {
	return section.parent != nil
}

// Keys will return back a list of all setting names in this Section
func (section *Section) Keys() []string {
	var keys []string
	for key := range section.values {
		keys = append(keys, key)
	}

	sort.Strings(keys)
	return keys
}

// Set will set a value (Primative or Section) to the provided name
func (section *Section) Set(name string, value Value) {
	section.values[name] = value
}

// SetBoolean will set the value for name as a bool
func (section *Section) SetBoolean(name string, value bool) {
	current, err := section.Get(name)

	// Exists just update the value/type
	if err == nil {
		current.UpdateValue(value)
	} else {
		section.values[name] = NewBoolean(value)
	}
}

// SetFloat will set the value for name as a float64
func (section *Section) SetFloat(name string, value float64) {
	current, err := section.Get(name)

	// Exists just update the value/type
	if err == nil {
		current.UpdateValue(value)
	} else {
		section.values[name] = NewFloat(value)
	}
}

// SetInteger will set the value for name as a int64
func (section *Section) SetInteger(name string, value int64) {
	current, err := section.Get(name)

	// Exists just update the value/type
	if err == nil {
		current.UpdateValue(value)
	} else {
		section.values[name] = NewInteger(value)
	}
}

// SetNull will set the value for name as nil
func (section *Section) SetNull(name string) {
	current, err := section.Get(name)

	// Already is a Null, nothing to do
	if err == nil && current.GetType() == NULL {
		return
	}
	section.Set(name, NewNull())
}

// SetString will set the value for name as a string
func (section *Section) SetString(name string, value string) {
	current, err := section.Get(name)

	// Exists just update the value/type
	if err == nil {
		current.UpdateValue(value)
	} else {
		section.Set(name, NewString(value))
	}
}

// Resolve will recursively try to fetch the provided value and will respond
// with an error if the name does not exist or tries to be resolved through
// a non-section value
func (section *Section) Resolve(name string) (Value, error) {
	// Used only in error state return value
	var value Value

	parts := strings.Split(name, ".")
	if len(parts) == 0 {
		return value, errors.New("no name provided")
	}

	var current Value
	current = section
	for _, part := range parts {
		if current.GetType() != SECTION {
			return value, errors.New("trying to resolve value from non-section")
		}

		nextCurrent, err := current.(*Section).Get(part)
		if err != nil {
			return value, errors.New("could not find value in section")
		}
		current = nextCurrent
	}
	return current, nil
}

// Merge merges the given section to current section. Settings from source
// section overwites the values in the current section
func (section *Section) Merge(source *Section) error {
	for _, key := range source.Keys() {
		sourceValue, _ := source.Get(key)
		targetValue, err := section.Get(key)

		// not found, so add it
		if err != nil {
			section.Set(key, sourceValue)
			continue
		}

		// found existing one and it's type SECTION, merge it
		if targetValue.GetType() == SECTION {
			// Source value have to be SECTION type here
			if sourceValue.GetType() != SECTION {
				return fmt.Errorf("source (%v) and target (%v) type doesn't match: %v",
					sourceValue.GetType(),
					targetValue.GetType(),
					key)
			}

			if err = targetValue.(*Section).Merge(sourceValue.(*Section)); err != nil {
				return err
			}

			continue
		}

		// found existing one, update it
		if err = targetValue.UpdateValue(sourceValue.GetValue()); err != nil {
			return fmt.Errorf("%v: %v", err, key)
		}
	}
	return nil
}

// ToJSON will convert this Section and all it's underlying values and Sections
// into JSON as a []byte
func (section *Section) ToJSON() ([]byte, error) {
	data := section.ToMap()
	return json.Marshal(data)
}

// ToMap will convert this Section and all it's underlying values and Sections into
// a map[string]interface{}
func (section *Section) ToMap() map[string]interface{} {
	output := make(map[string]interface{})

	for key, value := range section.values {
		if value.GetType() == SECTION {
			output[key] = value.(*Section).ToMap()
		} else {
			output[key] = value.GetValue()
		}
	}
	return output
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type Config struct {
	Location         string
	ChecksUrl        string
	MeasurementsUrl  string
	MeasurementsUser string
	MeasurementsPass string
}

type Check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type Measurement struct {
	Check             Check   `json:"check"`
	Id                string  `json:"id"`
	Location          string  `json:"location"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func GetEnvWithDefault(env string, def string) string {
	tmp := os.Getenv(env)

	if tmp == "" {
		return def
	}

	return tmp
}

func (c *Check) Measure(config Config) Measurement {
	var m Measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.Check = *c
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func MeasureLoop(config Config, checks chan Check, measurements chan Measurement) {
	for {
		c := <-checks
		m := c.Measure(config)

		measurements <- m
	}
}

func Record(config Config, payload []Measurement) {
	s, err := json.Marshal(&payload)
	if err != nil {
		panic(err)
	}

	body := bytes.NewBuffer(s)
	req, err := http.NewRequest("POST", config.MeasurementsUrl, body)
	if err != nil {
		panic(err)
	}

	req.Header.Add("Content-Type", "application/json")

	if config.MeasurementsUser != "" {
		req.SetBasicAuth(config.MeasurementsUser, config.MeasurementsPass)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		panic(err)
	}

	fmt.Printf("fn=Record http_code=%d\n", resp.StatusCode)
	resp.Body.Close()
}

func RecordLoop(config Config, measurements chan Measurement) {
	tickChan := time.NewTicker(time.Millisecond * 1000).C
	payload := make([]Measurement, 0, 100)

	for {
		select {
		case m := <-measurements:
			payload = append(payload, m)
		case <-tickChan:
			l := len(payload)
			fmt.Printf("fn=RecordLoop payload_size=%d\n", l)

			if l > 0 {
				Record(config, payload)
				payload = make([]Measurement, 0, 100)
			}
		}
	}
}

func GetChecks(config Config) []Check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []Check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func ScheduleLoop(check Check, checks chan Check) {
	for {
		checks <- check
		time.Sleep(1000 * time.Millisecond)
	}
}

func main() {
	var config Config
	config.Location = GetEnvWithDefault("LOCATION", "undefined")
	config.ChecksUrl = GetEnvWithDefault("CHECKS_URL", "https://s3.amazonaws.com/canary-public-data/data.json")
	config.MeasurementsUrl = GetEnvWithDefault("MEASUREMENTS_URL", "http://localhost:5000/measurements")

	u, err := url.Parse(config.MeasurementsUrl)
	if err != nil {
		panic(err)
	}

	if u.User != nil {
		config.MeasurementsUser = u.User.Username()
		config.MeasurementsPass, _ = u.User.Password()
	}

	measurerCount, err := strconv.Atoi(GetEnvWithDefault("MEASURER_COUNT", "1"))
	if err != nil {
		panic(err)
	}

	recorderCount, err := strconv.Atoi(GetEnvWithDefault("RECORDER_COUNT", "1"))
	if err != nil {
		panic(err)
	}

	check_list := GetChecks(config)

	checks := make(chan Check)
	measurements := make(chan Measurement)

	for i := 0; i < measurerCount; i++ {
		go MeasureLoop(config, checks, measurements)
	}
	for i := 0; i < recorderCount; i++ {
		go RecordLoop(config, measurements)
	}

	for _, c := range check_list {
		go ScheduleLoop(c, checks)
	}
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	<-sigs
}
<commit_msg>use flag to get config, use select {} instead of sigs to block at the end<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type Config struct {
	Location         string
	ChecksUrl        string
	MeasurementsUrl  string
	MeasurementsUser string
	MeasurementsPass string
	MeasurerCount    int
	RecorderCount    int
}

type Check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type Measurement struct {
	Check             Check   `json:"check"`
	Id                string  `json:"id"`
	Location          string  `json:"location"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func (c *Check) Measure(config Config) Measurement {
	var m Measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.Check = *c
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func MeasureLoop(config Config, checks chan Check, measurements chan Measurement) {
	for {
		c := <-checks
		m := c.Measure(config)

		measurements <- m
	}
}

func Record(config Config, payload []Measurement) {
	s, err := json.Marshal(&payload)
	if err != nil {
		panic(err)
	}

	body := bytes.NewBuffer(s)
	req, err := http.NewRequest("POST", config.MeasurementsUrl, body)
	if err != nil {
		panic(err)
	}

	req.Header.Add("Content-Type", "application/json")

	if config.MeasurementsUser != "" {
		req.SetBasicAuth(config.MeasurementsUser, config.MeasurementsPass)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		panic(err)
	}

	fmt.Printf("fn=Record http_code=%d\n", resp.StatusCode)
	resp.Body.Close()
}

func RecordLoop(config Config, measurements chan Measurement) {
	tickChan := time.NewTicker(time.Millisecond * 1000).C
	payload := make([]Measurement, 0, 100)

	for {
		select {
		case m := <-measurements:
			payload = append(payload, m)
		case <-tickChan:
			l := len(payload)
			fmt.Printf("fn=RecordLoop payload_size=%d\n", l)

			if l > 0 {
				Record(config, payload)
				payload = make([]Measurement, 0, 100)
			}
		}
	}
}

func GetChecks(config Config) []Check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []Check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func ScheduleLoop(check Check, checks chan Check) {
	for {
		checks <- check
		time.Sleep(1000 * time.Millisecond)
	}
}

func main() {
	config := Config{}
	flag.StringVar(&config.Location, "location", "undefined", "location of this sensor")
	flag.StringVar(&config.ChecksUrl, "checks_url", "https://s3.amazonaws.com/canary-public-data/checks.json", "URL for check data")
	flag.StringVar(&config.MeasurementsUrl, "measurements_url", "http://localhost:5000/measurements", "URL to POST measurements to")
	flag.IntVar(&config.MeasurerCount, "measurer_count", 1, "number of measurers to run")
	flag.IntVar(&config.RecorderCount, "recorder_count", 1, "number of recorders to run")
	flag.Parse()

	u, err := url.Parse(config.MeasurementsUrl)
	if err != nil {
		panic(err)
	}

	if u.User != nil {
		config.MeasurementsUser = u.User.Username()
		config.MeasurementsPass, _ = u.User.Password()
	}

	check_list := GetChecks(config)

	checks := make(chan Check)
	measurements := make(chan Measurement)

	for i := 0; i < config.MeasurerCount; i++ {
		go MeasureLoop(config, checks, measurements)
	}

	for i := 0; i < config.RecorderCount; i++ {
		go RecordLoop(config, measurements)
	}

	for _, c := range check_list {
		go ScheduleLoop(c, checks)
	}

	select {}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/stripe/sequins/backend"
	"github.com/stripe/sequins/index"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

type sequinsOptions struct {
	LocalPath           string
	CheckForSuccessFile bool
}

type sequins struct {
	options      sequinsOptions
	backend      backend.Backend
	indexMonitor index.IndexReference
	http         *http.Server
	started      time.Time
	updated      time.Time
	reloadLock   sync.Mutex
}

type status struct {
	Path    string `json:"path"`
	Started int64  `json:"started"`
	Updated int64  `json:"updated"`
	Count   int    `json:"count"`
}

func newSequins(backend backend.Backend, options sequinsOptions) *sequins {
	return &sequins{
		options:    options,
		backend:    backend,
		reloadLock: sync.Mutex{},
	}
}

func (s *sequins) init() error {
	err := s.refresh()
	if err != nil {
		return err
	}

	now := time.Now()
	s.started = now
	s.updated = now

	return nil
}

func (s *sequins) start(address string) error {
	// TODO: we may need a more graceful way of shutting down, since this will
	// cause requests that start processing after this runs to 500
	// However, this may not be a problem, since you have to shift traffic to
	// another instance before shutting down anyway, otherwise you'd have downtime

	defer s.indexMonitor.Replace(nil).Close()

	log.Printf("Listening on %s", address)
	return http.ListenAndServe(address, s)
}

func (s *sequins) reloadLatest() error {
	err := s.refresh()
	if err != nil {
		return err
	}

	s.updated = time.Now()

	return nil
}

func (s *sequins) refresh() error {
	s.reloadLock.Lock()
	defer s.reloadLock.Unlock()

	version, err := s.backend.LatestVersion(s.options.CheckForSuccessFile)
	if err != nil {
		return err
	}

	// We can use unsafe ref, since closing the index would not affect the version string
	var currentVersion string
	currentIndex := s.indexMonitor.UnsafeGet()
	if currentIndex != nil {
		currentVersion = currentIndex.Version
	}

	if version != currentVersion {
		path := filepath.Join(s.options.LocalPath, version)

		err := os.Mkdir(path, 0700|os.ModeDir)
		if err != nil && !os.IsExist(err) {
			return err
		}

		if os.IsExist(err) {
			log.Printf("Version %s is already downloaded", version)
		} else {
			log.Printf("Downloading version %s from %s", version, s.backend.DisplayPath(version))
			err = s.backend.Download(version, path)
			if err != nil {
				return err
			}
		}

		log.Printf("Preparing version %s at %s", version, path)
		index := index.New(path, version)
		err = index.Load()
		if err != nil {
			return fmt.Errorf("Error while indexing: %s", err)
		}

		log.Printf("Switching to version %s!", version)

		oldIndex := s.indexMonitor.Replace(index)
		if oldIndex != nil {
			oldIndex.Close()
		}
	} else {
		log.Printf("%s is already the newest version, so not reloading.", version)
	}

	return nil
}

func (s *sequins) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" {
		index := s.indexMonitor.Get()
		count, err := index.Count()
		currentVersion := index.Version
		s.indexMonitor.Release(index)

		if err != nil {
			log.Fatal(err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}

		status := status{
			Path:    s.backend.DisplayPath(currentVersion),
			Started: s.started.Unix(),
			Updated: s.updated.Unix(),
			Count:   count,
		}

		jsonBytes, err := json.Marshal(status)
		if err != nil {
			log.Fatal(err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}

		w.Write(jsonBytes)
		return
	}

	key := strings.TrimPrefix(r.URL.Path, "/")

	currentIndex := s.indexMonitor.Get()
	res, err := currentIndex.Get(key)
	s.indexMonitor.Release(currentIndex)

	if err == index.ErrNotFound {
		w.WriteHeader(http.StatusNotFound)
	} else if err != nil {
		log.Fatal(fmt.Errorf("Error fetching value for %s: %s", key, err))
		w.WriteHeader(http.StatusInternalServerError)
	} else {
		// Explicitly unset Content-Type, so ServeContent doesn't try to do any
		// sniffing.
		w.Header()["Content-Type"] = nil

		http.ServeContent(w, r, key, s.updated, bytes.NewReader(res))
	}
}
<commit_msg>Fix minor naming inconsistency<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/stripe/sequins/backend"
	"github.com/stripe/sequins/index"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

type sequinsOptions struct {
	LocalPath           string
	CheckForSuccessFile bool
}

type sequins struct {
	options      sequinsOptions
	backend      backend.Backend
	indexReference index.IndexReference
	http         *http.Server
	started      time.Time
	updated      time.Time
	reloadLock   sync.Mutex
}

type status struct {
	Path    string `json:"path"`
	Started int64  `json:"started"`
	Updated int64  `json:"updated"`
	Count   int    `json:"count"`
}

func newSequins(backend backend.Backend, options sequinsOptions) *sequins {
	return &sequins{
		options:    options,
		backend:    backend,
		reloadLock: sync.Mutex{},
	}
}

func (s *sequins) init() error {
	err := s.refresh()
	if err != nil {
		return err
	}

	now := time.Now()
	s.started = now
	s.updated = now

	return nil
}

func (s *sequins) start(address string) error {
	// TODO: we may need a more graceful way of shutting down, since this will
	// cause requests that start processing after this runs to 500
	// However, this may not be a problem, since you have to shift traffic to
	// another instance before shutting down anyway, otherwise you'd have downtime

	defer s.indexReference.Replace(nil).Close()

	log.Printf("Listening on %s", address)
	return http.ListenAndServe(address, s)
}

func (s *sequins) reloadLatest() error {
	err := s.refresh()
	if err != nil {
		return err
	}

	s.updated = time.Now()

	return nil
}

func (s *sequins) refresh() error {
	s.reloadLock.Lock()
	defer s.reloadLock.Unlock()

	version, err := s.backend.LatestVersion(s.options.CheckForSuccessFile)
	if err != nil {
		return err
	}

	// We can use unsafe ref, since closing the index would not affect the version string
	var currentVersion string
	currentIndex := s.indexReference.UnsafeGet()
	if currentIndex != nil {
		currentVersion = currentIndex.Version
	}

	if version != currentVersion {
		path := filepath.Join(s.options.LocalPath, version)

		err := os.Mkdir(path, 0700|os.ModeDir)
		if err != nil && !os.IsExist(err) {
			return err
		}

		if os.IsExist(err) {
			log.Printf("Version %s is already downloaded", version)
		} else {
			log.Printf("Downloading version %s from %s", version, s.backend.DisplayPath(version))
			err = s.backend.Download(version, path)
			if err != nil {
				return err
			}
		}

		log.Printf("Preparing version %s at %s", version, path)
		index := index.New(path, version)
		err = index.Load()
		if err != nil {
			return fmt.Errorf("Error while indexing: %s", err)
		}

		log.Printf("Switching to version %s!", version)

		oldIndex := s.indexReference.Replace(index)
		if oldIndex != nil {
			oldIndex.Close()
		}
	} else {
		log.Printf("%s is already the newest version, so not reloading.", version)
	}

	return nil
}

func (s *sequins) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/" {
		index := s.indexReference.Get()
		count, err := index.Count()
		currentVersion := index.Version
		s.indexReference.Release(index)

		if err != nil {
			log.Fatal(err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}

		status := status{
			Path:    s.backend.DisplayPath(currentVersion),
			Started: s.started.Unix(),
			Updated: s.updated.Unix(),
			Count:   count,
		}

		jsonBytes, err := json.Marshal(status)
		if err != nil {
			log.Fatal(err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}

		w.Write(jsonBytes)
		return
	}

	key := strings.TrimPrefix(r.URL.Path, "/")

	currentIndex := s.indexReference.Get()
	res, err := currentIndex.Get(key)
	s.indexReference.Release(currentIndex)

	if err == index.ErrNotFound {
		w.WriteHeader(http.StatusNotFound)
	} else if err != nil {
		log.Fatal(fmt.Errorf("Error fetching value for %s: %s", key, err))
		w.WriteHeader(http.StatusInternalServerError)
	} else {
		// Explicitly unset Content-Type, so ServeContent doesn't try to do any
		// sniffing.
		w.Header()["Content-Type"] = nil

		http.ServeContent(w, r, key, s.updated, bytes.NewReader(res))
	}
}
<|endoftext|>


<commit_before>package micro

import (
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"

	"github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/config/cmd"
	"github.com/micro/go-micro/v2/debug/profile"
	"github.com/micro/go-micro/v2/debug/profile/http"
	"github.com/micro/go-micro/v2/debug/profile/pprof"
	"github.com/micro/go-micro/v2/debug/service/handler"
	"github.com/micro/go-micro/v2/debug/stats"
	"github.com/micro/go-micro/v2/debug/trace"
	"github.com/micro/go-micro/v2/plugin"
	"github.com/micro/go-micro/v2/server"
	"github.com/micro/go-micro/v2/util/log"
	"github.com/micro/go-micro/v2/util/wrapper"
)

type service struct {
	opts Options

	once sync.Once
}

func newService(opts ...Option) Service {
	options := newOptions(opts...)

	// service name
	serviceName := options.Server.Options().Name

	// wrap client to inject From-Service header on any calls
	options.Client = wrapper.FromService(serviceName, options.Client)
	options.Client = wrapper.TraceCall(serviceName, trace.DefaultTracer, options.Client)

	// wrap the server to provide handler stats
	options.Server.Init(
		server.WrapHandler(wrapper.HandlerStats(stats.DefaultStats)),
		server.WrapHandler(wrapper.TraceHandler(trace.DefaultTracer)),
	)

	return &service{
		opts: options,
	}
}

func (s *service) Name() string {
	return s.opts.Server.Options().Name
}

// Init initialises options. Additionally it calls cmd.Init
// which parses command line flags. cmd.Init is only called
// on first Init.
func (s *service) Init(opts ...Option) {
	// process options
	for _, o := range opts {
		o(&s.opts)
	}

	s.once.Do(func() {
		// setup the plugins
		for _, p := range strings.Split(os.Getenv("MICRO_PLUGIN"), ",") {
			if len(p) == 0 {
				continue
			}

			// load the plugin
			c, err := plugin.Load(p)
			if err != nil {
				log.Fatal(err)
			}

			// initialise the plugin
			if err := plugin.Init(c); err != nil {
				log.Fatal(err)
			}
		}

		// set cmd name
		if len(s.opts.Cmd.App().Name) == 0 {
			s.opts.Cmd.App().Name = s.Server().Options().Name
		}

		// Initialise the command flags, overriding new service
		_ = s.opts.Cmd.Init(
			cmd.Broker(&s.opts.Broker),
			cmd.Registry(&s.opts.Registry),
			cmd.Transport(&s.opts.Transport),
			cmd.Client(&s.opts.Client),
			cmd.Server(&s.opts.Server),
		)
	})
}

func (s *service) Options() Options {
	return s.opts
}

func (s *service) Client() client.Client {
	return s.opts.Client
}

func (s *service) Server() server.Server {
	return s.opts.Server
}

func (s *service) String() string {
	return "micro"
}

func (s *service) Start() error {
	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	if err := s.opts.Server.Start(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	return nil
}

func (s *service) Stop() error {
	var gerr error

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	if err := s.opts.Server.Stop(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	return gerr
}

func (s *service) Run() error {
	// register the debug handler
	s.opts.Server.Handle(
		s.opts.Server.NewHandler(
			handler.NewHandler(),
			server.InternalHandler(true),
		),
	)

	// start the profiler
	// TODO: set as an option to the service, don't just use pprof
	if prof := os.Getenv("MICRO_DEBUG_PROFILE"); len(prof) > 0 {
		var profiler profile.Profile

		// to view mutex contention
		runtime.SetMutexProfileFraction(5)
		// to view blocking profile
		runtime.SetBlockProfileRate(1)

		switch prof {
		case "http":
			profiler = http.NewProfile()
		default:
			service := s.opts.Server.Options().Name
			version := s.opts.Server.Options().Version
			id := s.opts.Server.Options().Id
			profiler = pprof.NewProfile(
				profile.Name(service + "." + version + "." + id),
			)
		}

		if err := profiler.Start(); err != nil {
			return err
		}
		defer profiler.Stop()
	}

	log.Logf("Starting [service] %s", s.Name())

	if err := s.Start(); err != nil {
		return err
	}

	ch := make(chan os.Signal, 1)
	if s.opts.Signal {
		signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT)
	}

	select {
	// wait on kill signal
	case <-ch:
	// wait on context cancel
	case <-s.opts.Context.Done():
	}

	return s.Stop()
}
<commit_msg>fatal on command error<commit_after>package micro

import (
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"

	"github.com/micro/go-micro/v2/client"
	"github.com/micro/go-micro/v2/config/cmd"
	"github.com/micro/go-micro/v2/debug/profile"
	"github.com/micro/go-micro/v2/debug/profile/http"
	"github.com/micro/go-micro/v2/debug/profile/pprof"
	"github.com/micro/go-micro/v2/debug/service/handler"
	"github.com/micro/go-micro/v2/debug/stats"
	"github.com/micro/go-micro/v2/debug/trace"
	"github.com/micro/go-micro/v2/plugin"
	"github.com/micro/go-micro/v2/server"
	"github.com/micro/go-micro/v2/util/log"
	"github.com/micro/go-micro/v2/util/wrapper"
)

type service struct {
	opts Options

	once sync.Once
}

func newService(opts ...Option) Service {
	options := newOptions(opts...)

	// service name
	serviceName := options.Server.Options().Name

	// wrap client to inject From-Service header on any calls
	options.Client = wrapper.FromService(serviceName, options.Client)
	options.Client = wrapper.TraceCall(serviceName, trace.DefaultTracer, options.Client)

	// wrap the server to provide handler stats
	options.Server.Init(
		server.WrapHandler(wrapper.HandlerStats(stats.DefaultStats)),
		server.WrapHandler(wrapper.TraceHandler(trace.DefaultTracer)),
	)

	return &service{
		opts: options,
	}
}

func (s *service) Name() string {
	return s.opts.Server.Options().Name
}

// Init initialises options. Additionally it calls cmd.Init
// which parses command line flags. cmd.Init is only called
// on first Init.
func (s *service) Init(opts ...Option) {
	// process options
	for _, o := range opts {
		o(&s.opts)
	}

	s.once.Do(func() {
		// setup the plugins
		for _, p := range strings.Split(os.Getenv("MICRO_PLUGIN"), ",") {
			if len(p) == 0 {
				continue
			}

			// load the plugin
			c, err := plugin.Load(p)
			if err != nil {
				log.Fatal(err)
			}

			// initialise the plugin
			if err := plugin.Init(c); err != nil {
				log.Fatal(err)
			}
		}

		// set cmd name
		if len(s.opts.Cmd.App().Name) == 0 {
			s.opts.Cmd.App().Name = s.Server().Options().Name
		}

		// Initialise the command flags, overriding new service
		if err := s.opts.Cmd.Init(
			cmd.Broker(&s.opts.Broker),
			cmd.Registry(&s.opts.Registry),
			cmd.Transport(&s.opts.Transport),
			cmd.Client(&s.opts.Client),
			cmd.Server(&s.opts.Server),
		); err != nil {
			log.Fatal(err)
		}
	})
}

func (s *service) Options() Options {
	return s.opts
}

func (s *service) Client() client.Client {
	return s.opts.Client
}

func (s *service) Server() server.Server {
	return s.opts.Server
}

func (s *service) String() string {
	return "micro"
}

func (s *service) Start() error {
	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	if err := s.opts.Server.Start(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	return nil
}

func (s *service) Stop() error {
	var gerr error

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	if err := s.opts.Server.Stop(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	return gerr
}

func (s *service) Run() error {
	// register the debug handler
	s.opts.Server.Handle(
		s.opts.Server.NewHandler(
			handler.NewHandler(),
			server.InternalHandler(true),
		),
	)

	// start the profiler
	// TODO: set as an option to the service, don't just use pprof
	if prof := os.Getenv("MICRO_DEBUG_PROFILE"); len(prof) > 0 {
		var profiler profile.Profile

		// to view mutex contention
		runtime.SetMutexProfileFraction(5)
		// to view blocking profile
		runtime.SetBlockProfileRate(1)

		switch prof {
		case "http":
			profiler = http.NewProfile()
		default:
			service := s.opts.Server.Options().Name
			version := s.opts.Server.Options().Version
			id := s.opts.Server.Options().Id
			profiler = pprof.NewProfile(
				profile.Name(service + "." + version + "." + id),
			)
		}

		if err := profiler.Start(); err != nil {
			return err
		}
		defer profiler.Stop()
	}

	log.Logf("Starting [service] %s", s.Name())

	if err := s.Start(); err != nil {
		return err
	}

	ch := make(chan os.Signal, 1)
	if s.opts.Signal {
		signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT)
	}

	select {
	// wait on kill signal
	case <-ch:
	// wait on context cancel
	case <-s.opts.Context.Done():
	}

	return s.Stop()
}
<|endoftext|>


<commit_before>package siesta

import (
	"net/http"
	"path"
	"regexp"
	"strings"
)

// Registered services keyed by base URI.
var services = map[string]*Service{}

// A Service is a container for routes with a common base URI.
// It also has two middleware chains, named "pre" and "post".
//
// The "pre" chain is run before the main handler. The first
// handler in the "pre" chain is guaranteed to run, but execution
// may quit anywhere else in the chain.
//
// If the "pre" chain executes completely, the main handler is executed.
// It is skipped otherwise.
//
// The "post" chain runs after the main handler, whether it is skipped
// or not. The first handler in the "post" chain is guaranteed to run, but
// execution may quit anywhere else in the chain.
type Service struct {
	baseURI string

	pre  []contextHandler
	post []contextHandler

	handlers map[*regexp.Regexp]contextHandler

	routes map[string]*node
}

// NewService returns a new Service with the given base URI
// or panics if the base URI has already been registered.
func NewService(baseURI string) *Service {
	if services[baseURI] != nil {
		panic("service already registered")
	}

	return &Service{
		baseURI:  path.Join("/", baseURI, "/"),
		handlers: make(map[*regexp.Regexp]contextHandler),
		routes:   map[string]*node{},
	}
}

func addToChain(f interface{}, chain []contextHandler) []contextHandler {
	m := toContextHandler(f)
	return append(chain, m)
}

// AddPre adds f to the end of the "pre" chain.
// It panics if f cannot be converted to a contextHandler (see Service.Route).
func (s *Service) AddPre(f interface{}) {
	s.pre = addToChain(f, s.pre)
}

// AddPost adds f to the end of the "post" chain.
// It panics if f cannot be converted to a contextHandler (see Service.Route).
func (s *Service) AddPost(f interface{}) {
	s.post = addToChain(f, s.post)
}

// Service satisfies the http.Handler interface.
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.ServeHTTPInContext(NewSiestaContext(), w, r)
}

// ServiceHTTPInContext serves an HTTP request within the Context c.
// A Service will run through both of its internal chains, quitting
// when requested.
func (s *Service) ServeHTTPInContext(c Context, w http.ResponseWriter, r *http.Request) {
	quit := false
	for _, m := range s.pre {
		m(c, w, r, func() {
			quit = true
		})

		if quit {
			// Break out of the "pre" loop, but
			// continue on.
			break
		}
	}

	if !quit {
		// The main handler is only run if we have not
		// been signaled to quit.

		if r.URL.Path != "/" {
			r.URL.Path = strings.TrimRight(r.URL.Path, "/")
		}

		var (
			handler contextHandler
			params  routeParams
		)

		// Lookup the tree for this method
		routeNode, ok := s.routes[r.Method]

		if ok {
			handler, params, _ = routeNode.getValue(r.URL.Path)
		}

		if handler == nil {
			http.NotFoundHandler().ServeHTTP(w, r)
		} else {
			r.ParseForm()
			for _, p := range params {
				r.Form.Set(p.Key, p.Value)
			}

			handler(c, w, r, func() {
				quit = true
			})
		}
	}

	for _, m := range s.post {
		m(c, w, r, func() {
			quit = true
		})

		if quit {
			return
		}
	}
}

// Route adds a new route to the Service.
// f must be a function with one of the following signatures:
//
//     func(http.ResponseWriter, *http.Request)
//     func(http.ResponseWriter, *http.Request, func())
//     func(Context, http.ResponseWriter, *http.Request)
//     func(Context, http.ResponseWriter, *http.Request, func())
//
// Note that Context is an interface type defined in this package.
// The last argument is a function which is called to signal the
// quitting of the current execution sequence.
func (s *Service) Route(verb, uriPath, usage string, f interface{}) {
	handler := toContextHandler(f)

	if n := s.routes[verb]; n == nil {
		s.routes[verb] = &node{}
	}

	s.routes[verb].addRoute(path.Join(s.baseURI, strings.TrimRight(uriPath, "/")), handler)
}

// Register registers s by adding it as a handler to the
// DefaultServeMux in the net/http package.
func (s *Service) Register() {
	http.Handle(s.baseURI, s)
}
<commit_msg>support custom "not found" handlers; #35<commit_after>package siesta

import (
	"net/http"
	"path"
	"regexp"
	"strings"
)

// Registered services keyed by base URI.
var services = map[string]*Service{}

// A Service is a container for routes with a common base URI.
// It also has two middleware chains, named "pre" and "post".
//
// The "pre" chain is run before the main handler. The first
// handler in the "pre" chain is guaranteed to run, but execution
// may quit anywhere else in the chain.
//
// If the "pre" chain executes completely, the main handler is executed.
// It is skipped otherwise.
//
// The "post" chain runs after the main handler, whether it is skipped
// or not. The first handler in the "post" chain is guaranteed to run, but
// execution may quit anywhere else in the chain.
type Service struct {
	baseURI string

	pre  []contextHandler
	post []contextHandler

	handlers map[*regexp.Regexp]contextHandler

	routes map[string]*node

	notFound contextHandler
}

// NewService returns a new Service with the given base URI
// or panics if the base URI has already been registered.
func NewService(baseURI string) *Service {
	if services[baseURI] != nil {
		panic("service already registered")
	}

	return &Service{
		baseURI:  path.Join("/", baseURI, "/"),
		handlers: make(map[*regexp.Regexp]contextHandler),
		routes:   map[string]*node{},
	}
}

func addToChain(f interface{}, chain []contextHandler) []contextHandler {
	m := toContextHandler(f)
	return append(chain, m)
}

// AddPre adds f to the end of the "pre" chain.
// It panics if f cannot be converted to a contextHandler (see Service.Route).
func (s *Service) AddPre(f interface{}) {
	s.pre = addToChain(f, s.pre)
}

// AddPost adds f to the end of the "post" chain.
// It panics if f cannot be converted to a contextHandler (see Service.Route).
func (s *Service) AddPost(f interface{}) {
	s.post = addToChain(f, s.post)
}

// Service satisfies the http.Handler interface.
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.ServeHTTPInContext(NewSiestaContext(), w, r)
}

// ServiceHTTPInContext serves an HTTP request within the Context c.
// A Service will run through both of its internal chains, quitting
// when requested.
func (s *Service) ServeHTTPInContext(c Context, w http.ResponseWriter, r *http.Request) {
	quit := false
	for _, m := range s.pre {
		m(c, w, r, func() {
			quit = true
		})

		if quit {
			// Break out of the "pre" loop, but
			// continue on.
			break
		}
	}

	if !quit {
		// The main handler is only run if we have not
		// been signaled to quit.

		if r.URL.Path != "/" {
			r.URL.Path = strings.TrimRight(r.URL.Path, "/")
		}

		var (
			handler contextHandler
			params  routeParams
		)

		// Lookup the tree for this method
		routeNode, ok := s.routes[r.Method]

		if ok {
			handler, params, _ = routeNode.getValue(r.URL.Path)
		}

		if handler == nil {
			if s.notFound != nil {
				s.notFound(c, w, r, func() {})
			}
		} else {
			r.ParseForm()
			for _, p := range params {
				r.Form.Set(p.Key, p.Value)
			}

			handler(c, w, r, func() {
				quit = true
			})
		}
	}

	for _, m := range s.post {
		m(c, w, r, func() {
			quit = true
		})

		if quit {
			return
		}
	}
}

// Route adds a new route to the Service.
// f must be a function with one of the following signatures:
//
//     func(http.ResponseWriter, *http.Request)
//     func(http.ResponseWriter, *http.Request, func())
//     func(Context, http.ResponseWriter, *http.Request)
//     func(Context, http.ResponseWriter, *http.Request, func())
//
// Note that Context is an interface type defined in this package.
// The last argument is a function which is called to signal the
// quitting of the current execution sequence.
func (s *Service) Route(verb, uriPath, usage string, f interface{}) {
	handler := toContextHandler(f)

	if n := s.routes[verb]; n == nil {
		s.routes[verb] = &node{}
	}

	s.routes[verb].addRoute(path.Join(s.baseURI, strings.TrimRight(uriPath, "/")), handler)
}

// SetNotFound sets the handler for all paths that do not
// match any existing routes. It accepts the same function
// signatures that Route does.
func (s *Service) SetNotFound(f interface{}) {
	handler := toContextHandler(f)
	s.notFound = handler
}

// Register registers s by adding it as a handler to the
// DefaultServeMux in the net/http package.
func (s *Service) Register() {
	http.Handle(s.baseURI, s)
}
<|endoftext|>


<commit_before>package micro

import (
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	log "github.com/micro/go-log"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/cmd"
	"github.com/micro/go-micro/metadata"
	"github.com/micro/go-micro/server"
)

type service struct {
	opts Options

	once sync.Once
}

func newService(opts ...Option) Service {
	options := newOptions(opts...)

	options.Client = &clientWrapper{
		options.Client,
		metadata.Metadata{
			HeaderPrefix + "From-Service": options.Server.Options().Name,
		},
	}

	return &service{
		opts: options,
	}
}

func (s *service) run(exit chan bool) {
	if s.opts.RegisterInterval <= time.Duration(0) {
		return
	}

	t := time.NewTicker(s.opts.RegisterInterval)

	for {
		select {
		case <-t.C:
			err := s.opts.Server.Register()
			if err != nil {
				log.Log("service run Server.Register err : ", err)
			}
		case <-exit:
			t.Stop()
			return
		}
	}
}

// Init initialises options. Additionally it calls cmd.Init
// which parses command line flags. cmd.Init is only called
// on first Init.
func (s *service) Init(opts ...Option) {
	// process options
	for _, o := range opts {
		o(&s.opts)
	}

	s.once.Do(func() {
		// Initialise the command flags, overriding new service
		s.opts.Cmd.Init(
			cmd.Broker(&s.opts.Broker),
			cmd.Registry(&s.opts.Registry),
			cmd.Transport(&s.opts.Transport),
			cmd.Client(&s.opts.Client),
			cmd.Server(&s.opts.Server),
		)
	})
}

func (s *service) Options() Options {
	return s.opts
}

func (s *service) Client() client.Client {
	return s.opts.Client
}

func (s *service) Server() server.Server {
	return s.opts.Server
}

func (s *service) String() string {
	return "go-micro"
}

func (s *service) Start() error {
	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	if err := s.opts.Server.Start(); err != nil {
		return err
	}

	if err := s.opts.Server.Register(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	return nil
}

func (s *service) Stop() error {
	var gerr error

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	if err := s.opts.Server.Deregister(); err != nil {
		return err
	}

	if err := s.opts.Server.Stop(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	return gerr
}

func (s *service) Run() error {
	if err := s.Start(); err != nil {
		return err
	}

	// start reg loop
	ex := make(chan bool)
	go s.run(ex)

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)

	select {
	// wait on kill signal
	case <-ch:
	// wait on context cancel
	case <-s.opts.Context.Done():
	}

	// exit reg loop
	close(ex)

	if err := s.Stop(); err != nil {
		return err
	}

	return nil
}
<commit_msg>Remove whitespace<commit_after>package micro

import (
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	log "github.com/micro/go-log"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/cmd"
	"github.com/micro/go-micro/metadata"
	"github.com/micro/go-micro/server"
)

type service struct {
	opts Options

	once sync.Once
}

func newService(opts ...Option) Service {
	options := newOptions(opts...)

	options.Client = &clientWrapper{
		options.Client,
		metadata.Metadata{
			HeaderPrefix + "From-Service": options.Server.Options().Name,
		},
	}

	return &service{
		opts: options,
	}
}

func (s *service) run(exit chan bool) {
	if s.opts.RegisterInterval <= time.Duration(0) {
		return
	}

	t := time.NewTicker(s.opts.RegisterInterval)

	for {
		select {
		case <-t.C:
			err := s.opts.Server.Register()
			if err != nil {
				log.Log("service run Server.Register error: ", err)
			}
		case <-exit:
			t.Stop()
			return
		}
	}
}

// Init initialises options. Additionally it calls cmd.Init
// which parses command line flags. cmd.Init is only called
// on first Init.
func (s *service) Init(opts ...Option) {
	// process options
	for _, o := range opts {
		o(&s.opts)
	}

	s.once.Do(func() {
		// Initialise the command flags, overriding new service
		s.opts.Cmd.Init(
			cmd.Broker(&s.opts.Broker),
			cmd.Registry(&s.opts.Registry),
			cmd.Transport(&s.opts.Transport),
			cmd.Client(&s.opts.Client),
			cmd.Server(&s.opts.Server),
		)
	})
}

func (s *service) Options() Options {
	return s.opts
}

func (s *service) Client() client.Client {
	return s.opts.Client
}

func (s *service) Server() server.Server {
	return s.opts.Server
}

func (s *service) String() string {
	return "go-micro"
}

func (s *service) Start() error {
	for _, fn := range s.opts.BeforeStart {
		if err := fn(); err != nil {
			return err
		}
	}

	if err := s.opts.Server.Start(); err != nil {
		return err
	}

	if err := s.opts.Server.Register(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStart {
		if err := fn(); err != nil {
			return err
		}
	}

	return nil
}

func (s *service) Stop() error {
	var gerr error

	for _, fn := range s.opts.BeforeStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	if err := s.opts.Server.Deregister(); err != nil {
		return err
	}

	if err := s.opts.Server.Stop(); err != nil {
		return err
	}

	for _, fn := range s.opts.AfterStop {
		if err := fn(); err != nil {
			gerr = err
		}
	}

	return gerr
}

func (s *service) Run() error {
	if err := s.Start(); err != nil {
		return err
	}

	// start reg loop
	ex := make(chan bool)
	go s.run(ex)

	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)

	select {
	// wait on kill signal
	case <-ch:
	// wait on context cancel
	case <-s.opts.Context.Done():
	}

	// exit reg loop
	close(ex)

	if err := s.Stop(); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package common

import (
	"log"
	"os"
	"os/signal"
	"syscall"
)

// Service interface
type Service interface {
	Start() error
	Stop() error
}

// Serve starts a service, and stops it if recieve INT or TERM signal.
func Serve(s Service) {
	signalCh := make(chan os.Signal, 1)
	exitCh := make(chan bool)

	go func() {
		sig := <-signalCh
		log.Printf("recieve signal: %s", sig)
		exitCh <- true
	}()

	// listening INT & TERM signal
	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		log.Printf("start service")
		err := s.Start()
		if err != nil {
			log.Printf("service ended with error: %s", err)
		} else {
			log.Printf("service ended")
		}
		exitCh <- true
	}()

	<-exitCh

	log.Printf("stopping service...")
	if err := s.Stop(); err != nil {
		log.Fatalf("servic ended with error: %s", err)
	}
	log.Printf("Bye-bye!")
}
<commit_msg>Added Name method<commit_after>package common

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
)

var logger *log.Logger

// Service interface
type Service interface {
	Name() string
	Start() error
	Stop() error
}

// Serve starts a service, and stops it if recieve INT or TERM signal.
func Serve(s Service) {

	logf := func(format string, v ...interface{}) {
		name := s.Name()
		msg := fmt.Sprintf(format, v...)
		msg = fmt.Sprintf("[Service %s] %s", name, msg)
		if logger != nil {
			logger.Println(msg)
		} else {
			log.Println(msg)
		}
	}

	signalCh := make(chan os.Signal, 1)
	exitCh := make(chan bool)

	go func() {
		sig := <-signalCh
		logf("recieve signal: %s", sig)
		exitCh <- true
	}()

	// listening INT & TERM signal
	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		logf("started")
		err := s.Start()
		if err != nil {
			logf("ended unexpectely: %s", err)
		} else {
			logf("ended")
		}
		exitCh <- true
	}()

	<-exitCh

	logf("stopping...")
	if err := s.Stop(); err != nil {
		logf("stopped with error: %s", err)
	}
	logf("Bye-bye!")
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/cheggaaa/pb"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"log"
	"math"
	"os"
	"os/user"
	"path"
	"path/filepath"
	"strings"
)

var subdivisions *int
var tolerance *int

var scratchDir string

type pictable [][][]uint64

func MkPictable(dx int, dy int) pictable {
	pic := make([][][]uint64, dx) /* type declaration */
	for i := range pic {
		pic[i] = make([][]uint64, dy) /* again the type? */
		for j := range pic[i] {
			pic[i][j] = []uint64{0, 0, 0}
		}
	}
	return pic
}

func absdiff(a uint64, b uint64) uint64 {
	return uint64(math.Abs(float64(a) - float64(b)))
}

func init() {
	subdivisions = flag.Int("subdivisions", 10, "Slices per axis")
	tolerance = flag.Int("tolerance", 100, "Color delta tolerance, higher = more tolerant")
	flag.Parse()

	if flag.NArg() < 1 {
		fmt.Println("usage: imgdedup [options] [<directories>/files]")
		flag.PrintDefaults()
		os.Exit(1)
	}

	image.RegisterFormat("png", "png", png.Decode, png.DecodeConfig)
	image.RegisterFormat("jpeg", "jpeg", jpeg.Decode, jpeg.DecodeConfig)
	image.RegisterFormat("gif", "gif", gif.Decode, gif.DecodeConfig)
}

func init() {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	scratchDir = path.Join(usr.HomeDir, ".imgdedup")

	if _, err := os.Stat(scratchDir); err != nil {
		if os.IsNotExist(err) {
			os.Mkdir(scratchDir, 0700)
		} else {
			log.Fatal(err)
		}
	}
}

func scanImg(file *os.File) (pictable, error) {
	m, _, err := image.Decode(file)
	if err != nil {
		return nil, err
	}
	bounds := m.Bounds()

	avgdata := MkPictable(*subdivisions, *subdivisions)

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			rX := int64(math.Floor((float64(x) / float64(bounds.Max.X)) * float64(*subdivisions)))
			rY := int64(math.Floor((float64(y) / float64(bounds.Max.Y)) * float64(*subdivisions)))

			r, g, b, _ := m.At(x, y).RGBA()
			avgdata[rX][rY][0] += uint64((float32(r) / 65535) * 255)
			avgdata[rX][rY][1] += uint64((float32(g) / 65535) * 255)
			avgdata[rX][rY][2] += uint64((float32(b) / 65535) * 255)
		}
	}

	divisor := uint64((bounds.Max.X / *subdivisions) * (bounds.Max.Y / *subdivisions))

	for rX := 0; rX < *subdivisions; rX++ {
		for rY := 0; rY < *subdivisions; rY++ {
			avgdata[rX][rY][0] = avgdata[rX][rY][0] / divisor
			avgdata[rX][rY][1] = avgdata[rX][rY][1] / divisor
			avgdata[rX][rY][2] = avgdata[rX][rY][2] / divisor
		}
	}

	return avgdata, nil
}

func loadCache(cachename string) (pictable, error) {

	file, err := os.Open(cachename)
	defer file.Close()
	if err != nil {
		return nil, err
	}
	r := bufio.NewReader(file)

	var avgdata pictable

	dec := json.NewDecoder(r)

	err = dec.Decode(&avgdata)
	if err != nil {
		return nil, err
	}

	return avgdata, nil
}

func storeCache(cachename string, avgdata *pictable) {
	fo, err := os.Create(cachename)
	defer fo.Close()
	if err != nil {
		panic(err)
	}

	enc := json.NewEncoder(fo)
	enc.Encode(avgdata)
}

func main() {

	imgdata := make(map[string]pictable)

	fileList := getFiles(flag.Args())

	bar := pb.StartNew(len(fileList))

	for _, imgpath := range fileList {

		bar.Increment()

		file, err := os.Open(imgpath)
		if err != nil {
			log.Fatal(err)
		}

		fExt := strings.ToLower(filepath.Ext(imgpath))
		if fExt == ".png" || fExt == ".jpg" || fExt == ".jpeg" || fExt == ".gif" {

			fi, err := file.Stat()
			if err != nil {
				log.Fatal(err)
			}

			h := md5.New()

			cacheUnit := imgpath+"|"+string(*subdivisions)+"|"+string(fi.Size())+string(fi.ModTime().Unix());

			io.WriteString(h, cacheUnit)
			cachename := path.Join(scratchDir, fmt.Sprintf("%x", h.Sum(nil))+".tmp")

			var avgdata pictable

			avgdata, err = loadCache(cachename)
			if err != nil {

				avgdata, err = scanImg(file)
				if err != nil {
					log.Print(imgpath, err)
					continue
				}

				storeCache(cachename, &avgdata)
			}

			imgdata[imgpath] = avgdata

			file.Close()

		} else {
			file.Close()
		}
	}

	bar.Finish()

	fileLength := len(fileList)

	for i := 0; i < fileLength-1; i++ {
		for j := i + 1; j < fileLength-1; j++ {

			filename1 := fileList[i]
			filename2 := fileList[j]

			avgdata1, ok1 := imgdata[filename1]
			avgdata2, ok2 := imgdata[filename2]

			if ok1 && ok2 {

				if filename1 == filename2 {
					continue
				}

				var xdiff uint64 = 0

				for rX := 0; rX < *subdivisions; rX++ {
					for rY := 0; rY < *subdivisions; rY++ {
						aa := avgdata1[rX][rY]
						bb := avgdata2[rX][rY]

						xdiff += absdiff(absdiff(absdiff(aa[0], bb[0]), absdiff(aa[1], bb[1])), absdiff(aa[2], bb[2]))
					}
				}

				if xdiff < uint64(*tolerance) {
					fmt.Println(filename1, filename2)
					fmt.Println(xdiff)
				}

			}

		}
	}

}

func getFiles(paths []string) []string {
	var fileList []string

	for _, imgpath := range paths {

		file, err := os.Open(imgpath)
		if err != nil {
			log.Fatal(err)
		}

		fi, err := file.Stat()
		if err != nil {
			log.Fatal(err)
		}

		switch mode := fi.Mode(); {
		case mode.IsDir():
			// fmt.Println("directory")
			filepath.Walk(imgpath, func(path string, f os.FileInfo, err error) error {

				submode := f.Mode()
				if submode.IsRegular() {
					fpath, _ := filepath.Abs(path)
					fileList = append(fileList, fpath)
				}

				return nil
			})
		case mode.IsRegular():
			// fmt.Println("file")
			fpath, _ := filepath.Abs(imgpath)
			fileList = append(fileList, fpath)
		}

		file.Close()

	}

	return fileList
}
<commit_msg>Divide by zero error on bad image fixed<commit_after>package main

import (
	"bufio"
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/cheggaaa/pb"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"log"
	"math"
	"os"
	"os/user"
	"path"
	"path/filepath"
	"strings"
)

var subdivisions *int
var tolerance *int

var scratchDir string

type pictable [][][]uint64

func MkPictable(dx int, dy int) pictable {
	pic := make([][][]uint64, dx) /* type declaration */
	for i := range pic {
		pic[i] = make([][]uint64, dy) /* again the type? */
		for j := range pic[i] {
			pic[i][j] = []uint64{0, 0, 0}
		}
	}
	return pic
}

func absdiff(a uint64, b uint64) uint64 {
	return uint64(math.Abs(float64(a) - float64(b)))
}

func init() {
	subdivisions = flag.Int("subdivisions", 10, "Slices per axis")
	tolerance = flag.Int("tolerance", 100, "Color delta tolerance, higher = more tolerant")
	flag.Parse()

	if flag.NArg() < 1 {
		fmt.Println("usage: imgdedup [options] [<directories>/files]")
		flag.PrintDefaults()
		os.Exit(1)
	}

	image.RegisterFormat("png", "png", png.Decode, png.DecodeConfig)
	image.RegisterFormat("jpeg", "jpeg", jpeg.Decode, jpeg.DecodeConfig)
	image.RegisterFormat("gif", "gif", gif.Decode, gif.DecodeConfig)
}

func init() {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	scratchDir = path.Join(usr.HomeDir, ".imgdedup")

	if _, err := os.Stat(scratchDir); err != nil {
		if os.IsNotExist(err) {
			os.Mkdir(scratchDir, 0700)
		} else {
			log.Fatal(err)
		}
	}
}

func scanImg(file *os.File) (pictable, error) {
	m, _, err := image.Decode(file)
	if err != nil {
		return nil, err
	}
	bounds := m.Bounds()

	avgdata := MkPictable(*subdivisions, *subdivisions)

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			rX := int64(math.Floor((float64(x) / float64(bounds.Max.X)) * float64(*subdivisions)))
			rY := int64(math.Floor((float64(y) / float64(bounds.Max.Y)) * float64(*subdivisions)))

			r, g, b, _ := m.At(x, y).RGBA()
			avgdata[rX][rY][0] += uint64((float32(r) / 65535) * 255)
			avgdata[rX][rY][1] += uint64((float32(g) / 65535) * 255)
			avgdata[rX][rY][2] += uint64((float32(b) / 65535) * 255)
		}
	}

	divisor := uint64((bounds.Max.X / *subdivisions) * (bounds.Max.Y / *subdivisions))
	if divisor == 0 {
		return nil, fmt.Errorf("Image dimensions %d x %d invalid", bounds.Max.X, bounds.Max.Y)
	}

	for rX := 0; rX < *subdivisions; rX++ {
		for rY := 0; rY < *subdivisions; rY++ {
			avgdata[rX][rY][0] = avgdata[rX][rY][0] / divisor
			avgdata[rX][rY][1] = avgdata[rX][rY][1] / divisor
			avgdata[rX][rY][2] = avgdata[rX][rY][2] / divisor
		}
	}

	return avgdata, nil
}

func loadCache(cachename string) (pictable, error) {

	file, err := os.Open(cachename)
	defer file.Close()
	if err != nil {
		return nil, err
	}
	r := bufio.NewReader(file)

	var avgdata pictable

	dec := json.NewDecoder(r)

	err = dec.Decode(&avgdata)
	if err != nil {
		return nil, err
	}

	return avgdata, nil
}

func storeCache(cachename string, avgdata *pictable) {
	fo, err := os.Create(cachename)
	defer fo.Close()
	if err != nil {
		panic(err)
	}

	enc := json.NewEncoder(fo)
	enc.Encode(avgdata)
}

func main() {

	imgdata := make(map[string]pictable)

	fileList := getFiles(flag.Args())

	bar := pb.StartNew(len(fileList))

	for _, imgpath := range fileList {

		bar.Increment()

		file, err := os.Open(imgpath)
		if err != nil {
			log.Fatal(err)
		}

		fExt := strings.ToLower(filepath.Ext(imgpath))
		if fExt == ".png" || fExt == ".jpg" || fExt == ".jpeg" || fExt == ".gif" {

			fi, err := file.Stat()
			if err != nil {
				log.Fatal(err)
			}

			h := md5.New()

			cacheUnit := imgpath + "|" + string(*subdivisions) + "|" + string(fi.Size()) + string(fi.ModTime().Unix())

			io.WriteString(h, cacheUnit)
			cachename := path.Join(scratchDir, fmt.Sprintf("%x", h.Sum(nil))+".tmp")

			var avgdata pictable

			avgdata, err = loadCache(cachename)
			if err != nil {

				avgdata, err = scanImg(file)
				if err != nil {
					log.Print(imgpath, err)
					continue
				}

				storeCache(cachename, &avgdata)
			}

			imgdata[imgpath] = avgdata

			file.Close()

		} else {
			file.Close()
		}
	}

	bar.Finish()

	fileLength := len(fileList)

	for i := 0; i < fileLength-1; i++ {
		for j := i + 1; j < fileLength-1; j++ {

			filename1 := fileList[i]
			filename2 := fileList[j]

			avgdata1, ok1 := imgdata[filename1]
			avgdata2, ok2 := imgdata[filename2]

			if ok1 && ok2 {

				if filename1 == filename2 {
					continue
				}

				var xdiff uint64 = 0

				for rX := 0; rX < *subdivisions; rX++ {
					for rY := 0; rY < *subdivisions; rY++ {
						aa := avgdata1[rX][rY]
						bb := avgdata2[rX][rY]

						xdiff += absdiff(absdiff(absdiff(aa[0], bb[0]), absdiff(aa[1], bb[1])), absdiff(aa[2], bb[2]))
					}
				}

				if xdiff < uint64(*tolerance) {
					fmt.Println(filename1, filename2)
					fmt.Println(xdiff)
				}

			}

		}
	}

}

func getFiles(paths []string) []string {
	var fileList []string

	for _, imgpath := range paths {

		file, err := os.Open(imgpath)
		if err != nil {
			log.Fatal(err)
		}

		fi, err := file.Stat()
		if err != nil {
			log.Fatal(err)
		}

		switch mode := fi.Mode(); {
		case mode.IsDir():
			// fmt.Println("directory")
			filepath.Walk(imgpath, func(path string, f os.FileInfo, err error) error {

				submode := f.Mode()
				if submode.IsRegular() {
					fpath, _ := filepath.Abs(path)
					fileList = append(fileList, fpath)
				}

				return nil
			})
		case mode.IsRegular():
			// fmt.Println("file")
			fpath, _ := filepath.Abs(imgpath)
			fileList = append(fileList, fpath)
		}

		file.Close()

	}

	return fileList
}
<|endoftext|>