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

package lzma

import (
	"io"
)

// bWriter is used to convert a standard io.Writer into an io.ByteWriter.
type bWriter struct {
	io.Writer
	a []byte
}

// newByteWriter transforms an io.Writer into an io.ByteWriter.
func newByteWriter(w io.Writer) io.ByteWriter {
	if b, ok := w.(io.ByteWriter); ok {
		return b
	}
	return &bWriter{w, make([]byte, 1)}
}

// WriteByte writes a single byte into the Writer.
func (b *bWriter) WriteByte(c byte) error {
	b.a[0] = c
	n, err := b.Write(b.a)
	switch {
	case n > 1:
		panic("n > 1 for writing a single byte")
	case n == 1:
		return nil
	case err == nil:
		panic("no error for n == 0")
	}
	return err
}

// bReader is used to convert an io.Reader into an io.ByteReader.
type bReader struct {
	io.Reader
	a []byte
}

// newByteReader transforms an io.Reader into an io.ByteReader.
func newByteReader(r io.Reader) io.ByteReader {
	if b, ok := r.(io.ByteReader); ok {
		return b
	}
	return &bReader{r, make([]byte, 1)}
}

// ReadByte reads a byte from the wrapped io.ByteReader.
func (b bReader) ReadByte() (byte, error) {
	n, err := b.Read(b.a)
	switch {
	case n > 1:
		panic("n < 1 for reading a single byte")
	case n == 1:
		return b.a[0], nil
	}
	return 0, err
}

// rangeEncoder implements range encoding of single bits. The low value can
// overflow therefore we need uint64. The cache value is used to handle
// overflows.
type rangeEncoder struct {
	w         io.ByteWriter
	nrange    uint32
	low       uint64
	cacheSize int64
	cache     byte
}

// newRangeEncoder creates a new range encoder.
func newRangeEncoder(w io.Writer) *rangeEncoder {
	return &rangeEncoder{
		w:         newByteWriter(w),
		nrange:    0xffffffff,
		cacheSize: 1}
}

// DirectEncodeBit encodes the least-significant bit of b with probability 1/2.
func (e *rangeEncoder) DirectEncodeBit(b uint32) error {
	// e.bitCounter++
	e.nrange >>= 1
	e.low += uint64(e.nrange) & (0 - (uint64(b) & 1))
	if err := e.normalize(); err != nil {
		return err
	}

	return nil
}

// EncodeBit encodes the least significant bit of b. The p value will be
// updated by the function depending on the bit encoded.
func (e *rangeEncoder) EncodeBit(b uint32, p *prob) error {
	// e.bitCounter++
	bound := p.bound(e.nrange)
	if b&1 == 0 {
		e.nrange = bound
		p.inc()
	} else {
		e.low += uint64(bound)
		e.nrange -= bound
		p.dec()
	}
	if err := e.normalize(); err != nil {
		return err
	}

	return nil
}

// closeLen returns the number of bytes required for closing.
func (e *rangeEncoder) closeLen() int64 {
	return e.cacheSize + 4
}

// Close writes a complete copy of the low value.
func (e *rangeEncoder) Close() error {
	for i := 0; i < 5; i++ {
		if err := e.shiftLow(); err != nil {
			return err
		}
	}
	return nil
}

// newRangeDecoder initializes a range decoder. It reads five bytes from the
// reader and therefore may return an error.
func newRangeDecoder(r io.Reader) (d *rangeDecoder, err error) {
	d = &rangeDecoder{r: newByteReader(r)}
	err = d.init()
	return
}

// possiblyAtEnd checks whether the decoder may be at the end of the stream.
func (d *rangeDecoder) possiblyAtEnd() bool {
	return d.code == 0
}

// DirectDecodeBit decodes a bit with probability 1/2. The return value b will
// contain the bit at the least-significant position. All other bits will be
// zero.
func (d *rangeDecoder) DirectDecodeBit() (b uint32, err error) {
	// d.bitCounter++
	d.nrange >>= 1
	d.code -= d.nrange
	t := 0 - (d.code >> 31)
	d.code += d.nrange & t

	// d.code will stay less then d.nrange

	if err = d.normalize(); err != nil {
		return 0, err
	}

	b = (t + 1) & 1

	return b, nil
}

// decodeBit decodes a single bit. The bit will be returned at the
// least-significant position. All other bits will be zero. The probability
// value will be updated.
func (d *rangeDecoder) DecodeBit(p *prob) (b uint32, err error) {
	// d.bitCounter++
	bound := p.bound(d.nrange)
	if d.code < bound {
		d.nrange = bound
		p.inc()
		b = 0
	} else {
		d.code -= bound
		d.nrange -= bound
		p.dec()
		b = 1
	}

	// d.code will stay less then d.nrange

	if err = d.normalize(); err != nil {
		return 0, err
	}

	return b, nil
}

// shiftLow shifts the low value for 8 bit. The shifted byte is written into
// the byte writer. The cache value is used to handle overflows.
func (e *rangeEncoder) shiftLow() error {
	if uint32(e.low) < 0xff000000 || (e.low>>32) != 0 {
		tmp := e.cache
		for {
			err := e.w.WriteByte(tmp + byte(e.low>>32))
			if err != nil {
				return err
			}
			tmp = 0xff
			e.cacheSize--
			if e.cacheSize <= 0 {
				if e.cacheSize < 0 {
					return negError{"cacheSize", e.cacheSize}
				}
				break
			}
		}
		e.cache = byte(uint32(e.low) >> 24)
	}
	e.cacheSize++
	e.low = uint64(uint32(e.low) << 8)
	return nil
}

// normalize handles shifts of nrange and low.
func (e *rangeEncoder) normalize() error {
	const top = 1 << 24
	if e.nrange >= top {
		return nil
	}
	e.nrange <<= 8
	return e.shiftLow()
}

// rangeDecoder decodes single bits of the range encoding stream.
type rangeDecoder struct {
	r      io.ByteReader
	nrange uint32
	code   uint32
}

// init initializes the range decoder, by reading from the byte reader.
func (d *rangeDecoder) init() error {
	d.nrange = 0xffffffff
	d.code = 0

	b, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	if b != 0 {
		return lzmaError{"first byte not zero"}
	}

	for i := 0; i < 4; i++ {
		if err = d.updateCode(); err != nil {
			return err
		}
	}

	if d.code >= d.nrange {
		return lzmaError{"newRangeDecoder: d.code >= d.nrange"}
	}

	return nil
}

// updateCode reads a new byte into the code.
func (d *rangeDecoder) updateCode() error {
	b, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	d.code = (d.code << 8) | uint32(b)
	return nil
}

