<commit_before>/*
 * Copyright 2015 The Kythe Authors. All rights reserved.
 *
 * 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 disksort implements sorting algorithms for sets of data too large to
// fit fully in-memory.  If the number of elements becomes to large, data are
// paged onto the disk.
package disksort // import "kythe.io/kythe/go/util/disksort"

import (
	"bufio"
	"container/heap"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"kythe.io/kythe/go/platform/delimited"
	"kythe.io/kythe/go/util/sortutil"

	"github.com/golang/snappy"
)

// Interface is the standard interface for disk sorting algorithms.  Each
// element in the set of data to be sorted is added to the sorter with Add.
// Once all elements are added, Read can then be called to retrieve each element
// sequentially in sorted order.  Once Read is called, no other operations on
// the sorter are allowed.
type Interface interface {
	// Add adds a new element to the set of data to be sorted.
	Add(i interface{}) error

	// Iterator returns an Iterator to read each of the elements previously added
	// to the set of data to be sorted.  Once Iterator is called, no more data may
	// be added to the sorter.  This is a lower-level version of Read.  Iterator
	// and Read may only be called once.
	Iterator() (Iterator, error)

	// Read calls f on each elements previously added the set of data to be
	// sorted.  If f returns an error, it is returned immediately and f is no
	// longer called.  Once Read is called, no more data may be added to the
	// sorter.  Iterator and Read may only be called once.
	Read(f func(interface{}) error) error
}

// An Iterator reads each element, in order, in a sorted dataset.
type Iterator interface {
	// Next returns the next ordered element.  If none exist, an io.EOF error is
	// returned.
	Next() (interface{}, error)

	// Close releases all of the Iterator's used resources.  Each Iterator must be
	// closed after the client's last call to Next or stray temporary files may be
	// left on disk.
	Close() error
}

// Marshaler is an interface to functions that can binary encode/decode
// elements.
type Marshaler interface {
	// Marshal binary encodes the given element.
	Marshal(interface{}) ([]byte, error)

	// Unmarshal decodes the given encoding of an element.
	Unmarshal([]byte) (interface{}, error)
}

type mergeSorter struct {
	opts MergeOptions

	buffer  []interface{}
	workDir string
	shards  []string

	bufferSize int

	finalized bool
}

// DefaultMaxInMemory is the default number of elements to keep in-memory during
// a merge sort.
const DefaultMaxInMemory = 32000

// DefaultMaxBytesInMemory is the default maximum total size of elements to keep
// in-memory during a merge sort.
const DefaultMaxBytesInMemory = 1024 * 1024 * 256

// MergeOptions specifies how to sort elements.
type MergeOptions struct {
	// Name is optionally used as part of the path for temporary file shards.
	Name string

	// Lesser is the comparison function for sorting the given elements.
	Lesser sortutil.Lesser
	// Marshaler is used for encoding/decoding elements in temporary file shards.
	Marshaler Marshaler

	// WorkDir is the directory used for writing temporary file shards.  If empty,
	// the default directory for temporary files is used.
	WorkDir string

	// MaxInMemory is the maximum number of elements to keep in-memory before
	// paging them to a temporary file shard.  If non-positive, DefaultMaxInMemory
	// is used.
	MaxInMemory int

	// MaxBytesInMemory is the maximum total size of elements to keep in-memory
	// before paging them to a temporary file shard.  An element's size is
	// determined by its `Size() int` method. If non-positive,
	// DefaultMaxBytesInMemory is used.
	MaxBytesInMemory int

	// CompressShards determines whether the temporary file shards should be
	// compressed.
	CompressShards bool
}

type sizer interface{ Size() int }

// NewMergeSorter returns a new disk sorter using a mergesort algorithm.
func NewMergeSorter(opts MergeOptions) (Interface, error) {
	if opts.Lesser == nil {
		return nil, errors.New("missing Lesser")
	} else if opts.Marshaler == nil {
		return nil, errors.New("missing Marshaler")
	}

	name := strings.Replace(opts.Name, string(filepath.Separator), ".", -1)
	if name == "" {
		name = "external.merge.sort"
	}
	dir, err := ioutil.TempDir(opts.WorkDir, name)
	if err != nil {
		return nil, fmt.Errorf("error creating temporary work directory: %v", err)
	}

	if opts.MaxInMemory <= 0 {
		opts.MaxInMemory = DefaultMaxInMemory
	}
	if opts.MaxBytesInMemory <= 0 {
		opts.MaxBytesInMemory = DefaultMaxBytesInMemory
	}

	return &mergeSorter{
		opts:    opts,
		buffer:  make([]interface{}, 0, opts.MaxInMemory),
		workDir: dir,
	}, nil
}

var (
	// ErrAlreadyFinalized is returned from Interface#Add and Interface#Read when
	// Interface#Read has already been called, freezing the sort's inputs/outputs.
	ErrAlreadyFinalized = errors.New("sorter already finalized")
)

// Add implements part of the Interface interface.
func (m *mergeSorter) Add(i interface{}) error {
	if m.finalized {
		return ErrAlreadyFinalized
	}

	m.buffer = append(m.buffer, i)
	if sizer, ok := i.(sizer); ok {
		m.bufferSize += sizer.Size()
	}

	if len(m.buffer) >= m.opts.MaxInMemory || m.bufferSize >= m.opts.MaxBytesInMemory {
		return m.dumpShard()
	}
	return nil
}

type mergeIterator struct {
	buffer []interface{}

	merger    *sortutil.ByLesser
	marshaler Marshaler
	workDir   string
}

const ioBufferSize = 2 << 15

// Iterator implements part of the Interface interface.
func (m *mergeSorter) Iterator() (iter Iterator, err error) {
	if m.finalized {
		return nil, ErrAlreadyFinalized
	}
	m.finalized = true // signal that further operations should fail

	it := &mergeIterator{workDir: m.workDir, marshaler: m.opts.Marshaler}

	if len(m.shards) == 0 {
		// Fast path for a single, in-memory shard
		it.buffer, m.buffer = m.buffer, nil
		sortutil.Sort(m.opts.Lesser, it.buffer)
		return it, nil
	}

	// This is a heap storing the head of each shard.
	merger := &sortutil.ByLesser{
		Lesser: &mergeElementLesser{Lesser: m.opts.Lesser},
	}
	it.merger = merger

	defer func() {
		// Try to cleanup on errors
		if err != nil {
			if cErr := it.Close(); cErr != nil {
				log.Printf("WARNING: error closing Iterator after error: %v", cErr)
			}
		}
	}()

	// Push all of the in-memory elements into the merger heap.
	for _, el := range m.buffer {
		heap.Push(merger, &mergeElement{el: el})
	}
	m.buffer = nil

	// Initialize the merger heap by reading the first element of each shard.
	for _, shard := range m.shards {
		f, err := os.OpenFile(shard, os.O_RDONLY, shardFileMode)
		if err != nil {
			return nil, fmt.Errorf("error opening shard %q: %v", shard, err)
		}

		var r io.Reader
		if m.opts.CompressShards {
			r = snappy.NewReader(f)
		} else {
			r = bufio.NewReaderSize(f, ioBufferSize)
		}

		rd := delimited.NewReader(r)
		first, err := rd.Next()
		if err != nil {
			f.Close()
			return nil, fmt.Errorf("error reading beginning of shard %q: %v", shard, err)
		}
		el, err := m.opts.Marshaler.Unmarshal(first)
		if err != nil {
			f.Close()
			return nil, fmt.Errorf("error unmarshaling beginning of shard %q: %v", shard, err)
		}

		heap.Push(merger, &mergeElement{el: el, rd: rd, f: f})
	}

	return it, nil
}

// Next implements part of the Iterator interface.
func (i *mergeIterator) Next() (interface{}, error) {
	if i.merger == nil {
		// Fast path for a single, in-memory shard
		if len(i.buffer) == 0 {
			return nil, io.EOF
		}
		val := i.buffer[0]
		i.buffer = i.buffer[1:]
		return val, nil
	}

	if i.merger.Len() == 0 {
		return nil, io.EOF
	}

	// While the merger heap is non-empty:
	//   el := pop the head of the heap
	//   pass it to the user-specific function
	//   push the next element el.rd to the merger heap
	x := heap.Pop(i.merger).(*mergeElement)
	el := x.el

	if x.rd != nil {
		// Read and parse the next value on the same shard
		rec, err := x.rd.Next()
		if err != nil {
			_ = x.f.Close()           // ignore errors (file is only open for reading)
			_ = os.Remove(x.f.Name()) // ignore errors (os.RemoveAll used in Close)
			if err != io.EOF {
				return nil, fmt.Errorf("error reading shard: %v", err)
			}
		} else {
			next, err := i.marshaler.Unmarshal(rec)
			if err != nil {
				return nil, fmt.Errorf("error unmarshaling element: %v", err)
			}

			// Reuse mergeElement, push it back onto the merger heap with the next value
			x.el = next
			heap.Push(i.merger, x)
		}
	}

	return el, nil
}

// Close implements part of the Iterator interface.
func (i *mergeIterator) Close() error {
	i.buffer = nil
	if i.merger != nil {
		for i.merger.Len() != 0 {
			x := heap.Pop(i.merger).(*mergeElement)
			_ = x.f.Close() // ignore errors (file is only open for reading)
		}
	}
	if rmErr := os.RemoveAll(i.workDir); rmErr != nil {
		return fmt.Errorf("error removing temporary directory %q: %v", i.workDir, rmErr)
	}
	return nil
}

// Read implements part of the Interface interface.
func (m *mergeSorter) Read(f func(i interface{}) error) (err error) {
	it, err := m.Iterator()
	if err != nil {
		return err
	}
	defer func() {
		if cErr := it.Close(); cErr != nil {
			if err == nil {
				err = cErr
			} else {
				log.Println("WARNING: error closing Iterator:", cErr)
			}
		}
	}()
	for {
		val, err := it.Next()
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
		if err := f(val); err != nil {
			return err
		}
	}
}

const shardFileMode = 0600 | os.ModeExclusive | os.ModeAppend | os.ModeTemporary | os.ModeSticky

func (m *mergeSorter) dumpShard() (err error) {
	defer func() {
		m.buffer = make([]interface{}, 0, m.opts.MaxInMemory)
		m.bufferSize = 0
	}()

	// Create a new shard file
	shardPath := filepath.Join(m.workDir, fmt.Sprintf("shard.%.6d", len(m.shards)))
	file, err := os.OpenFile(shardPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, shardFileMode)
	if err != nil {
		return fmt.Errorf("error creating shard: %v", err)
	}
	defer func() {
		replaceErrIfNil(&err, "error closing shard: %v", file.Close())
	}()

	// Buffer writing to the shard
	var buf interface {
		io.Writer
		Flush() error
	}
	if m.opts.CompressShards {
		buf = snappy.NewBufferedWriter(file)
	} else {
		buf = bufio.NewWriterSize(file, ioBufferSize)
	}

	defer func() {
		replaceErrIfNil(&err, "error flushing shard: %v", buf.Flush())
	}()

	// Sort the in-memory buffer of elements
	sortutil.Sort(m.opts.Lesser, m.buffer)

	// Write each element of the in-memory to shard file, in sorted order
	wr := delimited.NewWriter(buf)
	for len(m.buffer) > 0 {
		rec, err := m.opts.Marshaler.Marshal(m.buffer[0])
		if err != nil {
			return fmt.Errorf("marshaling error: %v", err)
		}
		if _, err := wr.WriteRecord(rec); err != nil {
			return fmt.Errorf("writing error: %v", err)
		}
		m.buffer = m.buffer[1:]
	}

	m.shards = append(m.shards, shardPath)
	return nil
}

func replaceErrIfNil(err *error, s string, newError error) {
	if newError != nil && *err == nil {
		*err = fmt.Errorf(s, newError)
	}
}

type mergeElement struct {
	el interface{}
	rd *delimited.Reader
	f  *os.File
}

type mergeElementLesser struct{ sortutil.Lesser }

// Less implements the sortutil.Lesser interface.
func (m *mergeElementLesser) Less(a, b interface{}) bool {
	x, y := a.(*mergeElement), b.(*mergeElement)
	return m.Lesser.Less(x.el, y.el)
}
<commit_msg>perf(post_processing): use heap.Fix in disksort (#4063)<commit_after>/*
 * Copyright 2015 The Kythe Authors. All rights reserved.
 *
 * 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 disksort implements sorting algorithms for sets of data too large to
// fit fully in-memory.  If the number of elements becomes to large, data are
// paged onto the disk.
package disksort // import "kythe.io/kythe/go/util/disksort"

import (
	"bufio"
	"container/heap"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"kythe.io/kythe/go/platform/delimited"
	"kythe.io/kythe/go/util/sortutil"

	"github.com/golang/snappy"
)

// Interface is the standard interface for disk sorting algorithms.  Each
// element in the set of data to be sorted is added to the sorter with Add.
// Once all elements are added, Read can then be called to retrieve each element
// sequentially in sorted order.  Once Read is called, no other operations on
// the sorter are allowed.
type Interface interface {
	// Add adds a new element to the set of data to be sorted.
	Add(i interface{}) error

	// Iterator returns an Iterator to read each of the elements previously added
	// to the set of data to be sorted.  Once Iterator is called, no more data may
	// be added to the sorter.  This is a lower-level version of Read.  Iterator
	// and Read may only be called once.
	Iterator() (Iterator, error)

	// Read calls f on each elements previously added the set of data to be
	// sorted.  If f returns an error, it is returned immediately and f is no
	// longer called.  Once Read is called, no more data may be added to the
	// sorter.  Iterator and Read may only be called once.
	Read(f func(interface{}) error) error
}

// An Iterator reads each element, in order, in a sorted dataset.
type Iterator interface {
	// Next returns the next ordered element.  If none exist, an io.EOF error is
	// returned.
	Next() (interface{}, error)

	// Close releases all of the Iterator's used resources.  Each Iterator must be
	// closed after the client's last call to Next or stray temporary files may be
	// left on disk.
	Close() error
}

// Marshaler is an interface to functions that can binary encode/decode
// elements.
type Marshaler interface {
	// Marshal binary encodes the given element.
	Marshal(interface{}) ([]byte, error)

	// Unmarshal decodes the given encoding of an element.
	Unmarshal([]byte) (interface{}, error)
}

type mergeSorter struct {
	opts MergeOptions

	buffer  []interface{}
	workDir string
	shards  []string

	bufferSize int

	finalized bool
}

// DefaultMaxInMemory is the default number of elements to keep in-memory during
// a merge sort.
const DefaultMaxInMemory = 32000

// DefaultMaxBytesInMemory is the default maximum total size of elements to keep
// in-memory during a merge sort.
const DefaultMaxBytesInMemory = 1024 * 1024 * 256

// MergeOptions specifies how to sort elements.
type MergeOptions struct {
	// Name is optionally used as part of the path for temporary file shards.
	Name string

	// Lesser is the comparison function for sorting the given elements.
	Lesser sortutil.Lesser
	// Marshaler is used for encoding/decoding elements in temporary file shards.
	Marshaler Marshaler

	// WorkDir is the directory used for writing temporary file shards.  If empty,
	// the default directory for temporary files is used.
	WorkDir string

	// MaxInMemory is the maximum number of elements to keep in-memory before
	// paging them to a temporary file shard.  If non-positive, DefaultMaxInMemory
	// is used.
	MaxInMemory int

	// MaxBytesInMemory is the maximum total size of elements to keep in-memory
	// before paging them to a temporary file shard.  An element's size is
	// determined by its `Size() int` method. If non-positive,
	// DefaultMaxBytesInMemory is used.
	MaxBytesInMemory int

	// CompressShards determines whether the temporary file shards should be
	// compressed.
	CompressShards bool
}

type sizer interface{ Size() int }

// NewMergeSorter returns a new disk sorter using a mergesort algorithm.
func NewMergeSorter(opts MergeOptions) (Interface, error) {
	if opts.Lesser == nil {
		return nil, errors.New("missing Lesser")
	} else if opts.Marshaler == nil {
		return nil, errors.New("missing Marshaler")
	}

	name := strings.Replace(opts.Name, string(filepath.Separator), ".", -1)
	if name == "" {
		name = "external.merge.sort"
	}
	dir, err := ioutil.TempDir(opts.WorkDir, name)
	if err != nil {
		return nil, fmt.Errorf("error creating temporary work directory: %v", err)
	}

	if opts.MaxInMemory <= 0 {
		opts.MaxInMemory = DefaultMaxInMemory
	}
	if opts.MaxBytesInMemory <= 0 {
		opts.MaxBytesInMemory = DefaultMaxBytesInMemory
	}

	return &mergeSorter{
		opts:    opts,
		buffer:  make([]interface{}, 0, opts.MaxInMemory),
		workDir: dir,
	}, nil
}

var (
	// ErrAlreadyFinalized is returned from Interface#Add and Interface#Read when
	// Interface#Read has already been called, freezing the sort's inputs/outputs.
	ErrAlreadyFinalized = errors.New("sorter already finalized")
)

// Add implements part of the Interface interface.
func (m *mergeSorter) Add(i interface{}) error {
	if m.finalized {
		return ErrAlreadyFinalized
	}

	m.buffer = append(m.buffer, i)
	if sizer, ok := i.(sizer); ok {
		m.bufferSize += sizer.Size()
	}

	if len(m.buffer) >= m.opts.MaxInMemory || m.bufferSize >= m.opts.MaxBytesInMemory {
		return m.dumpShard()
	}
	return nil
}

type mergeIterator struct {
	buffer []interface{}

	merger    *sortutil.ByLesser
	marshaler Marshaler
	workDir   string
}

const ioBufferSize = 2 << 15

// Iterator implements part of the Interface interface.
func (m *mergeSorter) Iterator() (iter Iterator, err error) {
	if m.finalized {
		return nil, ErrAlreadyFinalized
	}
	m.finalized = true // signal that further operations should fail

	it := &mergeIterator{workDir: m.workDir, marshaler: m.opts.Marshaler}

	if len(m.shards) == 0 {
		// Fast path for a single, in-memory shard
		it.buffer, m.buffer = m.buffer, nil
		sortutil.Sort(m.opts.Lesser, it.buffer)
		return it, nil
	}

	// This is a heap storing the head of each shard.
	merger := &sortutil.ByLesser{
		Lesser: &mergeElementLesser{Lesser: m.opts.Lesser},
	}
	it.merger = merger

	defer func() {
		// Try to cleanup on errors
		if err != nil {
			if cErr := it.Close(); cErr != nil {
				log.Printf("WARNING: error closing Iterator after error: %v", cErr)
			}
		}
	}()

	// Push all of the in-memory elements into the merger heap.
	for _, el := range m.buffer {
		heap.Push(merger, &mergeElement{el: el})
	}
	m.buffer = nil

	// Initialize the merger heap by reading the first element of each shard.
	for _, shard := range m.shards {
		f, err := os.OpenFile(shard, os.O_RDONLY, shardFileMode)
		if err != nil {
			return nil, fmt.Errorf("error opening shard %q: %v", shard, err)
		}

		var r io.Reader
		if m.opts.CompressShards {
			r = snappy.NewReader(f)
		} else {
			r = bufio.NewReaderSize(f, ioBufferSize)
		}

		rd := delimited.NewReader(r)
		first, err := rd.Next()
		if err != nil {
			f.Close()
			return nil, fmt.Errorf("error reading beginning of shard %q: %v", shard, err)
		}
		el, err := m.opts.Marshaler.Unmarshal(first)
		if err != nil {
			f.Close()
			return nil, fmt.Errorf("error unmarshaling beginning of shard %q: %v", shard, err)
		}

		heap.Push(merger, &mergeElement{el: el, rd: rd, f: f})
	}

	return it, nil
}

// Next implements part of the Iterator interface.
func (i *mergeIterator) Next() (interface{}, error) {
	if i.merger == nil {
		// Fast path for a single, in-memory shard
		if len(i.buffer) == 0 {
			return nil, io.EOF
		}
		val := i.buffer[0]
		i.buffer = i.buffer[1:]
		return val, nil
	}

	if i.merger.Len() == 0 {
		return nil, io.EOF
	}

	// While the merger heap is non-empty:
	//   x := peek the head of the heap
	//   pass x.el to the user-specific function
	//   read the next element in x.rd; fix the merger heap order
	x := i.merger.Slice[0].(*mergeElement)
	el := x.el

	if x.rd == nil {
		heap.Pop(i.merger)
	} else {
		// Read and parse the next value on the same shard
		rec, err := x.rd.Next()
		if err != nil {
			_ = x.f.Close()           // ignore errors (file is only open for reading)
			_ = os.Remove(x.f.Name()) // ignore errors (os.RemoveAll used in Close)
			heap.Pop(i.merger)
			if err != io.EOF {
				return nil, fmt.Errorf("error reading shard: %v", err)
			}
		} else {
			next, err := i.marshaler.Unmarshal(rec)
			if err != nil {
				return nil, fmt.Errorf("error unmarshaling element: %v", err)
			}

			// Reuse mergeElement, reorder it in the merger heap with the next value
			x.el = next
			heap.Fix(i.merger, 0)
		}
	}

	return el, nil
}

// Close implements part of the Iterator interface.
func (i *mergeIterator) Close() error {
	i.buffer = nil
	if i.merger != nil {
		for _, x := range i.merger.Slice {
			el := x.(*mergeElement)
			if el.f != nil {
				el.f.Close() // ignore errors (file is only open for reading)
			}
		}
		i.merger = nil
	}
	if rmErr := os.RemoveAll(i.workDir); rmErr != nil {
		return fmt.Errorf("error removing temporary directory %q: %v", i.workDir, rmErr)
	}
	return nil
}

// Read implements part of the Interface interface.
func (m *mergeSorter) Read(f func(i interface{}) error) (err error) {
	it, err := m.Iterator()
	if err != nil {
		return err
	}
	defer func() {
		if cErr := it.Close(); cErr != nil {
			if err == nil {
				err = cErr
			} else {
				log.Println("WARNING: error closing Iterator:", cErr)
			}
		}
	}()
	for {
		val, err := it.Next()
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}
		if err := f(val); err != nil {
			return err
		}
	}
}

const shardFileMode = 0600 | os.ModeExclusive | os.ModeAppend | os.ModeTemporary | os.ModeSticky

func (m *mergeSorter) dumpShard() (err error) {
	defer func() {
		m.buffer = make([]interface{}, 0, m.opts.MaxInMemory)
		m.bufferSize = 0
	}()

	// Create a new shard file
	shardPath := filepath.Join(m.workDir, fmt.Sprintf("shard.%.6d", len(m.shards)))
	file, err := os.OpenFile(shardPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, shardFileMode)
	if err != nil {
		return fmt.Errorf("error creating shard: %v", err)
	}
	defer func() {
		replaceErrIfNil(&err, "error closing shard: %v", file.Close())
	}()

	// Buffer writing to the shard
	var buf interface {
		io.Writer
		Flush() error
	}
	if m.opts.CompressShards {
		buf = snappy.NewBufferedWriter(file)
	} else {
		buf = bufio.NewWriterSize(file, ioBufferSize)
	}

	defer func() {
		replaceErrIfNil(&err, "error flushing shard: %v", buf.Flush())
	}()

	// Sort the in-memory buffer of elements
	sortutil.Sort(m.opts.Lesser, m.buffer)

	// Write each element of the in-memory to shard file, in sorted order
	wr := delimited.NewWriter(buf)
	for len(m.buffer) > 0 {
		rec, err := m.opts.Marshaler.Marshal(m.buffer[0])
		if err != nil {
			return fmt.Errorf("marshaling error: %v", err)
		}
		if _, err := wr.WriteRecord(rec); err != nil {
			return fmt.Errorf("writing error: %v", err)
		}
		m.buffer = m.buffer[1:]
	}

	m.shards = append(m.shards, shardPath)
	return nil
}

func replaceErrIfNil(err *error, s string, newError error) {
	if newError != nil && *err == nil {
		*err = fmt.Errorf(s, newError)
	}
}

type mergeElement struct {
	el interface{}
	rd *delimited.Reader
	f  *os.File
}

type mergeElementLesser struct{ sortutil.Lesser }

// Less implements the sortutil.Lesser interface.
func (m *mergeElementLesser) Less(a, b interface{}) bool {
	x, y := a.(*mergeElement), b.(*mergeElement)
	return m.Lesser.Less(x.el, y.el)
}
<|endoftext|>


<commit_before>package machinery

import (
	"errors"
	"fmt"
	"log"
	"reflect"

	"github.com/RichardKnop/machinery/v1/backends"
	"github.com/RichardKnop/machinery/v1/signatures"
	"github.com/RichardKnop/machinery/v1/utils"
)

// Worker represents a single worker process
type Worker struct {
	server      *Server
	ConsumerTag string
}

// Launch starts a new worker process. The worker subscribes
// to the default queue and processes incoming registered tasks
func (worker *Worker) Launch() error {
	cnf := worker.server.GetConfig()
	broker := worker.server.GetBroker()

	log.Printf("Launching a worker with the following settings:")
	log.Printf("- Broker: %s", cnf.Broker)
	log.Printf("- ResultBackend: %s", cnf.ResultBackend)
	log.Printf("- Exchange: %s", cnf.Exchange)
	log.Printf("- ExchangeType: %s", cnf.ExchangeType)
	log.Printf("- DefaultQueue: %s", cnf.DefaultQueue)
	log.Printf("- BindingKey: %s", cnf.BindingKey)

	errorsChan := make(chan error)

	go func() {
		for {
			retry, err := broker.StartConsuming(worker.ConsumerTag, worker)

			if retry {
				log.Printf("Going to retry launching the worker. Error: %v", err)
			} else {
				errorsChan <- err // stop the goroutine
				return
			}
		}
	}()

	return <-errorsChan
}

// Quit tears down the running worker process
func (worker *Worker) Quit() {
	worker.server.GetBroker().StopConsuming()
}

// Process handles received tasks and triggers success/error callbacks
func (worker *Worker) Process(signature *signatures.TaskSignature) error {
	// If the task is not registered with this worker, do not continue
	// but only return nil as we do not want to restart the worker process
	if !worker.server.IsTaskRegistered(signature.Name) {
		return nil
	}

	task, err := worker.server.GetRegisteredTask(signature.Name)
	if err != nil {
		return nil
	}

	backend := worker.server.GetBackend()
	// Update task state to RECEIVED
	if err := backend.SetStateReceived(signature); err != nil {
		return fmt.Errorf("Set State Received: %v", err)
	}

	// Get task args and reflect them to proper types
	reflectedTask := reflect.ValueOf(task)
	reflectedArgs, err := reflectArgs(signature.Args)
	if err != nil {
		worker.finalizeError(signature, err)
		return fmt.Errorf("Reflect task args: %v", err)
	}

	// Update task state to STARTED
	if err := backend.SetStateStarted(signature); err != nil {
		return fmt.Errorf("Set State Started: %v", err)
	}

	// Call the task passing in the correct arguments
	results, err := tryCall(reflectedTask, reflectedArgs)

	if err != nil {
		return worker.finalizeError(signature, err)
	}

	return worker.finalizeSuccess(signature, results[0])
}

// Converts []TaskArg to []reflect.Value
func reflectArgs(args []signatures.TaskArg) ([]reflect.Value, error) {
	argValues := make([]reflect.Value, len(args))

	for i, arg := range args {
		argValue, err := utils.ReflectValue(arg.Type, arg.Value)
		if err != nil {
			return nil, err
		}
		argValues[i] = argValue
	}

	return argValues, nil
}

// Attempts to call the task with the supplied arguments.
//
// `err` is set in the return value in two cases:
// 1. The reflected function invocation panics (e.g. due to a mismatched
//    argument list).
// 2. The task func itself returns a non-nil error.
func tryCall(f reflect.Value, args []reflect.Value) (results []reflect.Value, err error) {
	defer func() {
		// Recover from panic and set err.
		if e := recover(); e != nil {
			switch e := e.(type) {
			default:
				err = errors.New("Invoking task caused a panic")
			case error:
				err = e
			case string:
				err = errors.New(e)
			}
		}
	}()

	results = f.Call(args)

	// If an error was returned by the task func, propagate it
	// to the caller via err.
	if !results[1].IsNil() {
		return nil, results[1].Interface().(error)
	}

	return results, err
}

func createTaskResult(value reflect.Value) *backends.TaskResult {
	return &backends.TaskResult{
		Type:  reflect.TypeOf(value.Interface()).String(),
		Value: value.Interface(),
	}
}

// Task succeeded, update state and trigger success callbacks
func (worker *Worker) finalizeSuccess(signature *signatures.TaskSignature, result reflect.Value) error {
	// Update task state to SUCCESS
	backend := worker.server.GetBackend()

	taskResult := createTaskResult(result)
	if err := backend.SetStateSuccess(signature, taskResult); err != nil {
		return fmt.Errorf("Set State Success: %v", err)
	}

	log.Printf("Processed %s. Result = %v", signature.UUID, taskResult.Value)

	// Trigger success callbacks
	for _, successTask := range signature.OnSuccess {
		if signature.Immutable == false {
			// Pass results of the task to success callbacks
			args := append([]signatures.TaskArg{{
				Type:  taskResult.Type,
				Value: taskResult.Value,
			}}, successTask.Args...)
			successTask.Args = args
		}

		worker.server.SendTask(successTask)
	}

	if signature.GroupUUID != "" {
		groupCompleted, err := worker.server.GetBackend().GroupCompleted(
			signature.GroupUUID,
			signature.GroupTaskCount,
		)
		if err != nil {
			return fmt.Errorf("GroupCompleted: %v", err)
		}
		if !groupCompleted {
			return nil
		}

		// Optionally trigger chord callback
		if signature.ChordCallback != nil {
			taskStates, err := worker.server.GetBackend().GroupTaskStates(
				signature.GroupUUID,
				signature.GroupTaskCount,
			)
			if err != nil {
				return nil
			}

			for _, taskState := range taskStates {
				if !taskState.IsSuccess() {
					return nil
				}

				if signature.ChordCallback.Immutable == false {
					// Pass results of the task to the chord callback
					signature.ChordCallback.Args = append(signature.ChordCallback.Args, signatures.TaskArg{
						Type:  taskState.Result.Type,
						Value: taskState.Result.Value,
					})
				}
			}

			_, err = worker.server.SendTask(signature.ChordCallback)
			if err != nil {
				return err
			}
		}

		// Purge group state if we are using AMQP backend and all tasks finished
		if worker.hasAMQPBackend() {
			err = worker.server.backend.PurgeGroupMeta(signature.GroupUUID)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Task failed, update state and trigger error callbacks
func (worker *Worker) finalizeError(signature *signatures.TaskSignature, err error) error {
	// Update task state to FAILURE
	backend := worker.server.GetBackend()
	if err := backend.SetStateFailure(signature, err.Error()); err != nil {
		return fmt.Errorf("Set State Failure: %v", err)
	}

	log.Printf("Failed processing %s. Error = %v", signature.UUID, err)

	// Trigger error callbacks
	for _, errorTask := range signature.OnError {
		// Pass error as a first argument to error callbacks
		args := append([]signatures.TaskArg{{
			Type:  reflect.TypeOf(err).String(),
			Value: reflect.ValueOf(err).Interface(),
		}}, errorTask.Args...)
		errorTask.Args = args
		worker.server.SendTask(errorTask)
	}

	return nil
}

// Returns true if the worker uses AMQP backend
func (worker *Worker) hasAMQPBackend() bool {
	_, ok := worker.server.backend.(*backends.AMQPBackend)
	return ok
}
<commit_msg>Fixed err return in worker<commit_after>package machinery

import (
	"errors"
	"fmt"
	"log"
	"reflect"

	"github.com/RichardKnop/machinery/v1/backends"
	"github.com/RichardKnop/machinery/v1/signatures"
	"github.com/RichardKnop/machinery/v1/utils"
)

// Worker represents a single worker process
type Worker struct {
	server      *Server
	ConsumerTag string
}

// Launch starts a new worker process. The worker subscribes
// to the default queue and processes incoming registered tasks
func (worker *Worker) Launch() error {
	cnf := worker.server.GetConfig()
	broker := worker.server.GetBroker()

	log.Printf("Launching a worker with the following settings:")
	log.Printf("- Broker: %s", cnf.Broker)
	log.Printf("- ResultBackend: %s", cnf.ResultBackend)
	log.Printf("- Exchange: %s", cnf.Exchange)
	log.Printf("- ExchangeType: %s", cnf.ExchangeType)
	log.Printf("- DefaultQueue: %s", cnf.DefaultQueue)
	log.Printf("- BindingKey: %s", cnf.BindingKey)

	errorsChan := make(chan error)

	go func() {
		for {
			retry, err := broker.StartConsuming(worker.ConsumerTag, worker)
			if err != nil {
				errorsChan <- err
			}

			if retry {
				log.Printf("Going to retry launching the worker. Error: %v", err)
			}
		}
	}()

	return <-errorsChan
}

// Quit tears down the running worker process
func (worker *Worker) Quit() {
	worker.server.GetBroker().StopConsuming()
}

// Process handles received tasks and triggers success/error callbacks
func (worker *Worker) Process(signature *signatures.TaskSignature) error {
	// If the task is not registered with this worker, do not continue
	// but only return nil as we do not want to restart the worker process
	if !worker.server.IsTaskRegistered(signature.Name) {
		return nil
	}

	task, err := worker.server.GetRegisteredTask(signature.Name)
	if err != nil {
		return nil
	}

	backend := worker.server.GetBackend()
	// Update task state to RECEIVED
	if err := backend.SetStateReceived(signature); err != nil {
		return fmt.Errorf("Set State Received: %v", err)
	}

	// Get task args and reflect them to proper types
	reflectedTask := reflect.ValueOf(task)
	reflectedArgs, err := reflectArgs(signature.Args)
	if err != nil {
		worker.finalizeError(signature, err)
		return fmt.Errorf("Reflect task args: %v", err)
	}

	// Update task state to STARTED
	if err := backend.SetStateStarted(signature); err != nil {
		return fmt.Errorf("Set State Started: %v", err)
	}

	// Call the task passing in the correct arguments
	results, err := tryCall(reflectedTask, reflectedArgs)

	if err != nil {
		return worker.finalizeError(signature, err)
	}

	return worker.finalizeSuccess(signature, results[0])
}

// Converts []TaskArg to []reflect.Value
func reflectArgs(args []signatures.TaskArg) ([]reflect.Value, error) {
	argValues := make([]reflect.Value, len(args))

	for i, arg := range args {
		argValue, err := utils.ReflectValue(arg.Type, arg.Value)
		if err != nil {
			return nil, err
		}
		argValues[i] = argValue
	}

	return argValues, nil
}

// Attempts to call the task with the supplied arguments.
//
// `err` is set in the return value in two cases:
// 1. The reflected function invocation panics (e.g. due to a mismatched
//    argument list).
// 2. The task func itself returns a non-nil error.
func tryCall(f reflect.Value, args []reflect.Value) (results []reflect.Value, err error) {
	defer func() {
		// Recover from panic and set err.
		if e := recover(); e != nil {
			switch e := e.(type) {
			default:
				err = errors.New("Invoking task caused a panic")
			case error:
				err = e
			case string:
				err = errors.New(e)
			}
		}
	}()

	results = f.Call(args)

	// If an error was returned by the task func, propagate it
	// to the caller via err.
	if !results[1].IsNil() {
		return nil, results[1].Interface().(error)
	}

	return results, err
}

func createTaskResult(value reflect.Value) *backends.TaskResult {
	return &backends.TaskResult{
		Type:  reflect.TypeOf(value.Interface()).String(),
		Value: value.Interface(),
	}
}

// Task succeeded, update state and trigger success callbacks
func (worker *Worker) finalizeSuccess(signature *signatures.TaskSignature, result reflect.Value) error {
	// Update task state to SUCCESS
	backend := worker.server.GetBackend()

	taskResult := createTaskResult(result)
	if err := backend.SetStateSuccess(signature, taskResult); err != nil {
		return fmt.Errorf("Set State Success: %v", err)
	}

	log.Printf("Processed %s. Result = %v", signature.UUID, taskResult.Value)

	// Trigger success callbacks
	for _, successTask := range signature.OnSuccess {
		if signature.Immutable == false {
			// Pass results of the task to success callbacks
			args := append([]signatures.TaskArg{{
				Type:  taskResult.Type,
				Value: taskResult.Value,
			}}, successTask.Args...)
			successTask.Args = args
		}

		worker.server.SendTask(successTask)
	}

	if signature.GroupUUID != "" {
		groupCompleted, err := worker.server.GetBackend().GroupCompleted(
			signature.GroupUUID,
			signature.GroupTaskCount,
		)
		if err != nil {
			return fmt.Errorf("GroupCompleted: %v", err)
		}
		if !groupCompleted {
			return nil
		}

		// Optionally trigger chord callback
		if signature.ChordCallback != nil {
			taskStates, err := worker.server.GetBackend().GroupTaskStates(
				signature.GroupUUID,
				signature.GroupTaskCount,
			)
			if err != nil {
				return nil
			}

			for _, taskState := range taskStates {
				if !taskState.IsSuccess() {
					return nil
				}

				if signature.ChordCallback.Immutable == false {
					// Pass results of the task to the chord callback
					signature.ChordCallback.Args = append(signature.ChordCallback.Args, signatures.TaskArg{
						Type:  taskState.Result.Type,
						Value: taskState.Result.Value,
					})
				}
			}

			_, err = worker.server.SendTask(signature.ChordCallback)
			if err != nil {
				return err
			}
		}

		// Purge group state if we are using AMQP backend and all tasks finished
		if worker.hasAMQPBackend() {
			err = worker.server.backend.PurgeGroupMeta(signature.GroupUUID)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// Task failed, update state and trigger error callbacks
func (worker *Worker) finalizeError(signature *signatures.TaskSignature, err error) error {
	// Update task state to FAILURE
	backend := worker.server.GetBackend()
	if err := backend.SetStateFailure(signature, err.Error()); err != nil {
		return fmt.Errorf("Set State Failure: %v", err)
	}

	log.Printf("Failed processing %s. Error = %v", signature.UUID, err)

	// Trigger error callbacks
	for _, errorTask := range signature.OnError {
		// Pass error as a first argument to error callbacks
		args := append([]signatures.TaskArg{{
			Type:  reflect.TypeOf(err).String(),
			Value: reflect.ValueOf(err).Interface(),
		}}, errorTask.Args...)
		errorTask.Args = args
		worker.server.SendTask(errorTask)
	}

	return nil
}

// Returns true if the worker uses AMQP backend
func (worker *Worker) hasAMQPBackend() bool {
	_, ok := worker.server.backend.(*backends.AMQPBackend)
	return ok
}
<|endoftext|>


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

package rpc

import (
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/getsentry/raven-go"
)

// ----------------------------------------------------------------------------
// Codec
// ----------------------------------------------------------------------------

// Codec creates a CodecRequest to process each request.
type Codec interface {
	NewRequest(*http.Request) []CodecRequest
}

// CodecRequest decodes a request and encodes a response using a specific
// serialization scheme.
type CodecRequest interface {
	// Reads the request and returns the RPC method name.
	Method() (string, error)
	// Reads the request filling the RPC method args.
	ReadRequest(interface{}) error
	// Writes the response using the RPC method reply.
	WriteResponse(http.ResponseWriter, interface{})
	// Writes an error produced by the server.
	WriteError(w http.ResponseWriter, status int, err error)
}

// ----------------------------------------------------------------------------
// Server
// ----------------------------------------------------------------------------

// NewServer returns a new RPC server.
func NewServer() *Server {
	return &Server{
		codecs:   make(map[string]Codec),
		services: new(serviceMap),
	}
}

// Server serves registered RPC services using registered codecs.
type Server struct {
	codecs   map[string]Codec
	services *serviceMap
}

// RegisterCodec adds a new codec to the server.
//
// Codecs are defined to process a given serialization scheme, e.g., JSON or
// XML. A codec is chosen based on the "Content-Type" header from the request,
// excluding the charset definition.
func (s *Server) RegisterCodec(codec Codec, contentType string) {
	s.codecs[strings.ToLower(contentType)] = codec
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method has three arguments: *http.Request, *args, *reply.
//    - All three arguments are pointers.
//    - The second and third arguments are exported or local.
//    - The method has return type error.
//
// All other methods are ignored.
func (s *Server) RegisterService(receiver interface{}, name string) error {
	return s.services.register(receiver, name)
}

// HasMethod returns true if the given method is registered.
//
// The method uses a dotted notation as in "Service.Method".
func (s *Server) HasMethod(method string) bool {
	if _, _, err := s.services.get(method); err == nil {
		return true
	}
	return false
}

// ServeHTTP
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	st := time.Now()
	if r.Method != "POST" {
		WriteError(w, 405, "rpc: POST method required, received "+r.Method)
		return
	}
	contentType := r.Header.Get("Content-Type")
	idx := strings.Index(contentType, ";")
	if idx != -1 {
		contentType = contentType[:idx]
	}
	codec := s.codecs[strings.ToLower(contentType)]
	if codec == nil {
		WriteError(w, 415, "rpc: unrecognized Content-Type: "+contentType)
		return
	}
	// Prevents Internet Explorer from MIME-sniffing a response away
	// from the declared content-type
	w.Header().Set("x-content-type-options", "nosniff")
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	// Create a new codec request.
	//codecReq := codec.NewRequest(r)
	codecRequests := codec.NewRequest(r)
	var wg sync.WaitGroup
	for _, cr := range codecRequests {
		// Get service method to be called.
		wg.Add(1)
		go func(codecReq CodecRequest) {
			defer wg.Done()
			method, errMethod := codecReq.Method()
			if errMethod != nil {
				codecReq.WriteError(w, 400, errMethod)
				return
			}
			serviceSpec, methodSpec, errGet := s.services.get(method)
			if errGet != nil {
				codecReq.WriteError(w, 400, errGet)
				return
			}
			// Decode the args.
			args := reflect.New(methodSpec.argsType)
			if errRead := codecReq.ReadRequest(args.Interface()); errRead != nil {
				codecReq.WriteError(w, 400, errRead)
				return
			}
			// Call the service method.
			reply := reflect.New(methodSpec.replyType)

			var errValue []reflect.Value

			errContext, errID := raven.CapturePanic(func() {
				errValue = methodSpec.method.Func.Call([]reflect.Value{
					serviceSpec.rcvr,
					reflect.ValueOf(r),
					args,
					reply,
				})
			}, nil)
			if errID != "" {
				codecReq.WriteError(w, http.StatusInternalServerError, fmt.Errorf(
					"rpc: panic occured and reported to sentry: %s\n", errID,
				))
				codecReq.WriteResponse(w, errContext)
				return
			}
			// Cast the result to error if needed.
			var errResult error
			errInter := errValue[0].Interface()
			if errInter != nil {
				errResult = errInter.(error)
			}
			// Encode the response.
			if errResult == nil {
				codecReq.WriteResponse(w, reply.Interface())
			} else {
				codecReq.WriteError(w, 400, errResult)
			}
			return
		}(cr)
	}
	wg.Wait()
	et := time.Now()
	fmt.Printf("Request from %s completed in %v\n", r.RemoteAddr, et.Sub(st))
	return
}

func WriteError(w http.ResponseWriter, status int, msg string) {
	w.WriteHeader(status)
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	fmt.Fprint(w, msg)
}
<commit_msg>Add per-method logging<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rpc

import (
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/getsentry/raven-go"
)

// ----------------------------------------------------------------------------
// Codec
// ----------------------------------------------------------------------------

// Codec creates a CodecRequest to process each request.
type Codec interface {
	NewRequest(*http.Request) []CodecRequest
}

// CodecRequest decodes a request and encodes a response using a specific
// serialization scheme.
type CodecRequest interface {
	// Reads the request and returns the RPC method name.
	Method() (string, error)
	// Reads the request filling the RPC method args.
	ReadRequest(interface{}) error
	// Writes the response using the RPC method reply.
	WriteResponse(http.ResponseWriter, interface{})
	// Writes an error produced by the server.
	WriteError(w http.ResponseWriter, status int, err error)
}

// ----------------------------------------------------------------------------
// Server
// ----------------------------------------------------------------------------

// NewServer returns a new RPC server.
func NewServer() *Server {
	return &Server{
		codecs:   make(map[string]Codec),
		services: new(serviceMap),
	}
}

// Server serves registered RPC services using registered codecs.
type Server struct {
	codecs   map[string]Codec
	services *serviceMap
}

// RegisterCodec adds a new codec to the server.
//
// Codecs are defined to process a given serialization scheme, e.g., JSON or
// XML. A codec is chosen based on the "Content-Type" header from the request,
// excluding the charset definition.
func (s *Server) RegisterCodec(codec Codec, contentType string) {
	s.codecs[strings.ToLower(contentType)] = codec
}

// RegisterService adds a new service to the server.
//
// The name parameter is optional: if empty it will be inferred from
// the receiver type name.
//
// Methods from the receiver will be extracted if these rules are satisfied:
//
//    - The receiver is exported (begins with an upper case letter) or local
//      (defined in the package registering the service).
//    - The method name is exported.
//    - The method has three arguments: *http.Request, *args, *reply.
//    - All three arguments are pointers.
//    - The second and third arguments are exported or local.
//    - The method has return type error.
//
// All other methods are ignored.
func (s *Server) RegisterService(receiver interface{}, name string) error {
	return s.services.register(receiver, name)
}

// HasMethod returns true if the given method is registered.
//
// The method uses a dotted notation as in "Service.Method".
func (s *Server) HasMethod(method string) bool {
	if _, _, err := s.services.get(method); err == nil {
		return true
	}
	return false
}

// ServeHTTP
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	st := time.Now()
	if r.Method != "POST" {
		WriteError(w, 405, "rpc: POST method required, received "+r.Method)
		return
	}
	contentType := r.Header.Get("Content-Type")
	idx := strings.Index(contentType, ";")
	if idx != -1 {
		contentType = contentType[:idx]
	}
	codec := s.codecs[strings.ToLower(contentType)]
	if codec == nil {
		WriteError(w, 415, "rpc: unrecognized Content-Type: "+contentType)
		return
	}
	// Prevents Internet Explorer from MIME-sniffing a response away
	// from the declared content-type
	w.Header().Set("x-content-type-options", "nosniff")
	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	// Create a new codec request.
	//codecReq := codec.NewRequest(r)
	codecRequests := codec.NewRequest(r)
	var wg sync.WaitGroup
	for _, cr := range codecRequests {
		// Get service method to be called.
		wg.Add(1)
		go func(codecReq CodecRequest) {
			defer wg.Done()
			method, errMethod := codecReq.Method()
			if errMethod != nil {
				codecReq.WriteError(w, 400, errMethod)
				return
			}
			serviceSpec, methodSpec, errGet := s.services.get(method)
			if errGet != nil {
				codecReq.WriteError(w, 400, errGet)
				return
			}
			// Decode the args.
			args := reflect.New(methodSpec.argsType)
			if errRead := codecReq.ReadRequest(args.Interface()); errRead != nil {
				codecReq.WriteError(w, 400, errRead)
				return
			}
			// Call the service method.
			reply := reflect.New(methodSpec.replyType)

			var errValue []reflect.Value

			errContext, errID := raven.CapturePanic(func() {
				errValue = methodSpec.method.Func.Call([]reflect.Value{
					serviceSpec.rcvr,
					reflect.ValueOf(r),
					args,
					reply,
				})
			}, nil)
			if errID != "" {
				codecReq.WriteError(w, http.StatusInternalServerError, fmt.Errorf(
					"rpc: panic occured and reported to sentry: %s\n", errID,
				))
				codecReq.WriteResponse(w, errContext)
				return
			}
			// Cast the result to error if needed.
			var errResult error
			errInter := errValue[0].Interface()
			if errInter != nil {
				errResult = errInter.(error)
			}
			// Encode the response.
			if errResult == nil {
				codecReq.WriteResponse(w, reply.Interface())
			} else {
				codecReq.WriteError(w, 400, errResult)
			}
			et := time.Now()
			fmt.Printf("Method call %s completed in %v\n", method, et.Sub(st))
			return
		}(cr)
	}
	wg.Wait()
	et := time.Now()
	fmt.Printf("Request from %s completed in %v\n", r.RemoteAddr, et.Sub(st))
	return
}

func WriteError(w http.ResponseWriter, status int, msg string) {
	w.WriteHeader(status)
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	fmt.Fprint(w, msg)
}
<|endoftext|>


<commit_before>package memberlist

import (
	"bytes"
	"fmt"
	"reflect"
	"testing"
	"time"
)

func TestEncodeDecode(t *testing.T) {
	msg := &ping{SeqNo: 100}
	buf, err := encode(pingMsg, msg)
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}
	var out ping
	if err := decode(buf.Bytes()[1:], &out); err != nil {
		t.Fatalf("unexpected err: %s", err)
	}
	if msg.SeqNo != out.SeqNo {
		t.Fatalf("bad sequence no")
	}
}

func TestRandomOffset(t *testing.T) {
	vals := make(map[int]struct{})
	for i := 0; i < 100; i++ {
		offset := randomOffset(2 << 30)
		if _, ok := vals[offset]; ok {
			t.Fatalf("got collision")
		}
		vals[offset] = struct{}{}
	}
}

func TestRandomOffset_Zero(t *testing.T) {
	offset := randomOffset(0)
	if offset != 0 {
		t.Fatalf("bad offset")
	}
}

func TestNotify(t *testing.T) {
	ch1 := make(chan *Node, 1)
	ch2 := make(chan *Node, 1)

	// Make sure ch1 is full
	ch1 <- &Node{Name: "test"}

	// Notify ch1
	n := &Node{Name: "Push"}
	notify(ch1, n)

	v := <-ch1
	if v.Name != "test" {
		t.Fatalf("bad name")
	}

	// Test receive
	select {
	case v := <-ch1:
		t.Fatalf("bad node %v", v)
	default:
	}

	// Test ch2
	notify(ch2, n)
	select {
	case v := <-ch2:
		if v != n {
			t.Fatalf("bad node %v", v)
		}
	default:
		t.Fatalf("nothing on channel")
	}
}

func TestSuspicionTimeout(t *testing.T) {
	timeout := suspicionTimeout(3, 10, time.Second)
	if timeout != 6*time.Second {
		t.Fatalf("bad timeout")
	}
}

func TestRetransmitLimit(t *testing.T) {
	lim := retransmitLimit(3, 0)
	if lim != 0 {
		t.Fatalf("bad val %v", lim)
	}
	lim = retransmitLimit(3, 1)
	if lim != 3 {
		t.Fatalf("bad val %v", lim)
	}
	lim = retransmitLimit(3, 99)
	if lim != 6 {
		t.Fatalf("bad val %v", lim)
	}
}

func TestShuffleNodes(t *testing.T) {
	orig := []*NodeState{
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
	}
	nodes := make([]*NodeState, len(orig))
	copy(nodes[:], orig[:])

	if !reflect.DeepEqual(nodes, orig) {
		t.Fatalf("should match")
	}

	shuffleNodes(nodes)

	if reflect.DeepEqual(nodes, orig) {
		t.Fatalf("should not match")
	}
}

func TestMoveDeadNodes(t *testing.T) {
	nodes := []*NodeState{
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
	}

	idx := moveDeadNodes(nodes)
	if idx != 3 {
		t.Fatalf("bad index")
	}
	for i := 0; i < idx; i++ {
		if nodes[i].State != StateAlive {
			t.Fatalf("Bad state %d", i)
		}
	}
	for i := idx; i < len(nodes); i++ {
		if nodes[i].State != StateDead {
			t.Fatalf("Bad state %d", i)
		}
	}
}

func TestKRandomNodes(t *testing.T) {
	nodes := []*NodeState{}
	for i := 0; i < 30; i++ {
		// Half the nodes are in a bad state
		state := StateAlive
		switch i % 3 {
		case 0:
			state = StateAlive
		case 1:
			state = StateSuspect
		case 2:
			state = StateDead
		}
		nodes = append(nodes, &NodeState{
			Node: Node{
				Name: fmt.Sprintf("test%d", i),
			},
			State: state,
		})
	}

	s1 := kRandomNodes(3, []string{"test0"}, nodes)
	s2 := kRandomNodes(3, []string{"test0"}, nodes)
	s3 := kRandomNodes(3, []string{"test0"}, nodes)

	if reflect.DeepEqual(s1, s2) {
		t.Fatalf("unexpected equal")
	}
	if reflect.DeepEqual(s1, s3) {
		t.Fatalf("unexpected equal")
	}
	if reflect.DeepEqual(s2, s3) {
		t.Fatalf("unexpected equal")
	}

	for _, s := range [][]*NodeState{s1, s2, s3} {
		if len(s) != 3 {
			t.Fatalf("bad len")
		}
		for _, n := range s {
			if n.Name == "test0" {
				t.Fatalf("Bad name")
			}
			if n.State != StateAlive {
				t.Fatalf("Bad state")
			}
		}
	}
}

func TestMakeCompoundMessage(t *testing.T) {
	msg := &ping{SeqNo: 100}
	buf, err := encode(pingMsg, msg)
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	msgs := [][]byte{buf.Bytes(), buf.Bytes(), buf.Bytes()}
	compound := makeCompoundMessage(msgs)

	if compound.Len() != 3*buf.Len()+3*compoundOverhead+compoundHeaderOverhead {
		t.Fatalf("bad len")
	}
}

func TestDecodeCompoundMessage(t *testing.T) {
	msg := &ping{SeqNo: 100}
	buf, err := encode(pingMsg, msg)
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	msgs := []*bytes.Buffer{buf, buf, buf}
	compound := makeCompoundMessage(msgs)

	trunc, parts, err := decodeCompoundMessage(compound.Bytes()[1:])
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}
	if trunc != 0 {
		t.Fatalf("should not truncate")
	}
	if len(parts) != 3 {
		t.Fatalf("bad parts")
	}
	for _, p := range parts {
		if len(p) != buf.Len() {
			t.Fatalf("bad part len")
		}
	}
}

func TestDecodeCompoundMessage_Trunc(t *testing.T) {
	msg := &ping{SeqNo: 100}
	buf, err := encode(pingMsg, msg)
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	msgs := []*bytes.Buffer{buf, buf, buf}
	compound := makeCompoundMessage(msgs)

	trunc, parts, err := decodeCompoundMessage(compound.Bytes()[1:30])
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}
	if trunc != 1 {
		t.Fatalf("truncate: %d", trunc)
	}
	if len(parts) != 2 {
		t.Fatalf("bad parts")
	}
	for _, p := range parts {
		if len(p) != buf.Len() {
			t.Fatalf("bad part len")
		}
	}
}
<commit_msg>Fixing util tests<commit_after>package memberlist

import (
	"fmt"
	"reflect"
	"testing"
	"time"
)

func TestEncodeDecode(t *testing.T) {
	msg := &ping{SeqNo: 100}
	buf, err := encode(pingMsg, msg)
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}
	var out ping
	if err := decode(buf.Bytes()[1:], &out); err != nil {
		t.Fatalf("unexpected err: %s", err)
	}
	if msg.SeqNo != out.SeqNo {
		t.Fatalf("bad sequence no")
	}
}

func TestRandomOffset(t *testing.T) {
	vals := make(map[int]struct{})
	for i := 0; i < 100; i++ {
		offset := randomOffset(2 << 30)
		if _, ok := vals[offset]; ok {
			t.Fatalf("got collision")
		}
		vals[offset] = struct{}{}
	}
}

func TestRandomOffset_Zero(t *testing.T) {
	offset := randomOffset(0)
	if offset != 0 {
		t.Fatalf("bad offset")
	}
}

func TestNotify(t *testing.T) {
	ch1 := make(chan *Node, 1)
	ch2 := make(chan *Node, 1)

	// Make sure ch1 is full
	ch1 <- &Node{Name: "test"}

	// Notify ch1
	n := &Node{Name: "Push"}
	notify(ch1, n)

	v := <-ch1
	if v.Name != "test" {
		t.Fatalf("bad name")
	}

	// Test receive
	select {
	case v := <-ch1:
		t.Fatalf("bad node %v", v)
	default:
	}

	// Test ch2
	notify(ch2, n)
	select {
	case v := <-ch2:
		if v != n {
			t.Fatalf("bad node %v", v)
		}
	default:
		t.Fatalf("nothing on channel")
	}
}

func TestSuspicionTimeout(t *testing.T) {
	timeout := suspicionTimeout(3, 10, time.Second)
	if timeout != 6*time.Second {
		t.Fatalf("bad timeout")
	}
}

func TestRetransmitLimit(t *testing.T) {
	lim := retransmitLimit(3, 0)
	if lim != 0 {
		t.Fatalf("bad val %v", lim)
	}
	lim = retransmitLimit(3, 1)
	if lim != 3 {
		t.Fatalf("bad val %v", lim)
	}
	lim = retransmitLimit(3, 99)
	if lim != 6 {
		t.Fatalf("bad val %v", lim)
	}
}

func TestShuffleNodes(t *testing.T) {
	orig := []*NodeState{
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
	}
	nodes := make([]*NodeState, len(orig))
	copy(nodes[:], orig[:])

	if !reflect.DeepEqual(nodes, orig) {
		t.Fatalf("should match")
	}

	shuffleNodes(nodes)

	if reflect.DeepEqual(nodes, orig) {
		t.Fatalf("should not match")
	}
}

func TestMoveDeadNodes(t *testing.T) {
	nodes := []*NodeState{
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateAlive,
		},
		&NodeState{
			State: StateDead,
		},
		&NodeState{
			State: StateAlive,
		},
	}

	idx := moveDeadNodes(nodes)
	if idx != 3 {
		t.Fatalf("bad index")
	}
	for i := 0; i < idx; i++ {
		if nodes[i].State != StateAlive {
			t.Fatalf("Bad state %d", i)
		}
	}
	for i := idx; i < len(nodes); i++ {
		if nodes[i].State != StateDead {
			t.Fatalf("Bad state %d", i)
		}
	}
}

func TestKRandomNodes(t *testing.T) {
	nodes := []*NodeState{}
	for i := 0; i < 30; i++ {
		// Half the nodes are in a bad state
		state := StateAlive
		switch i % 3 {
		case 0:
			state = StateAlive
		case 1:
			state = StateSuspect
		case 2:
			state = StateDead
		}
		nodes = append(nodes, &NodeState{
			Node: Node{
				Name: fmt.Sprintf("test%d", i),
			},
			State: state,
		})
	}

	s1 := kRandomNodes(3, []string{"test0"}, nodes)
	s2 := kRandomNodes(3, []string{"test0"}, nodes)
	s3 := kRandomNodes(3, []string{"test0"}, nodes)

	if reflect.DeepEqual(s1, s2) {
		t.Fatalf("unexpected equal")
	}
	if reflect.DeepEqual(s1, s3) {
		t.Fatalf("unexpected equal")
	}
	if reflect.DeepEqual(s2, s3) {
		t.Fatalf("unexpected equal")
	}

	for _, s := range [][]*NodeState{s1, s2, s3} {
		if len(s) != 3 {
			t.Fatalf("bad len")
		}
		for _, n := range s {
			if n.Name == "test0" {
				t.Fatalf("Bad name")
			}
			if n.State != StateAlive {
				t.Fatalf("Bad state")
			}
		}
	}
}

func TestMakeCompoundMessage(t *testing.T) {
	msg := &ping{SeqNo: 100}
	buf, err := encode(pingMsg, msg)
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	msgs := [][]byte{buf.Bytes(), buf.Bytes(), buf.Bytes()}
	compound := makeCompoundMessage(msgs)

	if compound.Len() != 3*buf.Len()+3*compoundOverhead+compoundHeaderOverhead {
		t.Fatalf("bad len")
	}
}

func TestDecodeCompoundMessage(t *testing.T) {
	msg := &ping{SeqNo: 100}
	buf, err := encode(pingMsg, msg)
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	msgs := [][]byte{buf.Bytes(), buf.Bytes(), buf.Bytes()}
	compound := makeCompoundMessage(msgs)

	trunc, parts, err := decodeCompoundMessage(compound.Bytes()[1:])
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}
	if trunc != 0 {
		t.Fatalf("should not truncate")
	}
	if len(parts) != 3 {
		t.Fatalf("bad parts")
	}
	for _, p := range parts {
		if len(p) != buf.Len() {
			t.Fatalf("bad part len")
		}
	}
}

func TestDecodeCompoundMessage_Trunc(t *testing.T) {
	msg := &ping{SeqNo: 100}
	buf, err := encode(pingMsg, msg)
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}

	msgs := [][]byte{buf.Bytes(), buf.Bytes(), buf.Bytes()}
	compound := makeCompoundMessage(msgs)

	trunc, parts, err := decodeCompoundMessage(compound.Bytes()[1:30])
	if err != nil {
		t.Fatalf("unexpected err: %s", err)
	}
	if trunc != 1 {
		t.Fatalf("truncate: %d", trunc)
	}
	if len(parts) != 2 {
		t.Fatalf("bad parts")
	}
	for _, p := range parts {
		if len(p) != buf.Len() {
			t.Fatalf("bad part len")
		}
	}
}
<|endoftext|>


<commit_before>package utils

// GitRepo type manages a VCS repository with Git
type GitRepo struct {
	// Local path to the repository
	Path string

	// Upstream URL of the Git repository
	Upstream string
}

// NewGitRepo creates a new Git repository
func NewGitRepo(path, upstream string) *GitRepo {
	return &GitRepo{
		Path:     path,
		Upstream: upstream,
	}
}
<commit_msg>utils: implement Fetch and Pull methods for GitRepo type<commit_after>package utils

import "os/exec"

// GitRepo type manages a VCS repository with Git
type GitRepo struct {
	// Local path to the repository
	Path string

	// Upstream URL of the Git repository
	Upstream string

	// Path to the Git tool
	git string
}

// NewGitRepo creates a new Git repository
func NewGitRepo(path, upstream string) (*GitRepo, error) {
	path, err := exec.LookPath("git")
	if err != nil {
		return nil, err
	}

	repo := &GitRepo{
		Path:     path,
		Upstream: upstream,
		cmd:      path,
	}

	return repo, nil
}

// Fetch fetches from the given remote
func (gr *GitRepo) Fetch(remote string) ([]byte, error) {
	return exec.Command(gr.git, "fetch", remote).CombinedOutput()
}

// Pull pulls from the given remote and merges changes into the
// local branch
func (gr *GitRepo) Pull(remote, branch string) ([]byte, error) {
	out, err := exec.Command(gr.git, "checkout", branch).CombinedOut()
	if err != nil {
		return out, err
	}

	return exec.Command(gr.git, "pull", remote).CombinedOutput()
}
<|endoftext|>


<commit_before>package gpool

import "testing"

func test_uuid_uniqueness(N int, get func() string, t *testing.T) {
	track := map[string]bool{}
	for i := 0; i < N; i++ {
		uid := get()
		if _, ok := track[uid]; ok {
			t.Fatal("uuid collision: ", uid)
		}
	}
}

func Test_uuid_Unique(t *testing.T) {
	N := 10000
	test_uuid_uniqueness(N, uuid, t)
}

func Test_uuid_Concurrent(t *testing.T) {
	N := 100
	block := make(chan bool)
	recv := make(chan string, N)

	for i := 0; i < N; i++ {
		go func() {
			<-block
			recv <- uuid()
		}()
	}
	close(block)
	test_uuid_uniqueness(N, func() string { return <-recv }, t)
}

func Benchmark_uuid(b *testing.B) {
	for i := 0; i < b.N; i++ {
		uuid()
	}
}
<commit_msg>uuid_uniqueness doesn't track uid<commit_after>package gpool

import "testing"

func test_uuid_uniqueness(N int, get func() string, t *testing.T) {
	track := map[string]bool{}
	for i := 0; i < N; i++ {
		uid := get()
		if _, ok := track[uid]; ok {
			t.Fatal("uuid collision: ", uid)
		}
		track[uid] = true
	}
}

func Test_uuid_Unique(t *testing.T) {
	N := 10000
	test_uuid_uniqueness(N, uuid, t)
}

func Test_uuid_Concurrent(t *testing.T) {
	N := 100
	block := make(chan bool)
	recv := make(chan string, N)

	for i := 0; i < N; i++ {
		go func() {
			<-block
			recv <- uuid()
		}()
	}
	close(block)
	test_uuid_uniqueness(N, func() string { return <-recv }, t)
}

func Benchmark_uuid(b *testing.B) {
	for i := 0; i < b.N; i++ {
		uuid()
	}
}
<|endoftext|>


<commit_before>package ebsvolume

import (
	"bytes"
	"context"
	"fmt"
	"sync"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"

	//"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	"github.com/hashicorp/packer-plugin-sdk/multistep"
	"github.com/hashicorp/packer-plugin-sdk/packer"
	"github.com/hashicorp/packer/builder/amazon/common"
)

// Define a mock struct to be used in unit tests for common aws steps.
type mockEC2Conn struct {
	ec2iface.EC2API
	Config *aws.Config

	// Counters to figure out what code path was taken
	copyImageCount       int
	describeImagesCount  int
	deregisterImageCount int
	deleteSnapshotCount  int
	waitCount            int

	lock sync.Mutex
}

func (m *mockEC2Conn) CreateSnapshot(input *ec2.CreateSnapshotInput) (*ec2.Snapshot, error) {
	snap := &ec2.Snapshot{
		// This isn't typical amazon format, but injecting the volume id into
		// this field lets us verify that the right volume was snapshotted with
		// a simple string comparison
		SnapshotId: aws.String(fmt.Sprintf("snap-of-%s", *input.VolumeId)),
	}

	return snap, nil
}

func (m *mockEC2Conn) WaitUntilSnapshotCompletedWithContext(aws.Context, *ec2.DescribeSnapshotsInput, ...request.WaiterOption) error {
	return nil
}

func getMockConn(config *common.AccessConfig, target string) (ec2iface.EC2API, error) {
	mockConn := &mockEC2Conn{
		Config: aws.NewConfig(),
	}
	return mockConn, nil
}

// Create statebag for running test
func tState(t *testing.T) multistep.StateBag {
	state := new(multistep.BasicStateBag)
	state.Put("ui", &packer.BasicUi{
		Reader: new(bytes.Buffer),
		Writer: new(bytes.Buffer),
	})
	// state.Put("amis", map[string]string{"us-east-1": "ami-12345"})
	// state.Put("snapshots", map[string][]string{"us-east-1": {"snap-0012345"}})
	conn, _ := getMockConn(&common.AccessConfig{}, "us-east-2")

	state.Put("ec2", conn)
	// Store a fake instance that contains a block device that matches the
	// volumes defined in the config above
	state.Put("instance", &ec2.Instance{
		InstanceId: aws.String("instance-id"),
		BlockDeviceMappings: []*ec2.InstanceBlockDeviceMapping{
			{
				DeviceName: aws.String("/dev/xvda"),
				Ebs: &ec2.EbsInstanceBlockDevice{
					VolumeId: aws.String("vol-1234"),
				},
			},
			{
				DeviceName: aws.String("/dev/xvdb"),
				Ebs: &ec2.EbsInstanceBlockDevice{
					VolumeId: aws.String("vol-5678"),
				},
			},
		},
	})
	return state
}

func TestStepSnapshot_run_simple(t *testing.T) {
	var b Builder
	config := testConfig() //from builder_test

	//Set some snapshot settings
	config["ebs_volumes"] = []map[string]interface{}{
		{
			"device_name":           "/dev/xvdb",
			"volume_size":           "32",
			"delete_on_termination": true,
			"snapshot_volume":       true,
		},
	}

	generatedData, warnings, err := b.Prepare(config)
	if len(warnings) > 0 {
		t.Fatalf("bad: %#v", warnings)
	}
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}
	if len(generatedData) == 0 {
		t.Fatalf("Generated data should not be empty")
	}

	state := tState(t)

	accessConfig := common.FakeAccessConfig()

	step := stepSnapshotEBSVolumes{
		PollingConfig: new(common.AWSPollingConfig),
		AccessConfig:  accessConfig,
		VolumeMapping: b.config.VolumeMappings,
		Ctx:           b.config.ctx,
	}

	step.Run(context.Background(), state)

	if len(step.snapshotMap) != 1 {
		t.Fatalf("Missing Snapshot from step")
	}

	if volmapping := step.snapshotMap["snap-of-vol-5678"]; volmapping == nil {
		t.Fatalf("Didn't snapshot correct volume: Map is %#v", step.snapshotMap)
	}
}

func TestStepSnapshot_run_no_snaps(t *testing.T) {
	var b Builder
	config := testConfig() //from builder_test

	//Set some snapshot settings
	config["ebs_volumes"] = []map[string]interface{}{
		{
			"device_name":           "/dev/xvdb",
			"volume_size":           "32",
			"delete_on_termination": true,
			"snapshot_volume":       false,
		},
	}

	generatedData, warnings, err := b.Prepare(config)
	if len(warnings) > 0 {
		t.Fatalf("bad: %#v", warnings)
	}
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}
	if len(generatedData) == 0 {
		t.Fatalf("Generated data should not be empty")
	}

	state := tState(t)

	accessConfig := common.FakeAccessConfig()

	step := stepSnapshotEBSVolumes{
		PollingConfig: new(common.AWSPollingConfig),
		AccessConfig:  accessConfig,
		VolumeMapping: b.config.VolumeMappings,
		Ctx:           b.config.ctx,
	}

	step.Run(context.Background(), state)

	if len(step.snapshotMap) != 0 {
		t.Fatalf("Shouldn't have snapshotted any volumes")
	}
}
<commit_msg>remove unused mock fields<commit_after>package ebsvolume

import (
	"bytes"
	"context"
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"

	//"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	"github.com/hashicorp/packer-plugin-sdk/multistep"
	"github.com/hashicorp/packer-plugin-sdk/packer"
	"github.com/hashicorp/packer/builder/amazon/common"
)

// Define a mock struct to be used in unit tests for common aws steps.
type mockEC2Conn struct {
	ec2iface.EC2API
	Config *aws.Config
}

func (m *mockEC2Conn) CreateSnapshot(input *ec2.CreateSnapshotInput) (*ec2.Snapshot, error) {
	snap := &ec2.Snapshot{
		// This isn't typical amazon format, but injecting the volume id into
		// this field lets us verify that the right volume was snapshotted with
		// a simple string comparison
		SnapshotId: aws.String(fmt.Sprintf("snap-of-%s", *input.VolumeId)),
	}

	return snap, nil
}

func (m *mockEC2Conn) WaitUntilSnapshotCompletedWithContext(aws.Context, *ec2.DescribeSnapshotsInput, ...request.WaiterOption) error {
	return nil
}

func getMockConn(config *common.AccessConfig, target string) (ec2iface.EC2API, error) {
	mockConn := &mockEC2Conn{
		Config: aws.NewConfig(),
	}
	return mockConn, nil
}

// Create statebag for running test
func tState(t *testing.T) multistep.StateBag {
	state := new(multistep.BasicStateBag)
	state.Put("ui", &packer.BasicUi{
		Reader: new(bytes.Buffer),
		Writer: new(bytes.Buffer),
	})
	// state.Put("amis", map[string]string{"us-east-1": "ami-12345"})
	// state.Put("snapshots", map[string][]string{"us-east-1": {"snap-0012345"}})
	conn, _ := getMockConn(&common.AccessConfig{}, "us-east-2")

	state.Put("ec2", conn)
	// Store a fake instance that contains a block device that matches the
	// volumes defined in the config above
	state.Put("instance", &ec2.Instance{
		InstanceId: aws.String("instance-id"),
		BlockDeviceMappings: []*ec2.InstanceBlockDeviceMapping{
			{
				DeviceName: aws.String("/dev/xvda"),
				Ebs: &ec2.EbsInstanceBlockDevice{
					VolumeId: aws.String("vol-1234"),
				},
			},
			{
				DeviceName: aws.String("/dev/xvdb"),
				Ebs: &ec2.EbsInstanceBlockDevice{
					VolumeId: aws.String("vol-5678"),
				},
			},
		},
	})
	return state
}

func TestStepSnapshot_run_simple(t *testing.T) {
	var b Builder
	config := testConfig() //from builder_test

	//Set some snapshot settings
	config["ebs_volumes"] = []map[string]interface{}{
		{
			"device_name":           "/dev/xvdb",
			"volume_size":           "32",
			"delete_on_termination": true,
			"snapshot_volume":       true,
		},
	}

	generatedData, warnings, err := b.Prepare(config)
	if len(warnings) > 0 {
		t.Fatalf("bad: %#v", warnings)
	}
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}
	if len(generatedData) == 0 {
		t.Fatalf("Generated data should not be empty")
	}

	state := tState(t)

	accessConfig := common.FakeAccessConfig()

	step := stepSnapshotEBSVolumes{
		PollingConfig: new(common.AWSPollingConfig),
		AccessConfig:  accessConfig,
		VolumeMapping: b.config.VolumeMappings,
		Ctx:           b.config.ctx,
	}

	step.Run(context.Background(), state)

	if len(step.snapshotMap) != 1 {
		t.Fatalf("Missing Snapshot from step")
	}

	if volmapping := step.snapshotMap["snap-of-vol-5678"]; volmapping == nil {
		t.Fatalf("Didn't snapshot correct volume: Map is %#v", step.snapshotMap)
	}
}

func TestStepSnapshot_run_no_snaps(t *testing.T) {
	var b Builder
	config := testConfig() //from builder_test

	//Set some snapshot settings
	config["ebs_volumes"] = []map[string]interface{}{
		{
			"device_name":           "/dev/xvdb",
			"volume_size":           "32",
			"delete_on_termination": true,
			"snapshot_volume":       false,
		},
	}

	generatedData, warnings, err := b.Prepare(config)
	if len(warnings) > 0 {
		t.Fatalf("bad: %#v", warnings)
	}
	if err != nil {
		t.Fatalf("should not have error: %s", err)
	}
	if len(generatedData) == 0 {
		t.Fatalf("Generated data should not be empty")
	}

	state := tState(t)

	accessConfig := common.FakeAccessConfig()

	step := stepSnapshotEBSVolumes{
		PollingConfig: new(common.AWSPollingConfig),
		AccessConfig:  accessConfig,
		VolumeMapping: b.config.VolumeMappings,
		Ctx:           b.config.ctx,
	}

	step.Run(context.Background(), state)

	if len(step.snapshotMap) != 0 {
		t.Fatalf("Shouldn't have snapshotted any volumes")
	}
}
<|endoftext|>


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

// This package implements a basic LISP interpretor for embedding in a go program for scripting.
// This file contains the type predicate primitive functions.

package golisp

func RegisterTypePredicatePrimitives() {
	MakePrimitiveFunction("list?", "1", IsPairImpl)
	MakePrimitiveFunction("pair?", "1", IsPairImpl)
	MakePrimitiveFunction("alist?", "1", IsAlistImpl)
	MakePrimitiveFunction("nil?", "1", NilPImpl)
	MakePrimitiveFunction("notnil?", "1", NotNilPImpl)
	MakePrimitiveFunction("symbol?", "1", IsSymbolImpl)
	MakePrimitiveFunction("string?", "1", IsStringImpl)
	MakePrimitiveFunction("integer?", "1", IsIntegerImpl)
	MakePrimitiveFunction("number?", "1", IsNumberImpl)
	MakePrimitiveFunction("float?", "1", IsFloatImpl)
	MakePrimitiveFunction("function?", "1", IsFunctionImpl)
	MakePrimitiveFunction("macro?", "1", IsMacroImpl)
}

func IsPairImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(PairP(Car(args))), nil
}

func IsAlistImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(AlistP(Car(args))), nil
}

func NilPImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(NilP(Car(args))), nil
}

func NotNilPImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(NotNilP(Car(args))), nil
}

func IsSymbolImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(SymbolP(Car(args))), nil
}

func IsStringImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(StringP(Car(args))), nil
}

func IsIntegerImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(IntegerP(Car(args))), nil
}

func IsNumberImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(NumberP(Car(args))), nil
}

func IsFloatImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(FloatP(Car(args))), nil
}

func IsFunctionImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(FunctionP(Car(args))), nil
}

func IsMacroImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(MacroP(Car(args))), nil
}
<commit_msg>Add null?, notnull?, frame? and bytearray?<commit_after>// Copyright 2014 SteelSeries ApS.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This package implements a basic LISP interpretor for embedding in a go program for scripting.
// This file contains the type predicate primitive functions.

package golisp

func RegisterTypePredicatePrimitives() {
	MakePrimitiveFunction("list?", "1", IsPairImpl)
	MakePrimitiveFunction("pair?", "1", IsPairImpl)
	MakePrimitiveFunction("alist?", "1", IsAlistImpl)
	MakePrimitiveFunction("nil?", "1", NilPImpl)
	MakePrimitiveFunction("null?", "1", NilPImpl)
	MakePrimitiveFunction("notnil?", "1", NotNilPImpl)
	MakePrimitiveFunction("notnull?", "1", NotNilPImpl)
	MakePrimitiveFunction("symbol?", "1", IsSymbolImpl)
	MakePrimitiveFunction("string?", "1", IsStringImpl)
	MakePrimitiveFunction("integer?", "1", IsIntegerImpl)
	MakePrimitiveFunction("number?", "1", IsNumberImpl)
	MakePrimitiveFunction("float?", "1", IsFloatImpl)
	MakePrimitiveFunction("function?", "1", IsFunctionImpl)
	MakePrimitiveFunction("macro?", "1", IsMacroImpl)
	MakePrimitiveFunction("frame?", "1", IsFrameImpl)
	MakePrimitiveFunction("bytearray?", "1", IsByteArrayImpl)
}

func IsPairImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(PairP(Car(args))), nil
}

func IsAlistImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(AlistP(Car(args))), nil
}

func NilPImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(NilP(Car(args))), nil
}

func NotNilPImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(NotNilP(Car(args))), nil
}

func IsSymbolImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(SymbolP(Car(args))), nil
}

func IsStringImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(StringP(Car(args))), nil
}

func IsIntegerImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(IntegerP(Car(args))), nil
}

func IsNumberImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(NumberP(Car(args))), nil
}

func IsFloatImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(FloatP(Car(args))), nil
}

func IsFunctionImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(FunctionP(Car(args))), nil
}

func IsMacroImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(MacroP(Car(args))), nil
}

func IsFrameImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(FrameP(Car(args))), nil
}

func IsByteArrayImpl(args *Data, env *SymbolTableFrame) (result *Data, err error) {
	return BooleanWithValue(ObjectP(Car(args)) && ObjectType(Car(args)) == "[]byte"), nil
}
<|endoftext|>


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

package commitlog

import (
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"github.com/m3db/m3db/clock"
	"github.com/m3db/m3db/ts"
	"github.com/m3db/m3x/log"
	"github.com/m3db/m3x/metrics"
	"github.com/m3db/m3x/time"
)

var (
	errCommitLogClosed    = errors.New("commit log is closed")
	errCommitLogQueueFull = errors.New("commit log queue is full")

	timeZero = time.Time{}
)

type newCommitLogWriterFn func(flushFn flushFn, opts Options) commitLogWriter

type commitLogFailFn func(err error)

type completionFn func(err error)

type commitLog struct {
	sync.RWMutex
	opts    Options
	nowFn   clock.NowFn
	metrics xmetrics.Scope
	log     xlog.Logger

	newCommitLogWriterFn newCommitLogWriterFn
	commitLogFailFn      commitLogFailFn
	writer               commitLogWriter

	// TODO(r): replace buffered channel with concurrent striped
	// circular buffer to avoid central write lock contention
	writes chan commitLogWrite

	flushMutex      sync.RWMutex
	lastFlushAt     time.Time
	pendingFlushFns []completionFn

	bitset         bitset
	writerExpireAt time.Time
	closed         bool
	closeErr       chan error
}

type valueType int

const (
	writeValueType valueType = iota
	flushValueType
)

type commitLogWrite struct {
	valueType valueType

	series       Series
	datapoint    ts.Datapoint
	unit         xtime.Unit
	annotation   ts.Annotation
	completionFn completionFn
}

// NewCommitLog creates a new commit log
func NewCommitLog(opts Options) CommitLog {
	iops := opts.GetInstrumentOptions()
	iops = iops.MetricsScope(iops.GetMetricsScope().SubScope("commitlog"))
	return &commitLog{
		opts:                 opts,
		nowFn:                opts.GetClockOptions().GetNowFn(),
		metrics:              iops.GetMetricsScope(),
		log:                  iops.GetLogger(),
		newCommitLogWriterFn: newCommitLogWriter,
		writes:               make(chan commitLogWrite, opts.GetBacklogQueueSize()),
		closeErr:             make(chan error),
	}
}

func (l *commitLog) Open() error {
	// Open the buffered commit log writer
	if err := l.openWriter(l.nowFn()); err != nil {
		return err
	}

	// Flush the info header to ensure we can write to disk
	if err := l.writer.Flush(); err != nil {
		return err
	}

	// NB(r): In the future we can introduce a commit log failure policy
	// similar to Cassandra's "stop", for example see:
	// https://github.com/apache/cassandra/blob/6dfc1e7eeba539774784dfd650d3e1de6785c938/conf/cassandra.yaml#L232
	// Right now it is a large amount of coordination to implement something similiar.
	var fails int64
	l.commitLogFailFn = func(err error) {
		currFails := atomic.AddInt64(&fails, 1)
		if currFails < 3 {
			go func() {
				time.Sleep(time.Second)
				atomic.AddInt64(&fails, -1)
			}()
			return
		}
		l.log.Fatalf("fatal commit log error: %v", err)
	}

	// Asynchronously write
	go l.write()

	flushInterval := l.opts.GetFlushInterval()
	if flushInterval > 0 {
		// Continually flush the commit log at given interval if set
		go l.flushEvery(flushInterval)
	}

	return nil
}

func (l *commitLog) flushEvery(interval time.Duration) {
	// Periodically flush the underlying commit log writer to cover
	// the case when writes stall for a considerable time
	var sleepForOverride time.Duration
	for {
		sleepFor := interval
		if sleepForOverride > 0 {
			sleepFor = sleepForOverride
			sleepForOverride = 0
		}

		time.Sleep(sleepFor)

		l.flushMutex.RLock()
		lastFlushAt := l.lastFlushAt
		l.flushMutex.RUnlock()

		sinceLastFlush := l.nowFn().Sub(lastFlushAt)
		if sinceLastFlush < interval {
			// Flushed already recently, sleep until we would next consider flushing
			sleepForOverride = interval - sinceLastFlush
			continue
		}

		// Request a flush
		l.RLock()
		if l.closed {
			l.RUnlock()
			return
		}
		l.writes <- commitLogWrite{valueType: flushValueType}
		l.RUnlock()
	}
}

func (l *commitLog) write() {
	var (
		write commitLogWrite
		err   error
		open  bool
	)
	for {
		write, open = <-l.writes
		if !open {
			break
		}

		// For writes requiring acks add to pending acks
		if write.completionFn != nil {
			l.pendingFlushFns = append(l.pendingFlushFns, write.completionFn)
		}

		if write.valueType == flushValueType {
			l.writer.Flush()
			continue
		}

		now := l.nowFn()
		if !now.Before(l.writerExpireAt) {
			if err = l.openWriter(now); err != nil {
				l.metrics.IncCounter("writes.errors", 1)
				l.metrics.IncCounter("writes.open-errors", 1)
				l.log.Errorf("failed to open commit log: %v", err)
				if l.commitLogFailFn != nil {
					l.commitLogFailFn(err)
				}
				continue
			}
		}

		err = l.writer.Write(write.series, write.datapoint, write.unit, write.annotation)
		if err != nil {
			l.metrics.IncCounter("writes.errors", 1)
			l.log.Errorf("failed to write to commit log: %v", err)
			if l.commitLogFailFn != nil {
				l.commitLogFailFn(err)
			}
			continue
		}

		l.metrics.IncCounter("writes.success", 1)
	}

	l.Lock()
	defer l.Unlock()

	writer := l.writer
	l.writer = nil
	l.closeErr <- writer.Close()
}

func (l *commitLog) onFlush(err error) {
	l.flushMutex.Lock()
	l.lastFlushAt = l.nowFn()
	l.flushMutex.Unlock()

	if err != nil {
		l.metrics.IncCounter("writes.errors", 1)
		l.metrics.IncCounter("writes.flush-errors", 1)
		l.log.Errorf("failed to flush commit log: %v", err)
		if l.commitLogFailFn != nil {
			l.commitLogFailFn(err)
		}
	}

	// onFlush only ever called by "write()" and "openWriter" or
	// before "write()" begins on "Open()" and there are no other
	// accessors of "pendingFlushFns" so it is safe to read and mutate
	// without a lock here
	if len(l.pendingFlushFns) == 0 {
		return
	}

	for i := range l.pendingFlushFns {
		l.pendingFlushFns[i](err)
	}

	l.pendingFlushFns = l.pendingFlushFns[:0]
}

func (l *commitLog) openWriter(now time.Time) error {
	if l.writer != nil {
		if err := l.writer.Close(); err != nil {
			l.metrics.IncCounter("writes.close-errors", 1)
			l.log.Errorf("failed to close commit log: %v", err)
			// If we failed to close then create a new commit log writer
			l.writer = nil
		}
	}
	if l.writer == nil {
		l.writer = l.newCommitLogWriterFn(l.onFlush, l.opts)
	}

	blockSize := l.opts.GetRetentionOptions().GetBlockSize()
	start := now.Truncate(blockSize)
	if err := l.writer.Open(start, blockSize); err != nil {
		return err
	}

	l.writerExpireAt = start.Add(blockSize)
	return nil
}

func (l *commitLog) Write(
	series Series,
	datapoint ts.Datapoint,
	unit xtime.Unit,
	annotation ts.Annotation,
) error {
	l.RLock()

	if l.closed {
		l.RUnlock()
		return errCommitLogClosed
	}

	var (
		wg     sync.WaitGroup
		result error
	)
	wg.Add(1)
	completion := func(err error) {
		result = err
		wg.Done()
	}

	var (
		write = commitLogWrite{
			series:       series,
			datapoint:    datapoint,
			unit:         unit,
			annotation:   annotation,
			completionFn: completion,
		}
		enqueued = false
	)
	select {
	case l.writes <- write:
		enqueued = true
	default:
	}

	l.RUnlock()

	if !enqueued {
		l.RUnlock()
		return errCommitLogQueueFull
	}

	wg.Wait()

	return result
}

func (l *commitLog) WriteBehind(
	series Series,
	datapoint ts.Datapoint,
	unit xtime.Unit,
	annotation ts.Annotation,
) error {
	l.RLock()

	if l.closed {
		l.RUnlock()
		return errCommitLogClosed
	}

	var (
		write = commitLogWrite{
			series:     series,
			datapoint:  datapoint,
			unit:       unit,
			annotation: annotation,
		}
		enqueued = false
	)
	select {
	case l.writes <- write:
		enqueued = true
	default:
	}

	if !enqueued {
		l.RUnlock()
		return errCommitLogQueueFull
	}

	l.RUnlock()
	return nil
}

func (l *commitLog) Iter() (Iterator, error) {
	return NewIterator(l.opts)
}

func (l *commitLog) Close() error {
	l.Lock()
	if l.closed {
		l.Unlock()
		return nil
	}
	l.closed = true
	close(l.writes)
	l.Unlock()
	// Receive the result of closing the writer from asynchronous writer
	return <-l.closeErr
}
<commit_msg>Emit metric for commit log write queue (#62)<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package commitlog

import (
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"github.com/m3db/m3db/clock"
	"github.com/m3db/m3db/ts"
	"github.com/m3db/m3x/log"
	"github.com/m3db/m3x/metrics"
	"github.com/m3db/m3x/time"
)

var (
	errCommitLogClosed    = errors.New("commit log is closed")
	errCommitLogQueueFull = errors.New("commit log queue is full")

	timeZero = time.Time{}
)

type newCommitLogWriterFn func(flushFn flushFn, opts Options) commitLogWriter

type commitLogFailFn func(err error)

type completionFn func(err error)

type commitLog struct {
	sync.RWMutex
	opts    Options
	nowFn   clock.NowFn
	metrics xmetrics.Scope
	log     xlog.Logger

	newCommitLogWriterFn newCommitLogWriterFn
	commitLogFailFn      commitLogFailFn
	writer               commitLogWriter

	// TODO(r): replace buffered channel with concurrent striped
	// circular buffer to avoid central write lock contention
	writes chan commitLogWrite

	flushMutex      sync.RWMutex
	lastFlushAt     time.Time
	pendingFlushFns []completionFn

	bitset         bitset
	writerExpireAt time.Time
	closed         bool
	closeErr       chan error
}

type valueType int

const (
	writeValueType valueType = iota
	flushValueType
)

type commitLogWrite struct {
	valueType valueType

	series       Series
	datapoint    ts.Datapoint
	unit         xtime.Unit
	annotation   ts.Annotation
	completionFn completionFn
}

// NewCommitLog creates a new commit log
func NewCommitLog(opts Options) CommitLog {
	iops := opts.GetInstrumentOptions()
	iops = iops.MetricsScope(iops.GetMetricsScope().SubScope("commitlog"))
	return &commitLog{
		opts:                 opts,
		nowFn:                opts.GetClockOptions().GetNowFn(),
		metrics:              iops.GetMetricsScope(),
		log:                  iops.GetLogger(),
		newCommitLogWriterFn: newCommitLogWriter,
		writes:               make(chan commitLogWrite, opts.GetBacklogQueueSize()),
		closeErr:             make(chan error),
	}
}

func (l *commitLog) Open() error {
	// Open the buffered commit log writer
	if err := l.openWriter(l.nowFn()); err != nil {
		return err
	}

	// Flush the info header to ensure we can write to disk
	if err := l.writer.Flush(); err != nil {
		return err
	}

	// NB(r): In the future we can introduce a commit log failure policy
	// similar to Cassandra's "stop", for example see:
	// https://github.com/apache/cassandra/blob/6dfc1e7eeba539774784dfd650d3e1de6785c938/conf/cassandra.yaml#L232
	// Right now it is a large amount of coordination to implement something similiar.
	var fails int64
	l.commitLogFailFn = func(err error) {
		currFails := atomic.AddInt64(&fails, 1)
		if currFails < 3 {
			go func() {
				time.Sleep(time.Second)
				atomic.AddInt64(&fails, -1)
			}()
			return
		}
		l.log.Fatalf("fatal commit log error: %v", err)
	}

	// Asynchronously write
	go l.write()

	flushInterval := l.opts.GetFlushInterval()
	if flushInterval > 0 {
		// Continually flush the commit log at given interval if set
		go l.flushEvery(flushInterval)
	}

	return nil
}

func (l *commitLog) flushEvery(interval time.Duration) {
	// Periodically flush the underlying commit log writer to cover
	// the case when writes stall for a considerable time
	var sleepForOverride time.Duration
	for {
		l.metrics.IncCounter("writes.queued", int64(len(l.writes)))

		sleepFor := interval
		if sleepForOverride > 0 {
			sleepFor = sleepForOverride
			sleepForOverride = 0
		}

		time.Sleep(sleepFor)

		l.flushMutex.RLock()
		lastFlushAt := l.lastFlushAt
		l.flushMutex.RUnlock()

		sinceLastFlush := l.nowFn().Sub(lastFlushAt)
		if sinceLastFlush < interval {
			// Flushed already recently, sleep until we would next consider flushing
			sleepForOverride = interval - sinceLastFlush
			continue
		}

		// Request a flush
		l.RLock()
		if l.closed {
			l.RUnlock()
			return
		}
		l.writes <- commitLogWrite{valueType: flushValueType}
		l.RUnlock()
	}
}

func (l *commitLog) write() {
	var (
		write commitLogWrite
		err   error
		open  bool
	)
	for {
		write, open = <-l.writes
		if !open {
			break
		}

		// For writes requiring acks add to pending acks
		if write.completionFn != nil {
			l.pendingFlushFns = append(l.pendingFlushFns, write.completionFn)
		}

		if write.valueType == flushValueType {
			l.writer.Flush()
			continue
		}

		now := l.nowFn()
		if !now.Before(l.writerExpireAt) {
			if err = l.openWriter(now); err != nil {
				l.metrics.IncCounter("writes.errors", 1)
				l.metrics.IncCounter("writes.open-errors", 1)
				l.log.Errorf("failed to open commit log: %v", err)
				if l.commitLogFailFn != nil {
					l.commitLogFailFn(err)
				}
				continue
			}
		}

		err = l.writer.Write(write.series, write.datapoint, write.unit, write.annotation)
		if err != nil {
			l.metrics.IncCounter("writes.errors", 1)
			l.log.Errorf("failed to write to commit log: %v", err)
			if l.commitLogFailFn != nil {
				l.commitLogFailFn(err)
			}
			continue
		}

		l.metrics.IncCounter("writes.success", 1)
	}

	l.Lock()
	defer l.Unlock()

	writer := l.writer
	l.writer = nil
	l.closeErr <- writer.Close()
}

func (l *commitLog) onFlush(err error) {
	l.flushMutex.Lock()
	l.lastFlushAt = l.nowFn()
	l.flushMutex.Unlock()

	if err != nil {
		l.metrics.IncCounter("writes.errors", 1)
		l.metrics.IncCounter("writes.flush-errors", 1)
		l.log.Errorf("failed to flush commit log: %v", err)
		if l.commitLogFailFn != nil {
			l.commitLogFailFn(err)
		}
	}

	// onFlush only ever called by "write()" and "openWriter" or
	// before "write()" begins on "Open()" and there are no other
	// accessors of "pendingFlushFns" so it is safe to read and mutate
	// without a lock here
	if len(l.pendingFlushFns) == 0 {
		return
	}

	for i := range l.pendingFlushFns {
		l.pendingFlushFns[i](err)
	}

	l.pendingFlushFns = l.pendingFlushFns[:0]
}

func (l *commitLog) openWriter(now time.Time) error {
	if l.writer != nil {
		if err := l.writer.Close(); err != nil {
			l.metrics.IncCounter("writes.close-errors", 1)
			l.log.Errorf("failed to close commit log: %v", err)
			// If we failed to close then create a new commit log writer
			l.writer = nil
		}
	}
	if l.writer == nil {
		l.writer = l.newCommitLogWriterFn(l.onFlush, l.opts)
	}

	blockSize := l.opts.GetRetentionOptions().GetBlockSize()
	start := now.Truncate(blockSize)
	if err := l.writer.Open(start, blockSize); err != nil {
		return err
	}

	l.writerExpireAt = start.Add(blockSize)
	return nil
}

func (l *commitLog) Write(
	series Series,
	datapoint ts.Datapoint,
	unit xtime.Unit,
	annotation ts.Annotation,
) error {
	l.RLock()

	if l.closed {
		l.RUnlock()
		return errCommitLogClosed
	}

	var (
		wg     sync.WaitGroup
		result error
	)
	wg.Add(1)
	completion := func(err error) {
		result = err
		wg.Done()
	}

	var (
		write = commitLogWrite{
			series:       series,
			datapoint:    datapoint,
			unit:         unit,
			annotation:   annotation,
			completionFn: completion,
		}
		enqueued = false
	)
	select {
	case l.writes <- write:
		enqueued = true
	default:
	}

	l.RUnlock()

	if !enqueued {
		l.RUnlock()
		return errCommitLogQueueFull
	}

	wg.Wait()

	return result
}

func (l *commitLog) WriteBehind(
	series Series,
	datapoint ts.Datapoint,
	unit xtime.Unit,
	annotation ts.Annotation,
) error {
	l.RLock()

	if l.closed {
		l.RUnlock()
		return errCommitLogClosed
	}

	var (
		write = commitLogWrite{
			series:     series,
			datapoint:  datapoint,
			unit:       unit,
			annotation: annotation,
		}
		enqueued = false
	)
	select {
	case l.writes <- write:
		enqueued = true
	default:
	}

	if !enqueued {
		l.RUnlock()
		return errCommitLogQueueFull
	}

	l.RUnlock()
	return nil
}

func (l *commitLog) Iter() (Iterator, error) {
	return NewIterator(l.opts)
}

func (l *commitLog) Close() error {
	l.Lock()
	if l.closed {
		l.Unlock()
		return nil
	}
	l.closed = true
	close(l.writes)
	l.Unlock()
	// Receive the result of closing the writer from asynchronous writer
	return <-l.closeErr
}
<|endoftext|>


<commit_before>package testutil

import (
	"github.com/hashicorp/consul/consul/structs"
	"testing"
	"time"
)

type testFn func() (bool, error)
type errorFn func(error)

func WaitForResult(test testFn, error errorFn) {
	retries := 1000

	for retries > 0 {
		time.Sleep(10 * time.Millisecond)
		retries--

		success, err := test()
		if success {
			return
		}

		if retries == 0 {
			error(err)
		}
	}
}

type rpcFn func(string, interface{}, interface{}) error

func WaitForLeader(t *testing.T, rpc rpcFn, dc string) structs.IndexedNodes {
	var out structs.IndexedNodes
	WaitForResult(func() (bool, error) {
		args := &structs.DCSpecificRequest{
			Datacenter: dc,
		}
		err := rpc("Catalog.ListNodes", args, &out)
		return out.QueryMeta.KnownLeader && out.Index > 0, err
	}, func(err error) {
		t.Fatalf("failed to find leader: %v", err)
	})
	return out
}
<commit_msg>Sleep for longer, but try less often<commit_after>package testutil

import (
	"testing"
	"time"

	"github.com/hashicorp/consul/consul/structs"
)

type testFn func() (bool, error)
type errorFn func(error)

func WaitForResult(test testFn, error errorFn) {
	retries := 100

	for retries > 0 {
		time.Sleep(100 * time.Millisecond)
		retries--

		success, err := test()
		if success {
			return
		}

		if retries == 0 {
			error(err)
		}
	}
}

type rpcFn func(string, interface{}, interface{}) error

func WaitForLeader(t *testing.T, rpc rpcFn, dc string) structs.IndexedNodes {
	var out structs.IndexedNodes
	WaitForResult(func() (bool, error) {
		args := &structs.DCSpecificRequest{
			Datacenter: dc,
		}
		err := rpc("Catalog.ListNodes", args, &out)
		return out.QueryMeta.KnownLeader && out.Index > 0, err
	}, func(err error) {
		t.Fatalf("failed to find leader: %v", err)
	})
	return out
}
<|endoftext|>