// normalize the top value and update the code value.
func (d *rangeDecoder) normalize() error {
	// assume d.code < d.nrange
	const top = 1 << 24
	if d.nrange < top {
		d.nrange <<= 8
		// d.code < d.nrange will be maintained
		if err := d.updateCode(); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>lzma: implemented a rangecodec using a byte write counter<commit_after>// Copyright 2015 Ulrich Kunitz. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package lzma

import (
	"io"
)

// byteWriteCounter is a ByteWriter that counts the bytes written.
type byteWriteCounter interface {
	io.ByteWriter
	Size() int64
}

// bwCounter provides a byteWriteCounter using a ByteWriter.
type bwCounter struct {
	BW io.ByteWriter
	N  int64
}

// newBWCounter converts a ByteWriter to a byteWriteCounter.
func newBWCounter(bw io.ByteWriter) *bwCounter {
	return &bwCounter{BW: bw}
}

// WriteByte writes a single byte to the bwCounter.
func (bwc *bwCounter) WriteByte(c byte) error {
	err := bwc.BW.WriteByte(c)
	if err == nil {
		bwc.N++
	}
	return err
}

// Size returns the number of bytes written.
func (bwc *bwCounter) Size() int64 { return bwc.N }

// wCounter implements a byteWriteCounter on top of a Writer.
type wCounter struct {
	W io.Writer
	N int64
	a []byte
}

// newWCounter converts a Writer to a wCounter.
func newWCounter(w io.Writer) *wCounter {
	return &wCounter{W: w, a: make([]byte, 1)}
}

// WriteByte writes a single byte into the wCounter.
func (wc *wCounter) WriteByte(c byte) error {
	wc.a[0] = c
	n, err := wc.W.Write(wc.a)
	switch {
	case n > 1:
		panic("n > 1 for writing a single byte")
	case n == 1:
		wc.N++
		return nil
	case err == nil:
		panic("no error for n == 0")
	}
	return err
}

// Size returns the total number of bytes written.
func (wc *wCounter) Size() int64 { return wc.N }

// newByteWriteCounter transforms an io.Writer into an byteWriteCounter
func newByteWriteCounter(w io.Writer) byteWriteCounter {
	if bw, ok := w.(io.ByteWriter); ok {
		return newBWCounter(bw)
	}
	return newWCounter(w)
}

// bReader is used to convert an io.Reader into an io.ByteReader.
type bReader struct {
	io.Reader
	a []byte
}

// newByteReader transforms an io.Reader into an io.ByteReader.
func newByteReader(r io.Reader) io.ByteReader {
	if b, ok := r.(io.ByteReader); ok {
		return b
	}
	return &bReader{r, make([]byte, 1)}
}

// ReadByte reads a byte from the wrapped io.ByteReader.
func (b bReader) ReadByte() (byte, error) {
	n, err := b.Read(b.a)
	switch {
	case n > 1:
		panic("n < 1 for reading a single byte")
	case n == 1:
		return b.a[0], nil
	}
	return 0, err
}

// rangeEncoder implements range encoding of single bits. The low value can
// overflow therefore we need uint64. The cache value is used to handle
// overflows.
type rangeEncoder struct {
	w         byteWriteCounter
	nrange    uint32
	low       uint64
	cacheSize int64
	cache     byte
}

// newRangeEncoder creates a new range encoder.
func newRangeEncoder(w io.Writer) *rangeEncoder {
	return &rangeEncoder{
		w:         newByteWriteCounter(w),
		nrange:    0xffffffff,
		cacheSize: 1}
}

// DirectEncodeBit encodes the least-significant bit of b with probability 1/2.
func (e *rangeEncoder) DirectEncodeBit(b uint32) error {
	// e.bitCounter++
	e.nrange >>= 1
	e.low += uint64(e.nrange) & (0 - (uint64(b) & 1))
	if err := e.normalize(); err != nil {
		return err
	}

	return nil
}

// EncodeBit encodes the least significant bit of b. The p value will be
// updated by the function depending on the bit encoded.
func (e *rangeEncoder) EncodeBit(b uint32, p *prob) error {
	// e.bitCounter++
	bound := p.bound(e.nrange)
	if b&1 == 0 {
		e.nrange = bound
		p.inc()
	} else {
		e.low += uint64(bound)
		e.nrange -= bound
		p.dec()
	}
	if err := e.normalize(); err != nil {
		return err
	}

	return nil
}

// closeLen returns the number of bytes required for closing.
func (e *rangeEncoder) closeLen() int64 {
	return e.cacheSize + 4
}

// Close writes a complete copy of the low value.
func (e *rangeEncoder) Close() error {
	for i := 0; i < 5; i++ {
		if err := e.shiftLow(); err != nil {
			return err
		}
	}
	return nil
}

// newRangeDecoder initializes a range decoder. It reads five bytes from the
// reader and therefore may return an error.
func newRangeDecoder(r io.Reader) (d *rangeDecoder, err error) {
	d = &rangeDecoder{r: newByteReader(r)}
	err = d.init()
	return
}

// possiblyAtEnd checks whether the decoder may be at the end of the stream.
func (d *rangeDecoder) possiblyAtEnd() bool {
	return d.code == 0
}

// DirectDecodeBit decodes a bit with probability 1/2. The return value b will
// contain the bit at the least-significant position. All other bits will be
// zero.
func (d *rangeDecoder) DirectDecodeBit() (b uint32, err error) {
	// d.bitCounter++
	d.nrange >>= 1
	d.code -= d.nrange
	t := 0 - (d.code >> 31)
	d.code += d.nrange & t

	// d.code will stay less then d.nrange

	if err = d.normalize(); err != nil {
		return 0, err
	}

	b = (t + 1) & 1

	return b, nil
}

// decodeBit decodes a single bit. The bit will be returned at the
// least-significant position. All other bits will be zero. The probability
// value will be updated.
func (d *rangeDecoder) DecodeBit(p *prob) (b uint32, err error) {
	// d.bitCounter++
	bound := p.bound(d.nrange)
	if d.code < bound {
		d.nrange = bound
		p.inc()
		b = 0
	} else {
		d.code -= bound
		d.nrange -= bound
		p.dec()
		b = 1
	}

	// d.code will stay less then d.nrange

	if err = d.normalize(); err != nil {
		return 0, err
	}

	return b, nil
}

// shiftLow shifts the low value for 8 bit. The shifted byte is written into
// the byte writer. The cache value is used to handle overflows.
func (e *rangeEncoder) shiftLow() error {
	if uint32(e.low) < 0xff000000 || (e.low>>32) != 0 {
		tmp := e.cache
		for {
			err := e.w.WriteByte(tmp + byte(e.low>>32))
			if err != nil {
				return err
			}
			tmp = 0xff
			e.cacheSize--
			if e.cacheSize <= 0 {
				if e.cacheSize < 0 {
					return negError{"cacheSize", e.cacheSize}
				}
				break
			}
		}
		e.cache = byte(uint32(e.low) >> 24)
	}
	e.cacheSize++
	e.low = uint64(uint32(e.low) << 8)
	return nil
}

// normalize handles shifts of nrange and low.
func (e *rangeEncoder) normalize() error {
	const top = 1 << 24
	if e.nrange >= top {
		return nil
	}
	e.nrange <<= 8
	return e.shiftLow()
}

// rangeDecoder decodes single bits of the range encoding stream.
type rangeDecoder struct {
	r      io.ByteReader
	nrange uint32
	code   uint32
}

// init initializes the range decoder, by reading from the byte reader.
func (d *rangeDecoder) init() error {
	d.nrange = 0xffffffff
	d.code = 0

	b, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	if b != 0 {
		return lzmaError{"first byte not zero"}
	}

	for i := 0; i < 4; i++ {
		if err = d.updateCode(); err != nil {
			return err
		}
	}

	if d.code >= d.nrange {
		return lzmaError{"newRangeDecoder: d.code >= d.nrange"}
	}

	return nil
}

// updateCode reads a new byte into the code.
func (d *rangeDecoder) updateCode() error {
	b, err := d.r.ReadByte()
	if err != nil {
		return err
	}
	d.code = (d.code << 8) | uint32(b)
	return nil
}

// normalize the top value and update the code value.
func (d *rangeDecoder) normalize() error {
	// assume d.code < d.nrange
	const top = 1 << 24
	if d.nrange < top {
		d.nrange <<= 8
		// d.code < d.nrange will be maintained
		if err := d.updateCode(); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package machine

import (
	"io/ioutil"
	"net"
	"strings"
)

const boot_id_path = "/proc/sys/kernel/random/boot_id"

type Addr struct {
	Addr    string `json:"addr"`
	Network string `json:"network"`
}

type Machine struct {
	BootId string
}

func (m *Machine) String() string {
	return m.BootId
}

func (m *Machine) GetAddresses() []Addr {
	var addrs []Addr
	ifs, err := net.Interfaces()

	if err != nil {
		panic(err)
	}

	shouldAppend := func(i net.Interface) bool {
		if (i.Flags & net.FlagLoopback) == net.FlagLoopback {
			return false
		}

		if (i.Flags & net.FlagUp) != net.FlagUp {
			return false
		}

		return true
	}

	for _, k := range ifs {
		if shouldAppend(k) != true {
			continue
		}
		kaddrs, _ := k.Addrs()
		for _, j := range kaddrs {
			if strings.HasPrefix(j.String(), "fe80::") == true {
				continue
			}
			addrs = append(addrs, Addr{j.String(), j.Network()})
		}
	}

	return addrs
}

func New(bootId string) (m *Machine) {
	m = &Machine{}

	if len(bootId) != 0 {
		m.BootId = bootId
	}

	id, err := ioutil.ReadFile(boot_id_path)
	if err != nil {
		panic(err)
	}
	m.BootId = strings.TrimSpace(string(id))

	return m
}
<commit_msg>fix(machine): Allow BootId to be provided to Machine.New<commit_after>package machine

import (
	"io/ioutil"
	"net"
	"strings"
)

const boot_id_path = "/proc/sys/kernel/random/boot_id"

type Addr struct {
	Addr    string `json:"addr"`
	Network string `json:"network"`
}

type Machine struct {
	BootId string
}

func (m *Machine) String() string {
	return m.BootId
}

func (m *Machine) GetAddresses() []Addr {
	var addrs []Addr
	ifs, err := net.Interfaces()

	if err != nil {
		panic(err)
	}

	shouldAppend := func(i net.Interface) bool {
		if (i.Flags & net.FlagLoopback) == net.FlagLoopback {
			return false
		}

		if (i.Flags & net.FlagUp) != net.FlagUp {
			return false
		}

		return true
	}

	for _, k := range ifs {
		if shouldAppend(k) != true {
			continue
		}
		kaddrs, _ := k.Addrs()
		for _, j := range kaddrs {
			if strings.HasPrefix(j.String(), "fe80::") == true {
				continue
			}
			addrs = append(addrs, Addr{j.String(), j.Network()})
		}
	}

	return addrs
}

func New(bootId string) (m *Machine) {
	m = &Machine{}

	if len(bootId) != 0 {
		m.BootId = bootId
	} else {
		id, err := ioutil.ReadFile(boot_id_path)
		if err != nil {
			panic(err)
		}
		m.BootId = strings.TrimSpace(string(id))
	}

	return m
}
<|endoftext|>


<commit_before>package manager

import (
	"net/url"
	"os"
	"path/filepath"

	"github.com/wasanx25/sreq/history"
	"github.com/wasanx25/sreq/fetcher"
	"github.com/wasanx25/sreq/pager"
	"github.com/wasanx25/sreq/parser"

	"github.com/AlecAivazis/survey"
	"github.com/wasanx25/goss"
)

type Manager struct {
	history *history.History
	pager   pager.Pager
}

func New(keyword, sort string) *Manager {
	historyFile := filepath.Join(os.Getenv("HOME"), ".config", "sreq", "sreq-history.toml")

	h := history.New(historyFile)
	// TODO
	p, _ := pager.New(keyword, sort)

	return &Manager{
		history: h,
		pager:   p,
	}
}

func (m *Manager) Run() error {
	var content fetcher.Content

loop:
	for {
		contents, err := fetcher.Fetch(m.pager.GetURL())
		if err != nil {
			return err
		}

		var contentsStr []string

		for _, c := range contents {
			contentsStr = append(contentsStr, c.GetTitle())
		}

		answer, err := m.viewSelector(contentsStr)
		content = m.getContent(answer, contents)

		if content != nil {
			break loop
		}

		m.pager.NextPage()
	}

	u := url.URL{
		Scheme: "https",
		Host:   "qiita.com",
		Path:   filepath.Join("api", "v2", "items", content.GetID()),
	}
	item, err := parser.Parse(u)

	if err != nil {
		return err
	}

	m.history.Write("keyword", item.URL, item.Title)

	err = goss.Run(item.Markdown)

	return err
}

func (m *Manager) viewSelector(contentsStr []string) (string, error) {
	contentsStr = append(contentsStr, "next")

	contQs := []*survey.Question{
		{
			Name: "content",
			Prompt: &survey.Select{
				Message: "Choice a content",
				Options: contentsStr,
			},
			Validate: survey.Required,
		},
	}

	answer := struct {
		Content string
	}{}
	err := survey.Ask(contQs, &answer)

	return answer.Content, err
}

func (m *Manager) getContent(answer string, contents []fetcher.Content) fetcher.Content {
	var content fetcher.Content

	switch answer {
	case "next":
		return nil
	default:
		for _, c := range contents {
			if answer == c.GetTitle() {
				content = c
				break
			}
		}
	}

	return content
}
<commit_msg>Sort import<commit_after>package manager

import (
	"net/url"
	"os"
	"path/filepath"

	"github.com/wasanx25/sreq/fetcher"
	"github.com/wasanx25/sreq/history"
	"github.com/wasanx25/sreq/pager"
	"github.com/wasanx25/sreq/parser"

	"github.com/AlecAivazis/survey"
	"github.com/wasanx25/goss"
)

type Manager struct {
	history *history.History
	pager   pager.Pager
}

func New(keyword, sort string) *Manager {
	historyFile := filepath.Join(os.Getenv("HOME"), ".config", "sreq", "sreq-history.toml")

	h := history.New(historyFile)
	// TODO
	p, _ := pager.New(keyword, sort)

	return &Manager{
		history: h,
		pager:   p,
	}
}

func (m *Manager) Run() error {
	var content fetcher.Content

loop:
	for {
		contents, err := fetcher.Fetch(m.pager.GetURL())
		if err != nil {
			return err
		}

		var contentsStr []string

		for _, c := range contents {
			contentsStr = append(contentsStr, c.GetTitle())
		}

		answer, err := m.viewSelector(contentsStr)
		content = m.getContent(answer, contents)

		if content != nil {
			break loop
		}

		m.pager.NextPage()
	}

	u := url.URL{
		Scheme: "https",
		Host:   "qiita.com",
		Path:   filepath.Join("api", "v2", "items", content.GetID()),
	}
	item, err := parser.Parse(u)

	if err != nil {
		return err
	}

	m.history.Write("keyword", item.URL, item.Title)

	err = goss.Run(item.Markdown)

	return err
}

func (m *Manager) viewSelector(contentsStr []string) (string, error) {
	contentsStr = append(contentsStr, "next")

	contQs := []*survey.Question{
		{
			Name: "content",
			Prompt: &survey.Select{
				Message: "Choice a content",
				Options: contentsStr,
			},
			Validate: survey.Required,
		},
	}

	answer := struct {
		Content string
	}{}
	err := survey.Ask(contQs, &answer)

	return answer.Content, err
}

func (m *Manager) getContent(answer string, contents []fetcher.Content) fetcher.Content {
	var content fetcher.Content

	switch answer {
	case "next":
		return nil
	default:
		for _, c := range contents {
			if answer == c.GetTitle() {
				content = c
				break
			}
		}
	}

	return content
}
<|endoftext|>


<commit_before>// Package manager provides interface between handler and lower level implementation
// such as geoloader.
package manager

// The implementation is currently rather naive.  Eviction is done based only on whether
// there is a pending request, and there are already the max number of datasets loaded.
// A later implementation will use LRU and dead time to make this determination.
//
// Behavior:
//   If a legacy dataset is requests, return the CurrentAnnotator instead.
//   If the requested dataset is loaded, return it.
//   If the requested dataset is loading, return ErrPendingAnnotatorLoad
//   If the dataset is not loaded or pending, check:
//      A: If there are already MaxPending loads in process:
//        Do nothing and reply with ErrPendingAnnotatorLoad (even though this isn't true)
//      B: If there is room to load it?
//       YES: start loading it, and return ErrPendingAnnotatorLoad
//        NO: kick out an existing dataset and return ErrPendingAnnotatorLoad.
//
// Please modify with extreme caution.  The lock MUST be held when ACCESSING any field
// of AnnotatorMap.

// Note that the system may evict up to the number of pending loads, so at any given time,
// there may only be MaxDatasetInMemory = MaxPending actually loaded.

// Also note that anyone holding an annotator will prevent it from being collected by the
// GC, so simply evicting it is not a guarantee that the memory will be reclaimed.

import (
	"errors"
	"log"
	"sync"
	"time"

	"github.com/m-lab/annotation-service/api"
	"github.com/m-lab/annotation-service/geoloader"
	"github.com/m-lab/annotation-service/metrics"
)

var (
	// These are vars instead of consts to facilitate testing.
	MaxDatasetInMemory = 12 // Limit on number of loaded datasets
	MaxPending         = 2  // Limit on number of concurrently loading datasets.

	// ErrNilDataset is returned when CurrentAnnotator is nil.
	ErrNilDataset = errors.New("CurrentAnnotator is nil")

	// ErrPendingAnnotatorLoad is returned when a new annotator is requested, but not yet loaded.
	ErrPendingAnnotatorLoad = errors.New("annotator is loading")

	// ErrAnnotatorLoadFailed is returned when a requested annotator has failed to load.
	ErrAnnotatorLoadFailed = errors.New("unable to load annoator")

	// These are UNEXPECTED errors!!

	// ErrNoAppropriateDataset is returned when directory is empty.
	ErrNoAppropriateDataset = errors.New("No Appropriate Dataset")
	// ErrGoroutineNotOwner indicates multithreaded code problem with reservation.
	ErrGoroutineNotOwner = errors.New("Goroutine not owner")
	// ErrMapEntryAlreadySet indicates multithreaded code problem setting map entry.
	ErrMapEntryAlreadySet = errors.New("Map entry already set")

	// TODO remove these and keep current in the map.
	// A mutex to make sure that we are not reading from the CurrentAnnotator
	// pointer while trying to update it
	currentDataMutex = &sync.RWMutex{}

	// CurrentAnnotator points to a GeoDataset struct containing the absolute
	// latest data for the annotator to search and reply with
	CurrentAnnotator api.Annotator

	// ArchivedLoader points to a AnnotatorMap struct containing the archived
	// Geolite2 and legacy dataset in memory.
	archivedAnnotator = NewAnnotatorMap(geoloader.ArchivedLoader)
)

// AnnotatorMap manages all loading and fetching of Annotators.
// TODO - should we call this AnnotatorCache?
// TODO - should this be a generic cache of interface{}?
//
// Synchronization:
//  All accesses must hold the mutex.  If an element is not found, the
//  goroutine may attempt to take responsibility for loading it by obtaining
//  the write lock, and writing an entry with a nil pointer.
// TODO - still need a strategy for dealing with persistent errors.
type AnnotatorMap struct {
	// Keys are filename of the datasets.
	annotators map[string]api.Annotator
	// Lock to be held when reading or writing the map.
	mutex      sync.RWMutex
	numPending int
	loader     func(string) (api.Annotator, error)
}

// NewAnnotatorMap creates a new map that will use the provided loader for loading new Annotators.
func NewAnnotatorMap(loader func(string) (api.Annotator, error)) *AnnotatorMap {
	return &AnnotatorMap{annotators: make(map[string]api.Annotator), loader: loader}
}

// NOTE: Should only be called by checkAndLoadAnnotator.
// The calling goroutine should "own" the responsibility for
// setting the annotator.
func (am *AnnotatorMap) setAnnotatorIfNil(key string, ann api.Annotator) error {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	old, ok := am.annotators[key]
	if !ok {
		log.Println("This should never happen", ErrGoroutineNotOwner)
		metrics.ErrorTotal.WithLabelValues("WrongOwner").Inc()
		return ErrGoroutineNotOwner
	}
	if old != nil {
		log.Println("This should never happen", ErrMapEntryAlreadySet)
		metrics.ErrorTotal.WithLabelValues("MapEntryAlreadySet").Inc()
		return ErrMapEntryAlreadySet
	}

	am.annotators[key] = ann
	metrics.LoadCount.Inc()
	metrics.PendingLoads.Dec()
	metrics.DatasetCount.Inc()
	am.numPending--
	log.Println("Loaded", key)
	return nil
}

// This creates a reservation for loading a dataset, IFF map entry is empty (not nil or populated)
//   If the dataset is not loaded or pending, check:
//      A: If there are already MaxPending loads in process:
//        Do nothing and reply false
//      B: If there is room to load it?
//       YES: make the reservation (by setting entry to nil) and return true.
//        NO: kick out an existing dataset and return false.
func (am *AnnotatorMap) maybeSetNil(key string) bool {
	am.mutex.Lock()
	defer am.mutex.Unlock()
	_, ok := am.annotators[key]
	if ok {
		// Another goroutine is already responsible for loading.
		return false
	}

	if am.numPending >= MaxPending {
		return false
	}
	// Check the number of datasets in memory. Given the memory
	// limit, some dataset may be removed from memory if needed.
	if len(am.annotators) >= MaxDatasetInMemory {
		for fileKey := range am.annotators {
			if am.annotators[fileKey] != nil {
				log.Println("removing dataset " + fileKey)
				delete(am.annotators, fileKey)
				metrics.EvictionCount.Inc()
				metrics.DatasetCount.Dec()
				break
			}
		}
		return false
	}

	// Place marker so that other requesters know it is loading.
	am.annotators[key] = nil
	metrics.PendingLoads.Inc()
	am.numPending++
	return true
}

// This synchronously attempts to set map entry to nil, and
// if successful, proceeds to asynchronously load the new dataset.
func (am *AnnotatorMap) checkAndLoadAnnotator(key string) {
	reserved := am.maybeSetNil(key)
	if reserved {
		// This goroutine now has exclusive ownership of the
		// map entry, and the responsibility for loading the annotator.
		go func(key string) {
			newAnn, err := am.loader(key)
			if err != nil {
				metrics.ErrorTotal.WithLabelValues(err.Error()).Inc()
				log.Println(err)
				return
			}
			// Set the new annotator value.  Entry should be nil.
			err = am.setAnnotatorIfNil(key, newAnn)
			if err != nil {
				metrics.ErrorTotal.WithLabelValues(err.Error()).Inc()
				log.Println(err)
			}
		}(key)
	}
}

// GetAnnotator gets the named annotator, if already in the map.
// If not already loaded, this will trigger loading, and return ErrPendingAnnotatorLoad
func (am *AnnotatorMap) GetAnnotator(key string) (api.Annotator, error) {
	am.mutex.RLock()
	ann, ok := am.annotators[key]
	am.mutex.RUnlock()

	if !ok {
		am.checkAndLoadAnnotator(key)
		metrics.RejectionCount.WithLabelValues("New Dataset").Inc()
		return nil, ErrPendingAnnotatorLoad
	}

	if ann == nil {
		// Another goroutine is already loading this entry.  Return error.
		metrics.RejectionCount.WithLabelValues("Dataset Pending").Inc()
		return nil, ErrPendingAnnotatorLoad
	}
	return ann, nil
}

// GetAnnotator returns the correct annotator to use for a given timestamp.
// TODO: Update to properly handle legacy datasets.
func GetAnnotator(date time.Time) (api.Annotator, error) {
	// key := strconv.FormatInt(date.Unix(), encodingBase)
	if date.After(geoloader.LatestDatasetDate()) {
		currentDataMutex.RLock()
		ann := CurrentAnnotator
		currentDataMutex.RUnlock()
		return ann, nil
	}

	filename := geoloader.BestAnnotatorName(date)

	if filename == "" {
		metrics.ErrorTotal.WithLabelValues("No Appropriate Dataset").Inc()
		return nil, errors.New("No Appropriate Dataset")
	}

	return archivedAnnotator.GetAnnotator(filename)
}

// InitDataset will update the filename list of archived dataset in memory
// and load the latest Geolite2 dataset in memory.
func InitDataset() {
	geoloader.UpdateArchivedFilenames()

	ann := geoloader.GetLatestData()
	currentDataMutex.Lock()
	CurrentAnnotator = ann
	currentDataMutex.Unlock()
}
<commit_msg>use-unload()<commit_after>// Package manager provides interface between handler and lower level implementation
// such as geoloader.
package manager

// The implementation is currently rather naive.  Eviction is done based only on whether
// there is a pending request, and there are already the max number of datasets loaded.
// A later implementation will use LRU and dead time to make this determination.
//
// Behavior:
//   If a legacy dataset is requests, return the CurrentAnnotator instead.
//   If the requested dataset is loaded, return it.
//   If the requested dataset is loading, return ErrPendingAnnotatorLoad
//   If the dataset is not loaded or pending, check:
//      A: If there are already MaxPending loads in process:
//        Do nothing and reply with ErrPendingAnnotatorLoad (even though this isn't true)
//      B: If there is room to load it?
//       YES: start loading it, and return ErrPendingAnnotatorLoad
//        NO: kick out an existing dataset and return ErrPendingAnnotatorLoad.
//
// Please modify with extreme caution.  The lock MUST be held when ACCESSING any field
// of AnnotatorMap.

// Note that the system may evict up to the number of pending loads, so at any given time,
// there may only be MaxDatasetInMemory = MaxPending actually loaded.

// Also note that anyone holding an annotator will prevent it from being collected by the
// GC, so simply evicting it is not a guarantee that the memory will be reclaimed.

import (
	"errors"
	"log"
	"sync"
	"time"

	"github.com/m-lab/annotation-service/api"
	"github.com/m-lab/annotation-service/geoloader"
	"github.com/m-lab/annotation-service/metrics"
)

var (
	// These are vars instead of consts to facilitate testing.
	MaxDatasetInMemory = 12 // Limit on number of loaded datasets
	MaxPending         = 2  // Limit on number of concurrently loading datasets.

	// ErrNilDataset is returned when CurrentAnnotator is nil.
	ErrNilDataset = errors.New("CurrentAnnotator is nil")

	// ErrPendingAnnotatorLoad is returned when a new annotator is requested, but not yet loaded.
	ErrPendingAnnotatorLoad = errors.New("annotator is loading")

	// ErrAnnotatorLoadFailed is returned when a requested annotator has failed to load.
	ErrAnnotatorLoadFailed = errors.New("unable to load annoator")

	// These are UNEXPECTED errors!!

	// ErrNoAppropriateDataset is returned when directory is empty.
	ErrNoAppropriateDataset = errors.New("No Appropriate Dataset")
	// ErrGoroutineNotOwner indicates multithreaded code problem with reservation.
	ErrGoroutineNotOwner = errors.New("Goroutine not owner")
	// ErrMapEntryAlreadySet indicates multithreaded code problem setting map entry.
	ErrMapEntryAlreadySet = errors.New("Map entry already set")

	// TODO remove these and keep current in the map.
	// A mutex to make sure that we are not reading from the CurrentAnnotator
	// pointer while trying to update it
	currentDataMutex = &sync.RWMutex{}

	// CurrentAnnotator points to a GeoDataset struct containing the absolute
	// latest data for the annotator to search and reply with
	CurrentAnnotator api.Annotator

	// ArchivedLoader points to a AnnotatorMap struct containing the archived
	// Geolite2 and legacy dataset in memory.
	archivedAnnotator = NewAnnotatorMap(geoloader.ArchivedLoader)
)

// AnnotatorMap manages all loading and fetching of Annotators.
// TODO - should we call this AnnotatorCache?
// TODO - should this be a generic cache of interface{}?
//
// Synchronization:
//  All accesses must hold the mutex.  If an element is not found, the
//  goroutine may attempt to take responsibility for loading it by obtaining
//  the write lock, and writing an entry with a nil pointer.
// TODO - still need a strategy for dealing with persistent errors.
type AnnotatorMap struct {
	// Keys are filename of the datasets.
	annotators map[string]api.Annotator
	// Lock to be held when reading or writing the map.
	mutex      sync.RWMutex
	numPending int
	loader     func(string) (api.Annotator, error)
}

// NewAnnotatorMap creates a new map that will use the provided loader for loading new Annotators.
func NewAnnotatorMap(loader func(string) (api.Annotator, error)) *AnnotatorMap {
	return &AnnotatorMap{annotators: make(map[string]api.Annotator), loader: loader}
}

// NOTE: Should only be called by checkAndLoadAnnotator.
// The calling goroutine should "own" the responsibility for
// setting the annotator.
func (am *AnnotatorMap) setAnnotatorIfNil(key string, ann api.Annotator) error {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	old, ok := am.annotators[key]
	if !ok {
		log.Println("This should never happen", ErrGoroutineNotOwner)
		metrics.ErrorTotal.WithLabelValues("WrongOwner").Inc()
		return ErrGoroutineNotOwner
	}
	if old != nil {
		log.Println("This should never happen", ErrMapEntryAlreadySet)
		metrics.ErrorTotal.WithLabelValues("MapEntryAlreadySet").Inc()
		return ErrMapEntryAlreadySet
	}

	am.annotators[key] = ann
	metrics.LoadCount.Inc()
	metrics.PendingLoads.Dec()
	metrics.DatasetCount.Inc()
	am.numPending--
	log.Println("Loaded", key)
	return nil
}

// This creates a reservation for loading a dataset, IFF map entry is empty (not nil or populated)
//   If the dataset is not loaded or pending, check:
//      A: If there are already MaxPending loads in process:
//        Do nothing and reply false
//      B: If there is room to load it?
//       YES: make the reservation (by setting entry to nil) and return true.
//        NO: kick out an existing dataset and return false.
func (am *AnnotatorMap) maybeSetNil(key string) bool {
	am.mutex.Lock()
	defer am.mutex.Unlock()
	_, ok := am.annotators[key]
	if ok {
		// Another goroutine is already responsible for loading.
		return false
	}

	if am.numPending >= MaxPending {
		return false
	}
	// Check the number of datasets in memory. Given the memory
	// limit, some dataset may be removed from memory if needed.
	if len(am.annotators) >= MaxDatasetInMemory {
		for fileKey := range am.annotators {
			ann, ok := am.annotators[fileKey]
			if ok {
				log.Println("removing dataset " + fileKey)
				ann.Unload()
				delete(am.annotators, fileKey)
				metrics.EvictionCount.Inc()
				metrics.DatasetCount.Dec()
				break
			}
		}
		return false
	}

	// Place marker so that other requesters know it is loading.
	am.annotators[key] = nil
	metrics.PendingLoads.Inc()
	am.numPending++
	return true
}

// This synchronously attempts to set map entry to nil, and
// if successful, proceeds to asynchronously load the new dataset.
func (am *AnnotatorMap) checkAndLoadAnnotator(key string) {
	reserved := am.maybeSetNil(key)
	if reserved {
		// This goroutine now has exclusive ownership of the
		// map entry, and the responsibility for loading the annotator.
		go func(key string) {
			newAnn, err := am.loader(key)
			if err != nil {
				metrics.ErrorTotal.WithLabelValues(err.Error()).Inc()
				log.Println(err)
				return
			}
			// Set the new annotator value.  Entry should be nil.
			err = am.setAnnotatorIfNil(key, newAnn)
			if err != nil {
				metrics.ErrorTotal.WithLabelValues(err.Error()).Inc()
				log.Println(err)
			}
		}(key)
	}
}

// GetAnnotator gets the named annotator, if already in the map.
// If not already loaded, this will trigger loading, and return ErrPendingAnnotatorLoad
func (am *AnnotatorMap) GetAnnotator(key string) (api.Annotator, error) {
	am.mutex.RLock()
	ann, ok := am.annotators[key]
	am.mutex.RUnlock()

	if !ok {
		am.checkAndLoadAnnotator(key)
		metrics.RejectionCount.WithLabelValues("New Dataset").Inc()
		return nil, ErrPendingAnnotatorLoad
	}

	if ann == nil {
		// Another goroutine is already loading this entry.  Return error.
		metrics.RejectionCount.WithLabelValues("Dataset Pending").Inc()
		return nil, ErrPendingAnnotatorLoad
	}
	return ann, nil
}

// GetAnnotator returns the correct annotator to use for a given timestamp.
// TODO: Update to properly handle legacy datasets.
func GetAnnotator(date time.Time) (api.Annotator, error) {
	// key := strconv.FormatInt(date.Unix(), encodingBase)
	if date.After(geoloader.LatestDatasetDate()) {
		currentDataMutex.RLock()
		ann := CurrentAnnotator
		currentDataMutex.RUnlock()
		return ann, nil
	}

	filename := geoloader.BestAnnotatorName(date)

	if filename == "" {
		metrics.ErrorTotal.WithLabelValues("No Appropriate Dataset").Inc()
		return nil, errors.New("No Appropriate Dataset")
	}

	return archivedAnnotator.GetAnnotator(filename)
}

// InitDataset will update the filename list of archived dataset in memory
// and load the latest Geolite2 dataset in memory.
func InitDataset() {
	geoloader.UpdateArchivedFilenames()

	ann := geoloader.GetLatestData()
	currentDataMutex.Lock()
	CurrentAnnotator = ann
	currentDataMutex.Unlock()
}
<|endoftext|>


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

package main

import (
	"fmt"
	"strings"

	log "github.com/couchbaselabs/clog"
)

// A janitor maintains PIndexes and Feeds, creating, deleting &
// hooking them up as necessary to try to match to latest plans.

const JANITOR_CLOSE_PINDEX = "janitor_close_pindex"

func (mgr *Manager) JanitorNOOP(msg string) {
	SyncWorkReq(mgr.janitorCh, WORK_NOOP, msg, nil)
}

func (mgr *Manager) JanitorKick(msg string) {
	SyncWorkReq(mgr.janitorCh, WORK_KICK, msg, nil)
}

func (mgr *Manager) JanitorLoop() {
	for m := range mgr.janitorCh {
		var err error
		if m.op == WORK_KICK {
			err = mgr.JanitorOnce(m.msg)
			if err != nil {
				log.Printf("error: JanitorOnce, err: %v", err)
				// Keep looping as perhaps it's a transient issue.
				// TODO: Need a better way to track warnings from errors,
				// and which need a future, rescheduled janitor kick.
			}
		} else if m.op == WORK_NOOP {
			// NOOP.
		} else if m.op == JANITOR_CLOSE_PINDEX {
			mgr.stopPIndex(m.obj.(*PIndex))
		} else {
			err = fmt.Errorf("error: unknown janitor op: %s, m: %#v", m.op, m)
		}
		if m.resCh != nil {
			if err != nil {
				m.resCh <- err
			}
			close(m.resCh)
		}
	}
}

func (mgr *Manager) JanitorOnce(reason string) error {
	log.Printf("janitor awakes, reason: %s", reason)

	if mgr.cfg == nil { // Can occur during testing.
		return fmt.Errorf("janitor skipped due to nil cfg")
	}

	planPIndexes, _, err := CfgGetPlanPIndexes(mgr.cfg)
	if err != nil {
		return fmt.Errorf("janitor skipped on CfgGetPlanPIndexes err: %v", err)
	}
	if planPIndexes == nil {
		// Might happen if janitor wins a race.
		// TODO: Should we reschedule another future janitor kick?
		return fmt.Errorf("janitor skipped on nil planPIndexes")
	}

	currFeeds, currPIndexes := mgr.CurrentMaps()

	addPlanPIndexes, removePIndexes :=
		CalcPIndexesDelta(mgr.uuid, currPIndexes, planPIndexes)

	log.Printf("janitor pindexes to add:")
	for _, ppi := range addPlanPIndexes {
		log.Printf("  %+v", ppi)
	}
	log.Printf("janitor pindexes to remove:")
	for _, pi := range removePIndexes {
		log.Printf("  %+v", pi)
	}

	// First, teardown pindexes that need to be removed.
	for _, removePIndex := range removePIndexes {
		log.Printf("janitor removing pindex: %s", removePIndex.Name)
		err = mgr.stopPIndex(removePIndex)
		if err != nil {
			return fmt.Errorf("error: janitor removing pindex: %s, err: %v",
				removePIndex.Name, err)
		}
	}
	// Then, (re-)create pindexes that we're missing.
	for _, addPlanPIndex := range addPlanPIndexes {
		log.Printf("janitor adding pindex: %s", addPlanPIndex.Name)
		err = mgr.startPIndex(addPlanPIndex)
		if err != nil {
			return fmt.Errorf("error: janitor adding pindex: %s, err: %v",
				addPlanPIndex.Name, err)
		}
	}

	currFeeds, currPIndexes = mgr.CurrentMaps()

	addFeeds, removeFeeds :=
		CalcFeedsDelta(mgr.uuid, planPIndexes, currFeeds, currPIndexes)

	log.Printf("janitor feeds to add:")
	for _, targetPIndexes := range addFeeds {
		if len(targetPIndexes) > 0 {
			log.Printf("  %s", FeedName(targetPIndexes[0]))
		}
	}
	log.Printf("janitor feeds to remove:")
	for _, removeFeed := range removeFeeds {
		log.Printf("  %s", removeFeed.Name())
	}

	// First, teardown feeds that need to be removed.
	for _, removeFeed := range removeFeeds {
		err = mgr.stopFeed(removeFeed)
		if err != nil {
			return fmt.Errorf("error: janitor removing feed name: %s, err: %v",
				removeFeed.Name(), err)
		}
	}
	// Then, (re-)create feeds that we're missing.
	for _, addFeedTargetPIndexes := range addFeeds {
		err = mgr.startFeed(addFeedTargetPIndexes)
		if err != nil {
			return fmt.Errorf("error: janitor adding feed, err: %v", err)
		}
	}

	return nil
}

// --------------------------------------------------------

// Functionally determine the delta of which pindexes need creation
// and which should be shut down on our local node (mgrUUID).
func CalcPIndexesDelta(mgrUUID string,
	currPIndexes map[string]*PIndex,
	wantedPlanPIndexes *PlanPIndexes) (
	addPlanPIndexes []*PlanPIndex,
	removePIndexes []*PIndex) {
	// Allocate our return arrays.
	addPlanPIndexes = make([]*PlanPIndex, 0)
	removePIndexes = make([]*PIndex, 0)

	// Just for fast transient lookups.
	mapWantedPlanPIndex := make(map[string]*PlanPIndex)
	mapRemovePIndex := make(map[string]*PIndex)

	// For each wanted plan pindex, if a pindex does not exist or is
	// different, then schedule to add.
	for _, wantedPlanPIndex := range wantedPlanPIndexes.PlanPIndexes {
	nodeUUIDs:
		for nodeUUID, nodeState := range wantedPlanPIndex.NodeUUIDs {
			if nodeUUID != mgrUUID || nodeState == "" {
				continue nodeUUIDs
			}

			mapWantedPlanPIndex[wantedPlanPIndex.Name] = wantedPlanPIndex

			currPIndex, exists := currPIndexes[wantedPlanPIndex.Name]
			if !exists {
				addPlanPIndexes = append(addPlanPIndexes, wantedPlanPIndex)
			} else if PIndexMatchesPlan(currPIndex, wantedPlanPIndex) == false {
				addPlanPIndexes = append(addPlanPIndexes, wantedPlanPIndex)
				removePIndexes = append(removePIndexes, currPIndex)
				mapRemovePIndex[currPIndex.Name] = currPIndex
			}

			break nodeUUIDs
		}
	}

	// For each existing pindex, if not part of wanted plan pindex,
	// then schedule for removal.
	for _, currPIndex := range currPIndexes {
		if _, exists := mapWantedPlanPIndex[currPIndex.Name]; !exists {
			if _, exists = mapRemovePIndex[currPIndex.Name]; !exists {
				removePIndexes = append(removePIndexes, currPIndex)
				mapRemovePIndex[currPIndex.Name] = currPIndex
			}
		}
	}

	return addPlanPIndexes, removePIndexes
}

// --------------------------------------------------------

// Functionally determine the delta of which feeds need creation and
// which should be shut down.
func CalcFeedsDelta(nodeUUID string, planPIndexes *PlanPIndexes,
	currFeeds map[string]Feed, pindexes map[string]*PIndex) (
	addFeeds [][]*PIndex, removeFeeds []Feed) {
	// Allocate result holders.
	addFeeds = make([][]*PIndex, 0)
	removeFeeds = make([]Feed, 0)

	// Group the writable pindexes by their feed names.  Non-writable
	// pindexes (such as paused) will have their feeds removed.
	groupedPIndexes := make(map[string][]*PIndex)
	for _, pindex := range pindexes {
		planPIndex, exists := planPIndexes.PlanPIndexes[pindex.Name]
		if exists && planPIndex != nil &&
			strings.Contains(planPIndex.NodeUUIDs[nodeUUID], PLAN_PINDEX_NODE_WRITE) {
			feedName := FeedName(pindex)
			arr, exists := groupedPIndexes[feedName]
			if !exists {
				arr = make([]*PIndex, 0)
			}
			groupedPIndexes[feedName] = append(arr, pindex)
		}
	}

	for feedName, feedPIndexes := range groupedPIndexes {
		if _, exists := currFeeds[feedName]; !exists {
			addFeeds = append(addFeeds, feedPIndexes)
		}
	}

	for currFeedName, currFeed := range currFeeds {
		if _, exists := groupedPIndexes[currFeedName]; !exists {
			removeFeeds = append(removeFeeds, currFeed)
		}
	}

	return addFeeds, removeFeeds
}

func FeedName(pindex *PIndex) string {
	// NOTE: We're depending on the IndexName/IndexUUID to
	// functionally "cover" the SourceType/SourceName/SourceUUID, so
	// we don't need to encode the source parts into the feed name.
	return pindex.IndexName + "_" + pindex.IndexUUID
}

// --------------------------------------------------------

func (mgr *Manager) startPIndex(planPIndex *PlanPIndex) error {
	pindex, err := NewPIndex(mgr, planPIndex.Name, NewUUID(),
		planPIndex.IndexType,
		planPIndex.IndexName,
		planPIndex.IndexUUID,
		planPIndex.IndexSchema,
		planPIndex.SourceType,
		planPIndex.SourceName,
		planPIndex.SourceUUID,
		planPIndex.SourcePartitions,
		mgr.PIndexPath(planPIndex.Name))
	if err != nil {
		return fmt.Errorf("error: NewPIndex, name: %s, err: %v",
			planPIndex.Name, err)
	}

	if err = mgr.registerPIndex(pindex); err != nil {
		close(pindex.Stream)
		return err
	}

	return nil
}

func (mgr *Manager) stopPIndex(pindex *PIndex) error {
	feeds, _ := mgr.CurrentMaps()
	for _, feed := range feeds {
		for _, stream := range feed.Streams() {
			if stream == pindex.Stream {
				if err := mgr.stopFeed(feed); err != nil {
					panic(fmt.Sprintf("error: could not stop feed, err: %v", err))
				}
			}
		}
	}

	pindexUnreg := mgr.unregisterPIndex(pindex.Name)
	if pindexUnreg != nil && pindexUnreg != pindex {
		panic("unregistered pindex isn't the one we're stopping")
	}

	close(pindex.Stream)

	return nil
}

// --------------------------------------------------------

func (mgr *Manager) startFeed(pindexes []*PIndex) error {
	if len(pindexes) <= 0 {
		return nil
	}

	pindexFirst := pindexes[0]
	feedName := FeedName(pindexFirst)

	streams := make(map[string]Stream)
	for _, pindex := range pindexes {
		if f := FeedName(pindex); f != feedName {
			panic(fmt.Sprintf("error: unexpected feedName: %s != %s", f, feedName))
		}
		if _, exists := streams[pindex.SourcePartitions]; exists {
			panic(fmt.Sprintf("error: SourcePartitions collision: %s",
				pindex.SourcePartitions))
		}

		streams[pindex.SourcePartitions] = pindex.Stream
	}

	return mgr.startFeedByType(feedName,
		pindexFirst.IndexName, pindexFirst.IndexUUID,
		pindexFirst.SourceType, pindexFirst.SourceName, pindexFirst.SourceUUID,
		streams)
}

func (mgr *Manager) stopFeed(feed Feed) error {
	feedUnreg := mgr.unregisterFeed(feed.Name())
	if feedUnreg != nil && feedUnreg != feed {
		panic("error: unregistered feed isn't the one we're closing")
	}

	// NOTE: We're depending on feed to synchronously close, so we
	// know it'll no longer be sending to any of its streams anymore.
	//
	return feed.Close()
}

// --------------------------------------------------------

// TODO: need way to track dead cows (non-beef)
// TODO: need a way to collect these errors so REST api
// can show them to user ("hey, perhaps you deleted a bucket
// and should delete these related full-text indexes?
// or the couchbase cluster is just down.");
// perhaps as specialized clog writer?

func (mgr *Manager) startFeedByType(feedName, indexName, indexUUID,
	sourceType, sourceName, sourceUUID string,
	streams map[string]Stream) error {
	if sourceType == "couchbase" {
		// TODO: Should default to DCP feed or projector feed one day.
		return mgr.startTAPFeed(feedName, indexName, indexUUID,
			sourceName, sourceUUID, streams)
	}

	if sourceType == "simple" {
		return mgr.startSimpleFeed(feedName, streams)
	}

	if sourceType == "nil" {
		return mgr.registerFeed(NewNILFeed(feedName, streams))
	}

	return fmt.Errorf("error: startFeed() got unknown source type: %s", sourceType)
}

func (mgr *Manager) startTAPFeed(feedName, indexName, indexUUID,
	bucketName, bucketUUID string, streams map[string]Stream) error {
	feed, err := NewTAPFeed(feedName, mgr.server, "default",
		bucketName, bucketUUID, EmptyPartitionFunc, streams)
	if err != nil {
		return fmt.Errorf("error: could not prepare TAP stream to server: %s,"+
			" bucketName: %s, indexName: %s, err: %v",
			mgr.server, bucketName, indexName, err)
	}
	if err = feed.Start(); err != nil {
		return fmt.Errorf("error: could not start tap feed, server: %s, err: %v",
			mgr.server, err)
	}
	if err = mgr.registerFeed(feed); err != nil {
		feed.Close()
		return err
	}
	return nil
}

func (mgr *Manager) startSimpleFeed(feedName string,
	streams map[string]Stream) error {
	feed, err := NewSimpleFeed(feedName, make(Stream), EmptyPartitionFunc, streams)
	if err != nil {
		return err
	}
	if err = feed.Start(); err != nil {
		return fmt.Errorf("error: could not start simple feed, server: %s, err: %v",
			mgr.server, err)
	}
	if err = mgr.registerFeed(feed); err != nil {
		feed.Close()
		return err
	}

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

package main

import (
	"fmt"
	"strings"

	log "github.com/couchbaselabs/clog"
)

// A janitor maintains PIndexes and Feeds, creating, deleting &
// hooking them up as necessary to try to match to latest plans.

const JANITOR_CLOSE_PINDEX = "janitor_close_pindex"

func (mgr *Manager) JanitorNOOP(msg string) {
	SyncWorkReq(mgr.janitorCh, WORK_NOOP, msg, nil)
}

func (mgr *Manager) JanitorKick(msg string) {
	SyncWorkReq(mgr.janitorCh, WORK_KICK, msg, nil)
}

func (mgr *Manager) JanitorLoop() {
	for m := range mgr.janitorCh {
		var err error
		if m.op == WORK_KICK {
			err = mgr.JanitorOnce(m.msg)
			if err != nil {
				// Keep looping as perhaps it's a transient issue.
				// TODO: perhaps need a rescheduled janitor kick.
				log.Printf("error: JanitorOnce, err: %v", err)
			}
		} else if m.op == WORK_NOOP {
			// NOOP.
		} else if m.op == JANITOR_CLOSE_PINDEX {
			mgr.stopPIndex(m.obj.(*PIndex))
		} else {
			err = fmt.Errorf("error: unknown janitor op: %s, m: %#v", m.op, m)
		}
		if m.resCh != nil {
			if err != nil {
				m.resCh <- err
			}
			close(m.resCh)
		}
	}
}

func (mgr *Manager) JanitorOnce(reason string) error {
	log.Printf("janitor awakes, reason: %s", reason)

	if mgr.cfg == nil { // Can occur during testing.
		return fmt.Errorf("janitor skipped due to nil cfg")
	}

	planPIndexes, _, err := CfgGetPlanPIndexes(mgr.cfg)
	if err != nil {
		return fmt.Errorf("janitor skipped on CfgGetPlanPIndexes err: %v", err)
	}
	if planPIndexes == nil {
		// Might happen if janitor wins an initialization race.
		return fmt.Errorf("janitor skipped on nil planPIndexes")
	}

	currFeeds, currPIndexes := mgr.CurrentMaps()

	addPlanPIndexes, removePIndexes :=
		CalcPIndexesDelta(mgr.uuid, currPIndexes, planPIndexes)

	log.Printf("janitor pindexes to add:")
	for _, ppi := range addPlanPIndexes {
		log.Printf("  %+v", ppi)
	}
	log.Printf("janitor pindexes to remove:")
	for _, pi := range removePIndexes {
		log.Printf("  %+v", pi)
	}

	// First, teardown pindexes that need to be removed.
	for _, removePIndex := range removePIndexes {
		log.Printf("janitor removing pindex: %s", removePIndex.Name)
		err = mgr.stopPIndex(removePIndex)
		if err != nil {
			return fmt.Errorf("error: janitor removing pindex: %s, err: %v",
				removePIndex.Name, err)
		}
	}
	// Then, (re-)create pindexes that we're missing.
	for _, addPlanPIndex := range addPlanPIndexes {
		log.Printf("janitor adding pindex: %s", addPlanPIndex.Name)
		err = mgr.startPIndex(addPlanPIndex)
		if err != nil {
			return fmt.Errorf("error: janitor adding pindex: %s, err: %v",
				addPlanPIndex.Name, err)
		}
	}

	currFeeds, currPIndexes = mgr.CurrentMaps()

	addFeeds, removeFeeds :=
		CalcFeedsDelta(mgr.uuid, planPIndexes, currFeeds, currPIndexes)

	log.Printf("janitor feeds to add:")
	for _, targetPIndexes := range addFeeds {
		if len(targetPIndexes) > 0 {
			log.Printf("  %s", FeedName(targetPIndexes[0]))
		}
	}
	log.Printf("janitor feeds to remove:")
	for _, removeFeed := range removeFeeds {
		log.Printf("  %s", removeFeed.Name())
	}

	// First, teardown feeds that need to be removed.
	for _, removeFeed := range removeFeeds {
		err = mgr.stopFeed(removeFeed)
		if err != nil {
			return fmt.Errorf("error: janitor removing feed name: %s, err: %v",
				removeFeed.Name(), err)
		}
	}
	// Then, (re-)create feeds that we're missing.
	for _, addFeedTargetPIndexes := range addFeeds {
		err = mgr.startFeed(addFeedTargetPIndexes)
		if err != nil {
			return fmt.Errorf("error: janitor adding feed, err: %v", err)
		}
	}

	return nil
}

// --------------------------------------------------------

// Functionally determine the delta of which pindexes need creation
// and which should be shut down on our local node (mgrUUID).
func CalcPIndexesDelta(mgrUUID string,
	currPIndexes map[string]*PIndex,
	wantedPlanPIndexes *PlanPIndexes) (
	addPlanPIndexes []*PlanPIndex,
	removePIndexes []*PIndex) {
	// Allocate our return arrays.
	addPlanPIndexes = make([]*PlanPIndex, 0)
	removePIndexes = make([]*PIndex, 0)

	// Just for fast transient lookups.
	mapWantedPlanPIndex := make(map[string]*PlanPIndex)
	mapRemovePIndex := make(map[string]*PIndex)

	// For each wanted plan pindex, if a pindex does not exist or is
	// different, then schedule to add.
	for _, wantedPlanPIndex := range wantedPlanPIndexes.PlanPIndexes {
	nodeUUIDs:
		for nodeUUID, nodeState := range wantedPlanPIndex.NodeUUIDs {
			if nodeUUID != mgrUUID || nodeState == "" {
				continue nodeUUIDs
			}

			mapWantedPlanPIndex[wantedPlanPIndex.Name] = wantedPlanPIndex

			currPIndex, exists := currPIndexes[wantedPlanPIndex.Name]
			if !exists {
				addPlanPIndexes = append(addPlanPIndexes, wantedPlanPIndex)
			} else if PIndexMatchesPlan(currPIndex, wantedPlanPIndex) == false {
				addPlanPIndexes = append(addPlanPIndexes, wantedPlanPIndex)
				removePIndexes = append(removePIndexes, currPIndex)
				mapRemovePIndex[currPIndex.Name] = currPIndex
			}

			break nodeUUIDs
		}
	}

	// For each existing pindex, if not part of wanted plan pindex,
	// then schedule for removal.
	for _, currPIndex := range currPIndexes {
		if _, exists := mapWantedPlanPIndex[currPIndex.Name]; !exists {
			if _, exists = mapRemovePIndex[currPIndex.Name]; !exists {
				removePIndexes = append(removePIndexes, currPIndex)
				mapRemovePIndex[currPIndex.Name] = currPIndex
			}
		}
	}

	return addPlanPIndexes, removePIndexes
}

// --------------------------------------------------------

// Functionally determine the delta of which feeds need creation and
// which should be shut down.
func CalcFeedsDelta(nodeUUID string, planPIndexes *PlanPIndexes,
	currFeeds map[string]Feed, pindexes map[string]*PIndex) (
	addFeeds [][]*PIndex, removeFeeds []Feed) {
	// Allocate result holders.
	addFeeds = make([][]*PIndex, 0)
	removeFeeds = make([]Feed, 0)

	// Group the writable pindexes by their feed names.  Non-writable
	// pindexes (such as paused) will have their feeds removed.
	groupedPIndexes := make(map[string][]*PIndex)
	for _, pindex := range pindexes {
		planPIndex, exists := planPIndexes.PlanPIndexes[pindex.Name]
		if exists && planPIndex != nil &&
			strings.Contains(planPIndex.NodeUUIDs[nodeUUID], PLAN_PINDEX_NODE_WRITE) {
			feedName := FeedName(pindex)
			arr, exists := groupedPIndexes[feedName]
			if !exists {
				arr = make([]*PIndex, 0)
			}
			groupedPIndexes[feedName] = append(arr, pindex)
		}
	}

	for feedName, feedPIndexes := range groupedPIndexes {
		if _, exists := currFeeds[feedName]; !exists {
			addFeeds = append(addFeeds, feedPIndexes)
		}
	}

	for currFeedName, currFeed := range currFeeds {
		if _, exists := groupedPIndexes[currFeedName]; !exists {
			removeFeeds = append(removeFeeds, currFeed)
		}
	}

	return addFeeds, removeFeeds
}

func FeedName(pindex *PIndex) string {
	// NOTE: We're depending on the IndexName/IndexUUID to
	// functionally "cover" the SourceType/SourceName/SourceUUID, so
	// we don't need to encode the source parts into the feed name.
	return pindex.IndexName + "_" + pindex.IndexUUID
}

// --------------------------------------------------------

func (mgr *Manager) startPIndex(planPIndex *PlanPIndex) error {
	pindex, err := NewPIndex(mgr, planPIndex.Name, NewUUID(),
		planPIndex.IndexType,
		planPIndex.IndexName,
		planPIndex.IndexUUID,
		planPIndex.IndexSchema,
		planPIndex.SourceType,
		planPIndex.SourceName,
		planPIndex.SourceUUID,
		planPIndex.SourcePartitions,
		mgr.PIndexPath(planPIndex.Name))
	if err != nil {
		return fmt.Errorf("error: NewPIndex, name: %s, err: %v",
			planPIndex.Name, err)
	}

	if err = mgr.registerPIndex(pindex); err != nil {
		close(pindex.Stream)
		return err
	}

	return nil
}

func (mgr *Manager) stopPIndex(pindex *PIndex) error {
	feeds, _ := mgr.CurrentMaps()
	for _, feed := range feeds {
		for _, stream := range feed.Streams() {
			if stream == pindex.Stream {
				if err := mgr.stopFeed(feed); err != nil {
					panic(fmt.Sprintf("error: could not stop feed, err: %v", err))
				}
			}
		}
	}

	pindexUnreg := mgr.unregisterPIndex(pindex.Name)
	if pindexUnreg != nil && pindexUnreg != pindex {
		panic("unregistered pindex isn't the one we're stopping")
	}

	close(pindex.Stream)

	return nil
}

// --------------------------------------------------------

func (mgr *Manager) startFeed(pindexes []*PIndex) error {
	if len(pindexes) <= 0 {
		return nil
	}

	pindexFirst := pindexes[0]
	feedName := FeedName(pindexFirst)

	streams := make(map[string]Stream)
	for _, pindex := range pindexes {
		if f := FeedName(pindex); f != feedName {
			panic(fmt.Sprintf("error: unexpected feedName: %s != %s", f, feedName))
		}
		if _, exists := streams[pindex.SourcePartitions]; exists {
			panic(fmt.Sprintf("error: SourcePartitions collision: %s",
				pindex.SourcePartitions))
		}

		streams[pindex.SourcePartitions] = pindex.Stream
	}

	return mgr.startFeedByType(feedName,
		pindexFirst.IndexName, pindexFirst.IndexUUID,
		pindexFirst.SourceType, pindexFirst.SourceName, pindexFirst.SourceUUID,
		streams)
}

func (mgr *Manager) stopFeed(feed Feed) error {
	feedUnreg := mgr.unregisterFeed(feed.Name())
	if feedUnreg != nil && feedUnreg != feed {
		panic("error: unregistered feed isn't the one we're closing")
	}

	// NOTE: We're depending on feed to synchronously close, so we
	// know it'll no longer be sending to any of its streams anymore.
	//
	return feed.Close()
}

// --------------------------------------------------------

// TODO: need way to track dead cows (non-beef)
// TODO: need a way to collect these errors so REST api
// can show them to user ("hey, perhaps you deleted a bucket
// and should delete these related full-text indexes?
// or the couchbase cluster is just down.");
// perhaps as specialized clog writer?

func (mgr *Manager) startFeedByType(feedName, indexName, indexUUID,
	sourceType, sourceName, sourceUUID string,
	streams map[string]Stream) error {
	if sourceType == "couchbase" {
		// TODO: Should default to DCP feed or projector feed one day.
		return mgr.startTAPFeed(feedName, indexName, indexUUID,
			sourceName, sourceUUID, streams)
	}

	if sourceType == "simple" {
		return mgr.startSimpleFeed(feedName, streams)
	}

	if sourceType == "nil" {
		return mgr.registerFeed(NewNILFeed(feedName, streams))
	}

	return fmt.Errorf("error: startFeed() got unknown source type: %s", sourceType)
}

func (mgr *Manager) startTAPFeed(feedName, indexName, indexUUID,
	bucketName, bucketUUID string, streams map[string]Stream) error {
	feed, err := NewTAPFeed(feedName, mgr.server, "default",
		bucketName, bucketUUID, EmptyPartitionFunc, streams)
	if err != nil {
		return fmt.Errorf("error: could not prepare TAP stream to server: %s,"+
			" bucketName: %s, indexName: %s, err: %v",
			mgr.server, bucketName, indexName, err)
	}
	if err = feed.Start(); err != nil {
		return fmt.Errorf("error: could not start tap feed, server: %s, err: %v",
			mgr.server, err)
	}
	if err = mgr.registerFeed(feed); err != nil {
		feed.Close()
		return err
	}
	return nil
}

func (mgr *Manager) startSimpleFeed(feedName string,
	streams map[string]Stream) error {
	feed, err := NewSimpleFeed(feedName, make(Stream), EmptyPartitionFunc, streams)
	if err != nil {
		return err
	}
	if err = feed.Start(); err != nil {
		return fmt.Errorf("error: could not start simple feed, server: %s, err: %v",
			mgr.server, err)
	}
	if err = mgr.registerFeed(feed); err != nil {
		feed.Close()
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Code generated by ""fitask" -type=RouteTableAssociation"; DO NOT EDIT

package awstasks

import (
	"encoding/json"

	"k8s.io/kops/upup/pkg/fi"
)

// RouteTableAssociation

// JSON marshalling boilerplate
type realRouteTableAssociation RouteTableAssociation

func (o *RouteTableAssociation) UnmarshalJSON(data []byte) error {
	var jsonName string
	if err := json.Unmarshal(data, &jsonName); err == nil {
		o.Name = &jsonName
		return nil
	}

	var r realRouteTableAssociation
	if err := json.Unmarshal(data, &r); err != nil {
		return err
	}
	*o = RouteTableAssociation(r)
	return nil
}

var _ fi.HasName = &RouteTableAssociation{}

func (e *RouteTableAssociation) GetName() *string {
	return e.Name
}

func (e *RouteTableAssociation) SetName(name string) {
	e.Name = &name
}

func (e *RouteTableAssociation) String() string {
	return fi.TaskAsString(e)
}
<commit_msg>fitask header fix<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Code generated by ""fitask" -type=RouteTableAssociation"; DO NOT EDIT

package awstasks

import (
	"encoding/json"

	"k8s.io/kops/upup/pkg/fi"
)

// RouteTableAssociation

// JSON marshalling boilerplate
type realRouteTableAssociation RouteTableAssociation

func (o *RouteTableAssociation) UnmarshalJSON(data []byte) error {
	var jsonName string
	if err := json.Unmarshal(data, &jsonName); err == nil {
		o.Name = &jsonName
		return nil
	}

	var r realRouteTableAssociation
	if err := json.Unmarshal(data, &r); err != nil {
		return err
	}
	*o = RouteTableAssociation(r)
	return nil
}

var _ fi.HasName = &RouteTableAssociation{}

func (e *RouteTableAssociation) GetName() *string {
	return e.Name
}

func (e *RouteTableAssociation) SetName(name string) {
	e.Name = &name
}

func (e *RouteTableAssociation) String() string {
	return fi.TaskAsString(e)
}
<|endoftext|>


<commit_before>package download_step

import (
	"archive/tar"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"
	"time"

	"github.com/cloudfoundry-incubator/executor/cacheddownloader"
	"github.com/cloudfoundry-incubator/executor/steps/emittable_error"
	"github.com/cloudfoundry-incubator/garden/warden"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/pivotal-golang/archiver/compressor"
	"github.com/pivotal-golang/archiver/extractor"
	"github.com/pivotal-golang/lager"
)

type DownloadStep struct {
	container        warden.Container
	model            models.DownloadAction
	cachedDownloader cacheddownloader.CachedDownloader
	extractor        extractor.Extractor
	tempDir          string
	logger           lager.Logger
}

func New(
	container warden.Container,
	model models.DownloadAction,
	cachedDownloader cacheddownloader.CachedDownloader,
	extractor extractor.Extractor,
	tempDir string,
	logger lager.Logger,
) *DownloadStep {
	logger = logger.Session("DownloadAction", lager.Data{
		"from": model.From,
		"to":   model.To,
	})
	return &DownloadStep{
		container:        container,
		model:            model,
		cachedDownloader: cachedDownloader,
		extractor:        extractor,
		tempDir:          tempDir,
		logger:           logger,
	}
}

func (step *DownloadStep) Perform() error {
	step.logger.Info("starting")

	//Stream this to the extractor + container when we have streaming support!
	downloadedPath, err := step.download()
	if err != nil {
		return err
	}

	defer os.Remove(downloadedPath)

	if step.model.Extract {
		return step.extract(downloadedPath)
	}

	return step.streamToContainer(downloadedPath)
}

func (step *DownloadStep) download() (string, error) {
	url, err := url.ParseRequestURI(step.model.From)
	if err != nil {
		step.logger.Error("parse-request-uri-error", err)
		return "", err
	}

	downloadedFile, err := step.cachedDownloader.Fetch(url, step.model.CacheKey)
	if err != nil {
		step.logger.Error("cached-downloader-fetch-error", err)
		return "", err
	}
	defer downloadedFile.Close()

	tempFile, err := ioutil.TempFile(step.tempDir, "downloaded")
	if err != nil {
		step.logger.Error("tempfile-create-error", err)
		return "", err
	}

	_, err = io.Copy(tempFile, downloadedFile)
	tempFile.Close()
	if err != nil {
		os.Remove(tempFile.Name())
		step.logger.Error("tempfile-copy-failed", err)
		return "", err
	}

	step.logger.Info("download-successful", lager.Data{
		"tempFile": tempFile.Name(),
	})
	return tempFile.Name(), nil
}

func (step *DownloadStep) extract(downloadedPath string) error {
	extractionDir, err := ioutil.TempDir(step.tempDir, "extracted")
	if err != nil {
		step.logger.Error("extract-failed-to-create-tempdir", err, lager.Data{
			"tempDir": step.tempDir,
		})
		return err
	}
	defer os.RemoveAll(extractionDir)

	err = step.extractor.Extract(downloadedPath, extractionDir)
	if err != nil {
		step.logger.Error("failed-to-extract", err, lager.Data{
			"extractionDir":  extractionDir,
			"downloadedPath": downloadedPath,
		})
		return emittable_error.New(err, "Extraction failed")
	}

	err = step.copyExtractedFiles(extractionDir, step.model.To)
	if err != nil {
		return emittable_error.New(err, "Copying into the container failed")
	}

	step.logger.Info("extract-successful")
	return nil
}

func (step *DownloadStep) streamToContainer(downloadedPath string) error {
	reader, writer := io.Pipe()

	go writeTarTo(filepath.Base(step.model.To), downloadedPath, writer)

	err := step.streamIn(filepath.Dir(step.model.To), reader)
	if err != nil {
		return err
	}

	step.logger.Info("stream-to-container-successful")
	return nil
}

func (step *DownloadStep) Cancel() {}

func (step *DownloadStep) Cleanup() {}

func (step *DownloadStep) copyExtractedFiles(source string, destination string) error {
	reader, writer := io.Pipe()

	go func() {
		err := compressor.WriteTar(source+string(filepath.Separator), writer)
		if err == nil {
			writer.Close()
		} else {
			step.logger.Error("wite-tar-failed", err, lager.Data{
				"source":      source,
				"destination": destination,
			})
			writer.CloseWithError(err)
		}
	}()

	return step.streamIn(destination, reader)
}

func (step *DownloadStep) streamIn(destination string, reader io.Reader) error {
	err := step.container.StreamIn(destination, reader)
	if err != nil {
		step.logger.Error("failed-to-stream-in", err, lager.Data{
			"destination": destination,
		})
		return emittable_error.New(err, "Copying into the container failed")
	}

	return err
}

func writeTarTo(name string, sourcePath string, destination *io.PipeWriter) {
	source, err := os.Open(sourcePath)
	if err != nil {
		destination.CloseWithError(err)
		return
	}
	defer source.Close()

	tarWriter := tar.NewWriter(destination)

	fileInfo, err := source.Stat()
	if err != nil {
		destination.CloseWithError(err)
		return
	}

	err = tarWriter.WriteHeader(&tar.Header{
		Name:       name,
		Size:       fileInfo.Size(),
		Mode:       0644,
		AccessTime: time.Now(),
		ChangeTime: time.Now(),
	})
	if err != nil {
		destination.CloseWithError(err)
		return
	}

	_, err = io.Copy(tarWriter, source)
	if err != nil {
		destination.CloseWithError(err)
		return
	}

	if err := tarWriter.Flush(); err != nil {
		destination.CloseWithError(err)
		return
	}

	destination.Close()
}
<commit_msg>Add cache key to lager session data<commit_after>package download_step

import (
	"archive/tar"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"
	"time"

	"github.com/cloudfoundry-incubator/executor/cacheddownloader"
	"github.com/cloudfoundry-incubator/executor/steps/emittable_error"
	"github.com/cloudfoundry-incubator/garden/warden"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/pivotal-golang/archiver/compressor"
	"github.com/pivotal-golang/archiver/extractor"
	"github.com/pivotal-golang/lager"
)

type DownloadStep struct {
	container        warden.Container
	model            models.DownloadAction
	cachedDownloader cacheddownloader.CachedDownloader
	extractor        extractor.Extractor
	tempDir          string
	logger           lager.Logger
}

func New(
	container warden.Container,
	model models.DownloadAction,
	cachedDownloader cacheddownloader.CachedDownloader,
	extractor extractor.Extractor,
	tempDir string,
	logger lager.Logger,
) *DownloadStep {
	logger = logger.Session("DownloadAction", lager.Data{
		"from":     model.From,
		"to":       model.To,
		"cacheKey": model.CacheKey,
	})
	return &DownloadStep{
		container:        container,
		model:            model,
		cachedDownloader: cachedDownloader,
		extractor:        extractor,
		tempDir:          tempDir,
		logger:           logger,
	}
}

func (step *DownloadStep) Perform() error {
	step.logger.Info("starting")

	//Stream this to the extractor + container when we have streaming support!
	downloadedPath, err := step.download()
	if err != nil {
		return err
	}

	defer os.Remove(downloadedPath)

	if step.model.Extract {
		return step.extract(downloadedPath)
	}

	return step.streamToContainer(downloadedPath)
}

func (step *DownloadStep) download() (string, error) {
	url, err := url.ParseRequestURI(step.model.From)
	if err != nil {
		step.logger.Error("parse-request-uri-error", err)
		return "", err
	}

	downloadedFile, err := step.cachedDownloader.Fetch(url, step.model.CacheKey)
	if err != nil {
		step.logger.Error("cached-downloader-fetch-error", err)
		return "", err
	}
	defer downloadedFile.Close()

	tempFile, err := ioutil.TempFile(step.tempDir, "downloaded")
	if err != nil {
		step.logger.Error("tempfile-create-error", err)
		return "", err
	}

	_, err = io.Copy(tempFile, downloadedFile)
	tempFile.Close()
	if err != nil {
		os.Remove(tempFile.Name())
		step.logger.Error("tempfile-copy-failed", err)
		return "", err
	}

	step.logger.Info("download-successful", lager.Data{
		"tempFile": tempFile.Name(),
	})
	return tempFile.Name(), nil
}

func (step *DownloadStep) extract(downloadedPath string) error {
	extractionDir, err := ioutil.TempDir(step.tempDir, "extracted")
	if err != nil {
		step.logger.Error("extract-failed-to-create-tempdir", err, lager.Data{
			"tempDir": step.tempDir,
		})
		return err
	}
	defer os.RemoveAll(extractionDir)

	err = step.extractor.Extract(downloadedPath, extractionDir)
	if err != nil {
		step.logger.Error("failed-to-extract", err, lager.Data{
			"extractionDir":  extractionDir,
			"downloadedPath": downloadedPath,
		})
		return emittable_error.New(err, "Extraction failed")
	}

	err = step.copyExtractedFiles(extractionDir, step.model.To)
	if err != nil {
		return emittable_error.New(err, "Copying into the container failed")
	}

	step.logger.Info("extract-successful")
	return nil
}

func (step *DownloadStep) streamToContainer(downloadedPath string) error {
	reader, writer := io.Pipe()

	go writeTarTo(filepath.Base(step.model.To), downloadedPath, writer)

	err := step.streamIn(filepath.Dir(step.model.To), reader)
	if err != nil {
		return err
	}

	step.logger.Info("stream-to-container-successful")
	return nil
}

func (step *DownloadStep) Cancel() {}

func (step *DownloadStep) Cleanup() {}

func (step *DownloadStep) copyExtractedFiles(source string, destination string) error {
	reader, writer := io.Pipe()

	go func() {
		err := compressor.WriteTar(source+string(filepath.Separator), writer)
		if err == nil {
			writer.Close()
		} else {
			step.logger.Error("wite-tar-failed", err, lager.Data{
				"source":      source,
				"destination": destination,
			})
			writer.CloseWithError(err)
		}
	}()

	return step.streamIn(destination, reader)
}

func (step *DownloadStep) streamIn(destination string, reader io.Reader) error {
	err := step.container.StreamIn(destination, reader)
	if err != nil {
		step.logger.Error("failed-to-stream-in", err, lager.Data{
			"destination": destination,
		})
		return emittable_error.New(err, "Copying into the container failed")
	}

	return err
}

func writeTarTo(name string, sourcePath string, destination *io.PipeWriter) {
	source, err := os.Open(sourcePath)
	if err != nil {
		destination.CloseWithError(err)
		return
	}
	defer source.Close()

	tarWriter := tar.NewWriter(destination)

	fileInfo, err := source.Stat()
	if err != nil {
		destination.CloseWithError(err)
		return
	}

	err = tarWriter.WriteHeader(&tar.Header{
		Name:       name,
		Size:       fileInfo.Size(),
		Mode:       0644,
		AccessTime: time.Now(),
		ChangeTime: time.Now(),
	})
	if err != nil {
		destination.CloseWithError(err)
		return
	}

	_, err = io.Copy(tarWriter, source)
	if err != nil {
		destination.CloseWithError(err)
		return
	}

	if err := tarWriter.Flush(); err != nil {
		destination.CloseWithError(err)
		return
	}

	destination.Close()
}
<|endoftext|>


<commit_before>package assured

import (
	"bytes"
	"context"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	kitlog "github.com/go-kit/kit/log"
	"github.com/stretchr/testify/require"
)

func TestClient(t *testing.T) {
	httpClient := &http.Client{}
	ctx := context.Background()
	settings := Settings{
		Logger:         kitlog.NewLogfmtLogger(ioutil.Discard),
		Port:           9091,
		TrackMadeCalls: true,
		HTTPClient:     *httpClient,
	}
	client := NewClient(ctx, settings)

	url := client.URL()
	require.Equal(t, "http://localhost:9091/when", url)

	require.NoError(t, client.Given(*testCall1()))
	require.NoError(t, client.Given(*testCall2()))
	require.NoError(t, client.Given(*testCall3()))

	req, err := http.NewRequest(http.MethodGet, url+"/test/assured", bytes.NewReader([]byte(`{"calling":"you"}`)))
	require.NoError(t, err)

	resp, err := httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusOK, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, []byte(`{"assured": true}`), body)

	req, err = http.NewRequest(http.MethodGet, url+"/test/assured", bytes.NewReader([]byte(`{"calling":"again"}`)))
	require.NoError(t, err)

	resp, err = httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusConflict, resp.StatusCode)
	body, err = ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, []byte("error"), body)

	req, err = http.NewRequest(http.MethodPost, url+"/teapot/assured", bytes.NewReader([]byte(`{"calling":"here"}`)))
	require.NoError(t, err)

	resp, err = httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusTeapot, resp.StatusCode)
	body, err = ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, []byte{}, body)

	calls, err := client.Verify("GET", "test/assured")
	require.NoError(t, err)
	require.Equal(t, []Call{
		Call{
			Method:     "GET",
			Path:       "test/assured",
			StatusCode: 200,
			Response:   []byte(`{"calling":"you"}`),
			Headers:    map[string]string{"Content-Length": "17", "User-Agent": "Go-http-client/1.1", "Accept-Encoding": "gzip"}},
		Call{
			Method:     "GET",
			Path:       "test/assured",
			StatusCode: 200,
			Response:   []byte(`{"calling":"again"}`),
			Headers:    map[string]string{"Content-Length": "19", "User-Agent": "Go-http-client/1.1", "Accept-Encoding": "gzip"}}}, calls)

	calls, err = client.Verify("POST", "teapot/assured")
	require.NoError(t, err)
	require.Equal(t, []Call{
		Call{
			Method:     "POST",
			Path:       "teapot/assured",
			StatusCode: 200,
			Response:   []byte(`{"calling":"here"}`),
			Headers:    map[string]string{"Content-Length": "18", "User-Agent": "Go-http-client/1.1", "Accept-Encoding": "gzip"}}}, calls)

	err = client.Clear("GET", "test/assured")
	require.NoError(t, err)

	calls, err = client.Verify("GET", "test/assured")
	require.NoError(t, err)
	require.Nil(t, calls)

	calls, err = client.Verify("POST", "teapot/assured")
	require.NoError(t, err)
	require.Equal(t, []Call{
		Call{
			Method:     "POST",
			Path:       "teapot/assured",
			StatusCode: 200,
			Response:   []byte(`{"calling":"here"}`),
			Headers:    map[string]string{"Content-Length": "18", "User-Agent": "Go-http-client/1.1", "Accept-Encoding": "gzip"}}}, calls)

	err = client.ClearAll()
	require.NoError(t, err)

	calls, err = client.Verify("GET", "test/assured")
	require.NoError(t, err)
	require.Nil(t, calls)

	calls, err = client.Verify("POST", "teapot/assured")
	require.NoError(t, err)
	require.Nil(t, calls)
}

func TestClientCallbacks(t *testing.T) {
	httpClient := http.Client{}
	called := false
	delayCalled := false
	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		require.NoError(t, err)
		require.Equal(t, []byte(`{"done":"here"}`), body)
		require.NotEmpty(t, r.Header.Get("x-info"))
		called = true
	}))
	delayTestServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		require.NoError(t, err)
		require.Equal(t, []byte(`{"wait":"there's more"}`), body)
		delayCalled = true
	}))
	client := NewDefaultClient()

	require.NoError(t, client.Given(Call{
		Path:   "test/assured",
		Method: "POST",
		Delay:  2,
		Callbacks: []Callback{
			Callback{
				Method:   "POST",
				Target:   testServer.URL,
				Response: []byte(`{"done":"here"}`),
				Headers:  map[string]string{"x-info": "important"},
			},
			Callback{
				Method:   "POST",
				Target:   delayTestServer.URL,
				Delay:    4,
				Response: []byte(`{"wait":"there's more"}`),
			},
		},
	}))

	req, err := http.NewRequest(http.MethodPost, client.URL()+"/test/assured", bytes.NewReader([]byte(`{"calling":"here"}`)))
	require.NoError(t, err)

	start := time.Now()
	_, err = httpClient.Do(req)
	require.NoError(t, err)

	require.True(t, time.Since(start) >= 2*time.Second, "response should be delayed 2 seconds")
	// allow go routine to finish
	time.Sleep(1 * time.Second)
	require.True(t, called, "callback was not hit")
	require.False(t, delayCalled, "delayed callback should not be hit yet")
	time.Sleep(2 * time.Second)
	require.True(t, delayCalled, "delayed callback was not hit")
}

func TestClientClose(t *testing.T) {
	client := NewDefaultClient()
	client2 := NewDefaultClient()

	require.NotEqual(t, client.URL(), client2.URL())

	require.NoError(t, client.Given(*testCall1()))
	require.NoError(t, client2.Given(*testCall1()))

	client.Close()
	err := client.Given(*testCall1())

	require.Error(t, err)
	require.Contains(t, err.Error(), `connection refused`)

	client2.Close()
	err = client2.Given(*testCall1())

	require.Error(t, err)
	require.Contains(t, err.Error(), `connection refused`)
}

func TestClientGivenNoMethod(t *testing.T) {
	httpClient := http.Client{}
	client := NewDefaultClient()

	err := client.Given(Call{Path: "NoMethodMan"})
	require.NoError(t, err)

	req, err := http.NewRequest(http.MethodGet, client.URL()+"/NoMethodMan", nil)
	require.NoError(t, err)

	resp, err := httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusOK, resp.StatusCode)
}

func TestClientGivenCallbackMissingTarget(t *testing.T) {
	call := Call{
		Method: "POST",
		Callbacks: []Callback{
			Callback{Method: "POST"},
		},
	}
	client := NewDefaultClient()

	err := client.Given(call)

	require.Error(t, err)
	require.Equal(t, "cannot stub callback without target", err.Error())
}

func TestClientGivenCallbackBadMethod(t *testing.T) {
	call := Call{
		Method: "POST",
		Callbacks: []Callback{
			Callback{Method: "\"", Target: "http://localhost/"},
		},
	}
	client := NewDefaultClient()

	err := client.Given(call)

	require.Error(t, err)
	require.Equal(t, "net/http: invalid method \"\\\"\"", err.Error())
}

func TestClientBadRequestFailure(t *testing.T) {
	client := NewDefaultClient()

	err := client.Given(Call{Method: "\"", Path: "goat/path"})

	require.Error(t, err)
	require.Equal(t, `net/http: invalid method "\""`, err.Error())

	err = client.Given(Call{Method: "\"", Path: "goat/path", Response: []byte("goats among men")})

	require.Error(t, err)
	require.Equal(t, `net/http: invalid method "\""`, err.Error())

	calls, err := client.Verify("\"", "goat/path")

	require.Error(t, err)
	require.Equal(t, `net/http: invalid method "\""`, err.Error())
	require.Nil(t, calls)

	err = client.Clear("\"", "goat/path")

	require.Error(t, err)
	require.Equal(t, `net/http: invalid method "\""`, err.Error())

	client.Port = -1
	err = client.ClearAll()

	require.Error(t, err)
	require.Equal(t, `Delete http://localhost:-1/clear: invalid URL port "-1"`, err.Error())
}

func TestClientVerifyHttpClientFailure(t *testing.T) {
	client := NewDefaultClient()
	client.Port = 1

	calls, err := client.Verify("GONE", "not/started")

	require.Error(t, err)
	require.Contains(t, err.Error(), `connection refused`)
	require.Nil(t, calls)
}

func TestClientVerifyResponseFailure(t *testing.T) {
	client := NewDefaultClient()
	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
	}))
	defer testServer.Close()
	index := strings.LastIndex(testServer.URL, ":")
	port, err := strconv.ParseInt(testServer.URL[index+1:], 10, 64)
	require.NoError(t, err)
	client.Port = int(port)

	calls, err := client.Verify("GONE", "not/started")

	require.Error(t, err)
	require.Equal(t, `failure to verify calls`, err.Error())
	require.Nil(t, calls)
}

func TestClientVerifyBodyFailure(t *testing.T) {
	client := NewDefaultClient()
	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode("ydob+dab")
	}))
	defer testServer.Close()
	index := strings.LastIndex(testServer.URL, ":")
	port, err := strconv.ParseInt(testServer.URL[index+1:], 10, 64)
	require.NoError(t, err)
	client.Port = int(port)

	calls, err := client.Verify("BODY", "bad+body")

	require.Error(t, err)
	require.Equal(t, `json: cannot unmarshal string into Go value of type []assured.Call`, err.Error())
	require.Nil(t, calls)
}

func TestClientPathSanitization(t *testing.T) {
	httpClient := &http.Client{}
	client := NewDefaultClient()

	require.NoError(t, client.Given(Call{Method: "GET", Path: "///yoyo/path///", StatusCode: http.StatusAccepted}))

	req, err := http.NewRequest(http.MethodGet, client.URL()+"/yoyo/path", nil)
	require.NoError(t, err)

	resp, err := httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusAccepted, resp.StatusCode)
}
<commit_msg>Fix client test.<commit_after>package assured

import (
	"bytes"
	"context"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	kitlog "github.com/go-kit/kit/log"
	"github.com/stretchr/testify/require"
)

func TestClient(t *testing.T) {
	httpClient := &http.Client{}
	ctx := context.Background()
	settings := Settings{
		Logger:         kitlog.NewLogfmtLogger(ioutil.Discard),
		Port:           9091,
		TrackMadeCalls: true,
		HTTPClient:     *httpClient,
	}
	client := NewClient(ctx, settings)

	url := client.URL()
	require.Equal(t, "http://localhost:9091/when", url)

	require.NoError(t, client.Given(*testCall1()))
	require.NoError(t, client.Given(*testCall2()))
	require.NoError(t, client.Given(*testCall3()))

	req, err := http.NewRequest(http.MethodGet, url+"/test/assured", bytes.NewReader([]byte(`{"calling":"you"}`)))
	require.NoError(t, err)

	resp, err := httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusOK, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, []byte(`{"assured": true}`), body)

	req, err = http.NewRequest(http.MethodGet, url+"/test/assured", bytes.NewReader([]byte(`{"calling":"again"}`)))
	require.NoError(t, err)

	resp, err = httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusConflict, resp.StatusCode)
	body, err = ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, []byte("error"), body)

	req, err = http.NewRequest(http.MethodPost, url+"/teapot/assured", bytes.NewReader([]byte(`{"calling":"here"}`)))
	require.NoError(t, err)

	resp, err = httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusTeapot, resp.StatusCode)
	body, err = ioutil.ReadAll(resp.Body)
	require.NoError(t, err)
	require.Equal(t, []byte{}, body)

	calls, err := client.Verify("GET", "test/assured")
	require.NoError(t, err)
	require.Equal(t, []Call{
		Call{
			Method:     "GET",
			Path:       "test/assured",
			StatusCode: 200,
			Response:   []byte(`{"calling":"you"}`),
			Headers:    map[string]string{"Content-Length": "17", "User-Agent": "Go-http-client/1.1", "Accept-Encoding": "gzip"}},
		Call{
			Method:     "GET",
			Path:       "test/assured",
			StatusCode: 200,
			Response:   []byte(`{"calling":"again"}`),
			Headers:    map[string]string{"Content-Length": "19", "User-Agent": "Go-http-client/1.1", "Accept-Encoding": "gzip"}}}, calls)

	calls, err = client.Verify("POST", "teapot/assured")
	require.NoError(t, err)
	require.Equal(t, []Call{
		Call{
			Method:     "POST",
			Path:       "teapot/assured",
			StatusCode: 200,
			Response:   []byte(`{"calling":"here"}`),
			Headers:    map[string]string{"Content-Length": "18", "User-Agent": "Go-http-client/1.1", "Accept-Encoding": "gzip"}}}, calls)

	err = client.Clear("GET", "test/assured")
	require.NoError(t, err)

	calls, err = client.Verify("GET", "test/assured")
	require.NoError(t, err)
	require.Nil(t, calls)

	calls, err = client.Verify("POST", "teapot/assured")
	require.NoError(t, err)
	require.Equal(t, []Call{
		Call{
			Method:     "POST",
			Path:       "teapot/assured",
			StatusCode: 200,
			Response:   []byte(`{"calling":"here"}`),
			Headers:    map[string]string{"Content-Length": "18", "User-Agent": "Go-http-client/1.1", "Accept-Encoding": "gzip"}}}, calls)

	err = client.ClearAll()
	require.NoError(t, err)

	calls, err = client.Verify("GET", "test/assured")
	require.NoError(t, err)
	require.Nil(t, calls)

	calls, err = client.Verify("POST", "teapot/assured")
	require.NoError(t, err)
	require.Nil(t, calls)
}

func TestClientCallbacks(t *testing.T) {
	httpClient := http.Client{}
	called := false
	delayCalled := false
	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		require.NoError(t, err)
		require.Equal(t, []byte(`{"done":"here"}`), body)
		require.NotEmpty(t, r.Header.Get("x-info"))
		called = true
	}))
	delayTestServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := ioutil.ReadAll(r.Body)
		require.NoError(t, err)
		require.Equal(t, []byte(`{"wait":"there's more"}`), body)
		delayCalled = true
	}))
	client := NewDefaultClient()

	require.NoError(t, client.Given(Call{
		Path:   "test/assured",
		Method: "POST",
		Delay:  2,
		Callbacks: []Callback{
			Callback{
				Method:   "POST",
				Target:   testServer.URL,
				Response: []byte(`{"done":"here"}`),
				Headers:  map[string]string{"x-info": "important"},
			},
			Callback{
				Method:   "POST",
				Target:   delayTestServer.URL,
				Delay:    4,
				Response: []byte(`{"wait":"there's more"}`),
			},
		},
	}))

	req, err := http.NewRequest(http.MethodPost, client.URL()+"/test/assured", bytes.NewReader([]byte(`{"calling":"here"}`)))
	require.NoError(t, err)

	start := time.Now()
	_, err = httpClient.Do(req)
	require.NoError(t, err)

	require.True(t, time.Since(start) >= 2*time.Second, "response should be delayed 2 seconds")
	// allow go routine to finish
	time.Sleep(1 * time.Second)
	require.True(t, called, "callback was not hit")
	require.False(t, delayCalled, "delayed callback should not be hit yet")
	time.Sleep(2 * time.Second)
	require.True(t, delayCalled, "delayed callback was not hit")
}

func TestClientClose(t *testing.T) {
	client := NewDefaultClient()
	client2 := NewDefaultClient()

	require.NotEqual(t, client.URL(), client2.URL())

	require.NoError(t, client.Given(*testCall1()))
	require.NoError(t, client2.Given(*testCall1()))

	client.Close()
	err := client.Given(*testCall1())

	require.Error(t, err)
	require.Contains(t, err.Error(), `connection refused`)

	client2.Close()
	err = client2.Given(*testCall1())

	require.Error(t, err)
	require.Contains(t, err.Error(), `connection refused`)
}

func TestClientGivenNoMethod(t *testing.T) {
	httpClient := http.Client{}
	client := NewDefaultClient()

	err := client.Given(Call{Path: "NoMethodMan"})
	require.NoError(t, err)

	req, err := http.NewRequest(http.MethodGet, client.URL()+"/NoMethodMan", nil)
	require.NoError(t, err)

	resp, err := httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusOK, resp.StatusCode)
}

func TestClientGivenCallbackMissingTarget(t *testing.T) {
	call := Call{
		Method: "POST",
		Callbacks: []Callback{
			Callback{Method: "POST"},
		},
	}
	client := NewDefaultClient()

	err := client.Given(call)

	require.Error(t, err)
	require.Equal(t, "cannot stub callback without target", err.Error())
}

func TestClientGivenCallbackBadMethod(t *testing.T) {
	call := Call{
		Method: "POST",
		Callbacks: []Callback{
			Callback{Method: "\"", Target: "http://localhost/"},
		},
	}
	client := NewDefaultClient()

	err := client.Given(call)

	require.Error(t, err)
	require.Equal(t, "net/http: invalid method \"\\\"\"", err.Error())
}

func TestClientBadRequestFailure(t *testing.T) {
	client := NewDefaultClient()

	err := client.Given(Call{Method: "\"", Path: "goat/path"})

	require.Error(t, err)
	require.Equal(t, `net/http: invalid method "\""`, err.Error())

	err = client.Given(Call{Method: "\"", Path: "goat/path", Response: []byte("goats among men")})

	require.Error(t, err)
	require.Equal(t, `net/http: invalid method "\""`, err.Error())

	calls, err := client.Verify("\"", "goat/path")

	require.Error(t, err)
	require.Equal(t, `net/http: invalid method "\""`, err.Error())
	require.Nil(t, calls)

	err = client.Clear("\"", "goat/path")

	require.Error(t, err)
	require.Equal(t, `net/http: invalid method "\""`, err.Error())

	client.Port = -1
	err = client.ClearAll()

	require.Error(t, err)
	require.Equal(t, `parse http://localhost:-1/clear: invalid port ":-1" after host`, err.Error())
}

func TestClientVerifyHttpClientFailure(t *testing.T) {
	client := NewDefaultClient()
	client.Port = 1

	calls, err := client.Verify("GONE", "not/started")

	require.Error(t, err)
	require.Contains(t, err.Error(), `connection refused`)
	require.Nil(t, calls)
}

func TestClientVerifyResponseFailure(t *testing.T) {
	client := NewDefaultClient()
	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
	}))
	defer testServer.Close()
	index := strings.LastIndex(testServer.URL, ":")
	port, err := strconv.ParseInt(testServer.URL[index+1:], 10, 64)
	require.NoError(t, err)
	client.Port = int(port)

	calls, err := client.Verify("GONE", "not/started")

	require.Error(t, err)
	require.Equal(t, `failure to verify calls`, err.Error())
	require.Nil(t, calls)
}

func TestClientVerifyBodyFailure(t *testing.T) {
	client := NewDefaultClient()
	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode("ydob+dab")
	}))
	defer testServer.Close()
	index := strings.LastIndex(testServer.URL, ":")
	port, err := strconv.ParseInt(testServer.URL[index+1:], 10, 64)
	require.NoError(t, err)
	client.Port = int(port)

	calls, err := client.Verify("BODY", "bad+body")

	require.Error(t, err)
	require.Equal(t, `json: cannot unmarshal string into Go value of type []assured.Call`, err.Error())
	require.Nil(t, calls)
}

func TestClientPathSanitization(t *testing.T) {
	httpClient := &http.Client{}
	client := NewDefaultClient()

	require.NoError(t, client.Given(Call{Method: "GET", Path: "///yoyo/path///", StatusCode: http.StatusAccepted}))

	req, err := http.NewRequest(http.MethodGet, client.URL()+"/yoyo/path", nil)
	require.NoError(t, err)

	resp, err := httpClient.Do(req)
	require.NoError(t, err)
	require.Equal(t, http.StatusAccepted, resp.StatusCode)
}
<|endoftext|>