<commit_before>package ebpf

import (
	"encoding/binary"
	"io"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"unsafe"

	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

const (
	perfTypeSoftware     = 1
	perfCountSWBPFOutput = 10
	perfSampleRaw        = 1 << 10
)

type perfEventMeta struct {
	_          [128]uint64 /* Pad to 1 k, ignore fields */
	dataHead   uint64      /* head in the data section */
	dataTail   uint64      /* user-space written tail */
	dataOffset uint64      /* where the buffer starts */
	dataSize   uint64      /* data buffer size */
}

type perfEventHeader struct {
	Type uint32
	Misc uint16
	Size uint16
}

// perfEventRing is a page of metadata followed by
// a variable number of pages which form a ring buffer.
type perfEventRing struct {
	fd   int
	meta *perfEventMeta
	mmap []byte
	ring []byte
}

func newPerfEventRing(cpu int, opts PerfReaderOptions) (*perfEventRing, error) {
	const flagWakeupWatermark = 1 << 14

	if opts.Watermark >= opts.PerCPUBuffer {
		return nil, errors.Errorf("Watermark must be smaller than PerCPUBuffer")
	}

	// Round to nearest page boundary and allocate
	// an extra page for meta data
	pageSize := os.Getpagesize()
	nPages := (opts.PerCPUBuffer + pageSize - 1) / pageSize
	size := (1 + nPages) * pageSize

	attr := perfEventAttr{
		perfType:                perfTypeSoftware,
		config:                  perfCountSWBPFOutput,
		flags:                   flagWakeupWatermark,
		sampleType:              perfSampleRaw,
		wakeupEventsOrWatermark: uint32(opts.Watermark),
	}

	fd, err := perfEventOpen(&attr, -1, cpu, -1, 0)
	if err != nil {
		return nil, err
	}

	if err := unix.SetNonblock(fd, true); err != nil {
		unix.Close(fd)
		return nil, err
	}

	mmap, err := unix.Mmap(fd, 0, size, unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED)
	if err != nil {
		unix.Close(fd)
		return nil, err
	}

	// This relies on the fact that we allocate an extra metadata page,
	// and that the struct is smaller than an OS page.
	// This use of unsafe.Pointer isn't explicitly sanctioned by the
	// documentation, since a byte is smaller than sampledPerfEvent.
	meta := (*perfEventMeta)(unsafe.Pointer(&mmap[0]))

	return &perfEventRing{
		fd:   fd,
		meta: meta,
		mmap: mmap,
		ring: mmap[meta.dataOffset : meta.dataOffset+meta.dataSize],
	}, nil
}

func (ring *perfEventRing) Close() {
	unix.Close(ring.fd)
	unix.Munmap(ring.mmap)
}

func readRecord(rd io.Reader) (*PerfSample, uint64, error) {
	const (
		perfRecordLost   = 2
		perfRecordSample = 9
	)

	var header perfEventHeader
	err := binary.Read(rd, nativeEndian, &header)
	if err == io.EOF {
		return nil, 0, nil
	}

	if err != nil {
		return nil, 0, errors.Wrap(err, "can't read event header")
	}

	switch header.Type {
	case perfRecordLost:
		lost, err := readLostRecords(rd)
		if err != nil {
			return nil, 0, err
		}

		return nil, lost, nil

	case perfRecordSample:
		sample, err := readSample(rd)
		if err != nil {
			return nil, 0, err
		}

		return sample, 0, nil

	default:
		return nil, 0, errors.Errorf("unknown event type %d", header.Type)
	}
}

func readLostRecords(rd io.Reader) (uint64, error) {
	var lostHeader struct {
		ID   uint64
		Lost uint64
	}

	err := binary.Read(rd, nativeEndian, &lostHeader)
	if err != nil {
		return 0, errors.Wrap(err, "can't read lost records header")
	}

	return lostHeader.Lost, nil
}

func readSample(rd io.Reader) (*PerfSample, error) {
	var size uint32
	if err := binary.Read(rd, nativeEndian, &size); err != nil {
		return nil, errors.Wrap(err, "can't read sample size")
	}

	data := make([]byte, int(size))
	_, err := io.ReadFull(rd, data)
	return &PerfSample{data}, errors.Wrap(err, "can't read sample")
}

// PerfSample is read from the kernel by PerfReader.
type PerfSample struct {
	// Data are padded with 0 to have a 64-bit alignment.
	// If you are using variable length samples you need to take
	// this into account.
	Data []byte
}

// PerfReader allows reading bpf_perf_event_output
// from user space.
type PerfReader struct {
	lostSamples uint64
	// Closing a PERF_EVENT_ARRAY removes all event fds
	// stored in it, so we keep a reference alive.
	array *Map

	// Eventfds for closing
	closeFd      int
	flushCloseFd int
	// Ensure we only close once
	closeOnce sync.Once
	// Channel to interrupt polling blocked on writing to consumer
	stopWriter chan struct{}
	// Channel closed when poll() is done
	closed chan struct{}

	// Error receives a write if the reader exits
	// due to an error.
	Error <-chan error

	// Samples is closed when the Reader exits.
	Samples <-chan *PerfSample
}

// PerfReaderOptions control the behaviour of the user
// space reader.
type PerfReaderOptions struct {
	// A map of type PerfEventArray. The reader takes ownership of the
	// map and takes care of closing it.
	Map *Map
	// Controls the size of the per CPU buffer in bytes. LostSamples() will
	// increase if the buffer is too small.
	PerCPUBuffer int
	// The reader will start processing samples once the per CPU buffer
	// exceeds this value. Must be smaller than PerCPUBuffer.
	Watermark int
}

// NewPerfReader creates a new reader with the given options.
//
// The value returned by LostSamples() will increase if the buffer
// isn't large enough to contain all incoming samples.
func NewPerfReader(opts PerfReaderOptions) (out *PerfReader, err error) {
	if opts.PerCPUBuffer < 1 {
		return nil, errors.New("PerCPUBuffer must be larger than 0")
	}

	// We can't create a ring for CPUs that aren't online, so use only the online (of possible) CPUs
	nCPU, err := onlineCPUs()
	if err != nil {
		return nil, errors.Wrap(err, "sampled perf event")
	}

	var (
		fds   []int
		rings = make(map[int]*perfEventRing)
	)

	defer func() {
		if err != nil {
			for _, ring := range rings {
				ring.Close()
			}
		}
	}()

	// bpf_perf_event_output checks which CPU an event is enabled on,
	// but doesn't allow using a wildcard like -1 to specify "all CPUs".
	// Hence we have to create a ring for each CPU.
	for i := 0; i < nCPU; i++ {
		ring, err := newPerfEventRing(i, opts)
		if err != nil {
			return nil, errors.Wrapf(err, "failed to create perf ring for CPU %d", i)
		}

		if err := opts.Map.Put(uint32(i), uint32(ring.fd)); err != nil {
			ring.Close()
			return nil, errors.Wrapf(err, "could't put event fd for CPU %d", i)
		}

		fds = append(fds, ring.fd)
		rings[ring.fd] = ring
	}

	closeFd, err := newEventFd()
	if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			unix.Close(closeFd)
		}
	}()
	fds = append(fds, closeFd)

	flushCloseFd, err := newEventFd()
	if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			unix.Close(flushCloseFd)
		}
	}()
	fds = append(fds, flushCloseFd)

	epollFd, err := newEpollFd(fds...)
	if err != nil {
		return nil, err
	}

	samples := make(chan *PerfSample, nCPU)
	errs := make(chan error, 1)

	out = &PerfReader{
		array:        opts.Map,
		closeFd:      closeFd,
		flushCloseFd: flushCloseFd,
		stopWriter:   make(chan struct{}),
		closed:       make(chan struct{}),
		Error:        errs,
		Samples:      samples,
	}
	runtime.SetFinalizer(out, (*PerfReader).Close)

	go out.poll(epollFd, rings, samples, errs)

	return out, nil
}

// LostSamples returns the number of samples dropped
// by the perf subsystem.
func (pr *PerfReader) LostSamples() uint64 {
	return atomic.LoadUint64(&pr.lostSamples)
}

// Close stops the reader, discarding any samples not yet written to 'Samples'.
//
// Calls to perf_event_output from eBPF programs will return
// ENOENT after calling this method.
func (pr *PerfReader) Close() (err error) {
	return pr.close(false)
}

// FlushAndClose stops the reader, flushing any samples to 'Samples'.
// Will block if no consumer reads from 'Samples'.
//
// Calls to perf_event_output from eBPF programs will return
// ENOENT after calling this method.
func (pr *PerfReader) FlushAndClose() error {
	return pr.close(true)
}

func (pr *PerfReader) close(flush bool) error {
	pr.closeOnce.Do(func() {
		runtime.SetFinalizer(pr, nil)

		// Interrupt polling so we don't deadlock if the consumer is dead
		if !flush {
			close(pr.stopWriter)
		}

		// Signal poll() via the event fd. Ignore the
		// write error since poll() may have exited
		// and closed the fd already
		var value [8]byte
		nativeEndian.PutUint64(value[:], 1)
		if flush {
			_, _ = unix.Write(pr.flushCloseFd, value[:])
		} else {
			_, _ = unix.Write(pr.closeFd, value[:])
		}
	})

	// Wait until poll is done
	<-pr.closed

	return nil
}

func (pr *PerfReader) poll(epollFd int, rings map[int]*perfEventRing, samples chan<- *PerfSample, errs chan<- error) {
	// last as it means we're done
	defer close(pr.closed)
	defer close(samples)
	defer pr.array.Close()
	defer unix.Close(epollFd)
	defer unix.Close(pr.closeFd)
	defer unix.Close(pr.flushCloseFd)
	defer func() {
		for _, ring := range rings {
			ring.Close()
		}
	}()

	epollEvents := make([]unix.EpollEvent, len(rings)+1)

	for {
		nEvents, err := unix.EpollWait(epollFd, epollEvents, -1)
		if err != nil {
			// Handle EINTR
			if temp, ok := err.(temporaryError); ok && temp.Temporary() {
				continue
			}

			errs <- err
			return
		}

		for _, event := range epollEvents[:nEvents] {
			fd := int(event.Fd)
			if fd == pr.closeFd {
				// We were woken by Close via the close fd
				return
			}

			if fd == pr.flushCloseFd {
				for _, ring := range rings {
					err := pr.flushRing(ring, samples)
					if err != nil {
						errs <- err
						return
					}
				}

				return
			}

			err := pr.flushRing(rings[fd], samples)
			if err != nil {
				errs <- err
				return
			}
		}
	}
}

func (pr *PerfReader) flushRing(ring *perfEventRing, samples chan<- *PerfSample) error {
	rd := newRingReader(ring.meta, ring.ring)
	defer rd.Close()

	var totalLost uint64

	for {
		sample, lost, err := readRecord(rd)
		if err != nil {
			return err
		}

		if lost > 0 {
			totalLost += lost
			continue
		}

		if sample == nil {
			break
		}

		select {
		case samples <- sample:
		case <-pr.stopWriter:
			break
		}
	}

	if totalLost > 0 {
		atomic.AddUint64(&pr.lostSamples, totalLost)
	}
	return nil
}

type ringReader struct {
	meta       *perfEventMeta
	head, tail uint64
	mask       uint64
	ring       []byte
}

func newRingReader(meta *perfEventMeta, ring []byte) *ringReader {
	return &ringReader{
		meta: meta,
		head: atomic.LoadUint64(&meta.dataHead),
		tail: atomic.LoadUint64(&meta.dataTail),
		// cap is always a power of two
		mask: uint64(cap(ring) - 1),
		ring: ring,
	}
}

func (rb *ringReader) Close() error {
	// Commit the new tail. This lets the kernel know that
	// the ring buffer has been consumed.
	atomic.StoreUint64(&rb.meta.dataTail, rb.tail)
	return nil
}

func (rb *ringReader) Read(p []byte) (int, error) {
	start := int(rb.tail & rb.mask)

	n := len(p)
	// Truncate if the read wraps in the ring buffer
	if remainder := cap(rb.ring) - start; n > remainder {
		n = remainder
	}

	// Truncate if there isn't enough data
	if remainder := int(rb.head - rb.tail); n > remainder {
		n = remainder
	}

	copy(p, rb.ring[start:start+n])
	rb.tail += uint64(n)

	if rb.tail == rb.head {
		return n, io.EOF
	}

	return n, nil
}

type temporaryError interface {
	Temporary() bool
}
<commit_msg>perf: set finalizer on perfEventRing<commit_after>package ebpf

import (
	"encoding/binary"
	"io"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"unsafe"

	"github.com/pkg/errors"
	"golang.org/x/sys/unix"
)

const (
	perfTypeSoftware     = 1
	perfCountSWBPFOutput = 10
	perfSampleRaw        = 1 << 10
)

type perfEventMeta struct {
	_          [128]uint64 /* Pad to 1 k, ignore fields */
	dataHead   uint64      /* head in the data section */
	dataTail   uint64      /* user-space written tail */
	dataOffset uint64      /* where the buffer starts */
	dataSize   uint64      /* data buffer size */
}

type perfEventHeader struct {
	Type uint32
	Misc uint16
	Size uint16
}

// perfEventRing is a page of metadata followed by
// a variable number of pages which form a ring buffer.
type perfEventRing struct {
	fd   int
	meta *perfEventMeta
	mmap []byte
	ring []byte
}

func newPerfEventRing(cpu int, opts PerfReaderOptions) (*perfEventRing, error) {
	const flagWakeupWatermark = 1 << 14

	if opts.Watermark >= opts.PerCPUBuffer {
		return nil, errors.Errorf("Watermark must be smaller than PerCPUBuffer")
	}

	// Round to nearest page boundary and allocate
	// an extra page for meta data
	pageSize := os.Getpagesize()
	nPages := (opts.PerCPUBuffer + pageSize - 1) / pageSize
	size := (1 + nPages) * pageSize

	attr := perfEventAttr{
		perfType:                perfTypeSoftware,
		config:                  perfCountSWBPFOutput,
		flags:                   flagWakeupWatermark,
		sampleType:              perfSampleRaw,
		wakeupEventsOrWatermark: uint32(opts.Watermark),
	}

	fd, err := perfEventOpen(&attr, -1, cpu, -1, 0)
	if err != nil {
		return nil, err
	}

	if err := unix.SetNonblock(fd, true); err != nil {
		unix.Close(fd)
		return nil, err
	}

	mmap, err := unix.Mmap(fd, 0, size, unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED)
	if err != nil {
		unix.Close(fd)
		return nil, err
	}

	// This relies on the fact that we allocate an extra metadata page,
	// and that the struct is smaller than an OS page.
	// This use of unsafe.Pointer isn't explicitly sanctioned by the
	// documentation, since a byte is smaller than sampledPerfEvent.
	meta := (*perfEventMeta)(unsafe.Pointer(&mmap[0]))

	ring := &perfEventRing{
		fd:   fd,
		meta: meta,
		mmap: mmap,
		ring: mmap[meta.dataOffset : meta.dataOffset+meta.dataSize],
	}
	runtime.SetFinalizer(ring, (*perfEventRing).Close)

	return ring, nil
}

func (ring *perfEventRing) Close() {
	runtime.SetFinalizer(ring, nil)
	unix.Close(ring.fd)
	unix.Munmap(ring.mmap)
}

func readRecord(rd io.Reader) (*PerfSample, uint64, error) {
	const (
		perfRecordLost   = 2
		perfRecordSample = 9
	)

	var header perfEventHeader
	err := binary.Read(rd, nativeEndian, &header)
	if err == io.EOF {
		return nil, 0, nil
	}

	if err != nil {
		return nil, 0, errors.Wrap(err, "can't read event header")
	}

	switch header.Type {
	case perfRecordLost:
		lost, err := readLostRecords(rd)
		if err != nil {
			return nil, 0, err
		}

		return nil, lost, nil

	case perfRecordSample:
		sample, err := readSample(rd)
		if err != nil {
			return nil, 0, err
		}

		return sample, 0, nil

	default:
		return nil, 0, errors.Errorf("unknown event type %d", header.Type)
	}
}

func readLostRecords(rd io.Reader) (uint64, error) {
	var lostHeader struct {
		ID   uint64
		Lost uint64
	}

	err := binary.Read(rd, nativeEndian, &lostHeader)
	if err != nil {
		return 0, errors.Wrap(err, "can't read lost records header")
	}

	return lostHeader.Lost, nil
}

func readSample(rd io.Reader) (*PerfSample, error) {
	var size uint32
	if err := binary.Read(rd, nativeEndian, &size); err != nil {
		return nil, errors.Wrap(err, "can't read sample size")
	}

	data := make([]byte, int(size))
	_, err := io.ReadFull(rd, data)
	return &PerfSample{data}, errors.Wrap(err, "can't read sample")
}

// PerfSample is read from the kernel by PerfReader.
type PerfSample struct {
	// Data are padded with 0 to have a 64-bit alignment.
	// If you are using variable length samples you need to take
	// this into account.
	Data []byte
}

// PerfReader allows reading bpf_perf_event_output
// from user space.
type PerfReader struct {
	lostSamples uint64
	// Closing a PERF_EVENT_ARRAY removes all event fds
	// stored in it, so we keep a reference alive.
	array *Map

	// Eventfds for closing
	closeFd      int
	flushCloseFd int
	// Ensure we only close once
	closeOnce sync.Once
	// Channel to interrupt polling blocked on writing to consumer
	stopWriter chan struct{}
	// Channel closed when poll() is done
	closed chan struct{}

	// Error receives a write if the reader exits
	// due to an error.
	Error <-chan error

	// Samples is closed when the Reader exits.
	Samples <-chan *PerfSample
}

// PerfReaderOptions control the behaviour of the user
// space reader.
type PerfReaderOptions struct {
	// A map of type PerfEventArray. The reader takes ownership of the
	// map and takes care of closing it.
	Map *Map
	// Controls the size of the per CPU buffer in bytes. LostSamples() will
	// increase if the buffer is too small.
	PerCPUBuffer int
	// The reader will start processing samples once the per CPU buffer
	// exceeds this value. Must be smaller than PerCPUBuffer.
	Watermark int
}

// NewPerfReader creates a new reader with the given options.
//
// The value returned by LostSamples() will increase if the buffer
// isn't large enough to contain all incoming samples.
func NewPerfReader(opts PerfReaderOptions) (out *PerfReader, err error) {
	if opts.PerCPUBuffer < 1 {
		return nil, errors.New("PerCPUBuffer must be larger than 0")
	}

	// We can't create a ring for CPUs that aren't online, so use only the online (of possible) CPUs
	nCPU, err := onlineCPUs()
	if err != nil {
		return nil, errors.Wrap(err, "sampled perf event")
	}

	var (
		fds   []int
		rings = make(map[int]*perfEventRing)
	)

	defer func() {
		if err != nil {
			for _, ring := range rings {
				ring.Close()
			}
		}
	}()

	// bpf_perf_event_output checks which CPU an event is enabled on,
	// but doesn't allow using a wildcard like -1 to specify "all CPUs".
	// Hence we have to create a ring for each CPU.
	for i := 0; i < nCPU; i++ {
		ring, err := newPerfEventRing(i, opts)
		if err != nil {
			return nil, errors.Wrapf(err, "failed to create perf ring for CPU %d", i)
		}

		if err := opts.Map.Put(uint32(i), uint32(ring.fd)); err != nil {
			ring.Close()
			return nil, errors.Wrapf(err, "could't put event fd for CPU %d", i)
		}

		fds = append(fds, ring.fd)
		rings[ring.fd] = ring
	}

	closeFd, err := newEventFd()
	if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			unix.Close(closeFd)
		}
	}()
	fds = append(fds, closeFd)

	flushCloseFd, err := newEventFd()
	if err != nil {
		return nil, err
	}
	defer func() {
		if err != nil {
			unix.Close(flushCloseFd)
		}
	}()
	fds = append(fds, flushCloseFd)

	epollFd, err := newEpollFd(fds...)
	if err != nil {
		return nil, err
	}

	samples := make(chan *PerfSample, nCPU)
	errs := make(chan error, 1)

	out = &PerfReader{
		array:        opts.Map,
		closeFd:      closeFd,
		flushCloseFd: flushCloseFd,
		stopWriter:   make(chan struct{}),
		closed:       make(chan struct{}),
		Error:        errs,
		Samples:      samples,
	}
	runtime.SetFinalizer(out, (*PerfReader).Close)

	go out.poll(epollFd, rings, samples, errs)

	return out, nil
}

// LostSamples returns the number of samples dropped
// by the perf subsystem.
func (pr *PerfReader) LostSamples() uint64 {
	return atomic.LoadUint64(&pr.lostSamples)
}

// Close stops the reader, discarding any samples not yet written to 'Samples'.
//
// Calls to perf_event_output from eBPF programs will return
// ENOENT after calling this method.
func (pr *PerfReader) Close() (err error) {
	return pr.close(false)
}

// FlushAndClose stops the reader, flushing any samples to 'Samples'.
// Will block if no consumer reads from 'Samples'.
//
// Calls to perf_event_output from eBPF programs will return
// ENOENT after calling this method.
func (pr *PerfReader) FlushAndClose() error {
	return pr.close(true)
}

func (pr *PerfReader) close(flush bool) error {
	pr.closeOnce.Do(func() {
		runtime.SetFinalizer(pr, nil)

		// Interrupt polling so we don't deadlock if the consumer is dead
		if !flush {
			close(pr.stopWriter)
		}

		// Signal poll() via the event fd. Ignore the
		// write error since poll() may have exited
		// and closed the fd already
		var value [8]byte
		nativeEndian.PutUint64(value[:], 1)
		if flush {
			_, _ = unix.Write(pr.flushCloseFd, value[:])
		} else {
			_, _ = unix.Write(pr.closeFd, value[:])
		}
	})

	// Wait until poll is done
	<-pr.closed

	return nil
}

func (pr *PerfReader) poll(epollFd int, rings map[int]*perfEventRing, samples chan<- *PerfSample, errs chan<- error) {
	// last as it means we're done
	defer close(pr.closed)
	defer close(samples)
	defer pr.array.Close()
	defer unix.Close(epollFd)
	defer unix.Close(pr.closeFd)
	defer unix.Close(pr.flushCloseFd)
	defer func() {
		for _, ring := range rings {
			ring.Close()
		}
	}()

	epollEvents := make([]unix.EpollEvent, len(rings)+1)

	for {
		nEvents, err := unix.EpollWait(epollFd, epollEvents, -1)
		if err != nil {
			// Handle EINTR
			if temp, ok := err.(temporaryError); ok && temp.Temporary() {
				continue
			}

			errs <- err
			return
		}

		for _, event := range epollEvents[:nEvents] {
			fd := int(event.Fd)
			if fd == pr.closeFd {
				// We were woken by Close via the close fd
				return
			}

			if fd == pr.flushCloseFd {
				for _, ring := range rings {
					err := pr.flushRing(ring, samples)
					if err != nil {
						errs <- err
						return
					}
				}

				return
			}

			err := pr.flushRing(rings[fd], samples)
			if err != nil {
				errs <- err
				return
			}
		}
	}
}

func (pr *PerfReader) flushRing(ring *perfEventRing, samples chan<- *PerfSample) error {
	rd := newRingReader(ring.meta, ring.ring)
	defer rd.Close()

	var totalLost uint64

	for {
		sample, lost, err := readRecord(rd)
		if err != nil {
			return err
		}

		if lost > 0 {
			totalLost += lost
			continue
		}

		if sample == nil {
			break
		}

		select {
		case samples <- sample:
		case <-pr.stopWriter:
			break
		}
	}

	if totalLost > 0 {
		atomic.AddUint64(&pr.lostSamples, totalLost)
	}
	return nil
}

type ringReader struct {
	meta       *perfEventMeta
	head, tail uint64
	mask       uint64
	ring       []byte
}

func newRingReader(meta *perfEventMeta, ring []byte) *ringReader {
	return &ringReader{
		meta: meta,
		head: atomic.LoadUint64(&meta.dataHead),
		tail: atomic.LoadUint64(&meta.dataTail),
		// cap is always a power of two
		mask: uint64(cap(ring) - 1),
		ring: ring,
	}
}

func (rb *ringReader) Close() error {
	// Commit the new tail. This lets the kernel know that
	// the ring buffer has been consumed.
	atomic.StoreUint64(&rb.meta.dataTail, rb.tail)
	return nil
}

func (rb *ringReader) Read(p []byte) (int, error) {
	start := int(rb.tail & rb.mask)

	n := len(p)
	// Truncate if the read wraps in the ring buffer
	if remainder := cap(rb.ring) - start; n > remainder {
		n = remainder
	}

	// Truncate if there isn't enough data
	if remainder := int(rb.head - rb.tail); n > remainder {
		n = remainder
	}

	copy(p, rb.ring[start:start+n])
	rb.tail += uint64(n)

	if rb.tail == rb.head {
		return n, io.EOF
	}

	return n, nil
}

type temporaryError interface {
	Temporary() bool
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"net/url"
	"os"

	"golang.org/x/net/context"

	"github.com/docker/docker/api/client/bundlefile"
	"github.com/docker/engine-api/client"
	"github.com/docker/engine-api/types"
	"github.com/docker/engine-api/types/swarm"
	"github.com/fatih/color"
	"github.com/urfave/cli"
)

var Replica1 uint64 = 1

type Services map[string]swarm.Service

func plan(c *cli.Context) error {
	stackName := c.Args().Get(0)

	if stackName == "" {
		return cli.NewExitError("Need to specify a stack name", 1)
	}
	dabLocation := c.String("dab")

	if dabLocation == "" {
		// Assume it is called as the stack name
		dabLocation = fmt.Sprintf("%s.dab", stackName)
	}

	var dabReader io.Reader
	if u, e := url.Parse(dabLocation); e == nil && u.IsAbs() {
		// DAB file seems to be remote, try to download it first
		return cli.NewExitError("Not implemented", 2)
	} else {
		if dabFile, err := os.Open(dabLocation); err != nil {
			return cli.NewExitError(err.Error(), 3)
		} else {
			dabReader = dabFile
		}
	}

	bundle, bundleErr := bundlefile.LoadFile(dabReader)
	if bundleErr != nil {
		return cli.NewExitError(bundleErr.Error(), 3)
	}

	swarm, swarmErr := client.NewEnvClient()
	if swarmErr != nil {
		return cli.NewExitError(swarmErr.Error(), 3)
	}

	services, servicesErr := swarm.ServiceList(context.Background(), types.ServiceListOptions{})
	if servicesErr != nil {
		return cli.NewExitError(servicesErr.Error(), 3)
	}

	expected := getBundleServicesSpec(bundle, stackName)
	current := getSwarmServicesSpecForStack(services, stackName)

	for n, es := range expected {
		if cs, found := current[n]; !found {
			color.Green("\n + %s << service will be added", n)
			// New service to add
		} else {
			color.Cyan("\n%s\n", es.Spec.Name)
			PrintServiceSpecDiff(cs.Spec, es.Spec)
		}
	}

	// Checks services to remove

	for n, _ := range current {
		if _, found := expected[n]; !found {
			color.Red("\n - %s << service will be removed", n)
		}
	}

	/*
		max := math.Max(len(current), len(expected))
		for i := 0; i < max; i++ {
			if i >= len(current) {
				// New service to add
			} else if i >= len(expected) {
				// Service to remove
			} else if current[i] != expected[i] {
				// Service to remove and service to add
			} else {
			}
		}*/

	return nil
}

func safeDereference(p *string) string {
	if p == nil {
		return ""
	}
	return *p
}

func getBundleServicesSpec(bundle *bundlefile.Bundlefile, stack string) Services {
	specs := Services{}

	for name, service := range bundle.Services {
		spec := swarm.ServiceSpec{
			TaskTemplate: swarm.TaskSpec{
				ContainerSpec: swarm.ContainerSpec{
					Image:   service.Image,
					Labels:  service.Labels,
					Command: service.Command,
					Args:    service.Args,
					Env:     service.Env,
					Dir:     safeDereference(service.WorkingDir),
					User:    safeDereference(service.User),
				},
			},
			Mode: swarm.ServiceMode{
				Replicated: &swarm.ReplicatedService{
					Replicas: &Replica1,
				},
			},
		}
		spec.Labels = map[string]string{"com.docker.stack.namespace": stack}
		spec.Name = fmt.Sprintf("%s_%s", stack, name)

		// Populate ports
		ports := []swarm.PortConfig{}
		for _, port := range service.Ports {
			p := swarm.PortConfig{
				TargetPort: port.Port,
				Protocol:   swarm.PortConfigProtocol(port.Protocol),
			}

			ports = append(ports, p)
		}
		if len(ports) > 0 {
			spec.EndpointSpec = &swarm.EndpointSpec{Ports: ports}
		}

		service := swarm.Service{}
		service.ID = spec.Name
		service.Spec = spec

		specs[spec.Name] = service
	}
	return specs
}

func getSwarmServicesSpecForStack(services []swarm.Service, stack string) Services {
	specs := Services{}

	for _, service := range services {
		if service.Spec.Labels["com.docker.stack.namespace"] == stack {
			specs[service.Spec.Name] = service
		}
	}

	return specs
}
<commit_msg>Add network population<commit_after>package main

import (
	"fmt"
	"io"
	"net/url"
	"os"

	"golang.org/x/net/context"

	"github.com/docker/docker/api/client/bundlefile"
	"github.com/docker/engine-api/client"
	"github.com/docker/engine-api/types"
	"github.com/docker/engine-api/types/swarm"
	"github.com/fatih/color"
	"github.com/urfave/cli"
)

var Replica1 uint64 = 1

type Services map[string]swarm.Service

func plan(c *cli.Context) error {
	stackName := c.Args().Get(0)

	if stackName == "" {
		return cli.NewExitError("Need to specify a stack name", 1)
	}
	dabLocation := c.String("dab")

	if dabLocation == "" {
		// Assume it is called as the stack name
		dabLocation = fmt.Sprintf("%s.dab", stackName)
	}

	var dabReader io.Reader
	if u, e := url.Parse(dabLocation); e == nil && u.IsAbs() {
		// DAB file seems to be remote, try to download it first
		return cli.NewExitError("Not implemented", 2)
	} else {
		if dabFile, err := os.Open(dabLocation); err != nil {
			return cli.NewExitError(err.Error(), 3)
		} else {
			dabReader = dabFile
		}
	}

	bundle, bundleErr := bundlefile.LoadFile(dabReader)
	if bundleErr != nil {
		return cli.NewExitError(bundleErr.Error(), 3)
	}

	swarm, swarmErr := client.NewEnvClient()
	if swarmErr != nil {
		return cli.NewExitError(swarmErr.Error(), 3)
	}

	services, servicesErr := swarm.ServiceList(context.Background(), types.ServiceListOptions{})
	if servicesErr != nil {
		return cli.NewExitError(servicesErr.Error(), 3)
	}

	expected := getBundleServicesSpec(bundle, stackName)
	current := getSwarmServicesSpecForStack(services, stackName)

	for n, es := range expected {
		if cs, found := current[n]; !found {
			color.Green("\n + %s << service will be added", n)
			// New service to add
		} else {
			color.Cyan("\n%s\n", es.Spec.Name)
			PrintServiceSpecDiff(cs.Spec, es.Spec)
		}
	}

	// Checks services to remove

	for n, _ := range current {
		if _, found := expected[n]; !found {
			color.Red("\n - %s << service will be removed", n)
		}
	}

	/*
		max := math.Max(len(current), len(expected))
		for i := 0; i < max; i++ {
			if i >= len(current) {
				// New service to add
			} else if i >= len(expected) {
				// Service to remove
			} else if current[i] != expected[i] {
				// Service to remove and service to add
			} else {
			}
		}*/

	return nil
}

func safeDereference(p *string) string {
	if p == nil {
		return ""
	}
	return *p
}

func getBundleServicesSpec(bundle *bundlefile.Bundlefile, stack string) Services {
	specs := Services{}

	for name, service := range bundle.Services {
		spec := swarm.ServiceSpec{
			TaskTemplate: swarm.TaskSpec{
				ContainerSpec: swarm.ContainerSpec{
					Image:   service.Image,
					Labels:  service.Labels,
					Command: service.Command,
					Args:    service.Args,
					Env:     service.Env,
					Dir:     safeDereference(service.WorkingDir),
					User:    safeDereference(service.User),
				},
			},
			Mode: swarm.ServiceMode{
				Replicated: &swarm.ReplicatedService{
					Replicas: &Replica1,
				},
			},
			Networks: convertNetworks(service.Networks, stack, name),
		}
		spec.Labels = map[string]string{"com.docker.stack.namespace": stack}
		spec.Name = fmt.Sprintf("%s_%s", stack, name)

		// Populate ports
		ports := []swarm.PortConfig{}
		for _, port := range service.Ports {
			p := swarm.PortConfig{
				TargetPort: port.Port,
				Protocol:   swarm.PortConfigProtocol(port.Protocol),
			}

			ports = append(ports, p)
		}
		if len(ports) > 0 {
			// Hardcode resolution mode to VIP as it's the default with dab
			spec.EndpointSpec = &swarm.EndpointSpec{Ports: ports, Mode: swarm.ResolutionMode("vip")}
		}

		service := swarm.Service{}
		service.ID = spec.Name
		service.Spec = spec

		specs[spec.Name] = service
	}
	return specs
}

func convertNetworks(networks []string, namespace string, name string) []swarm.NetworkAttachmentConfig {
	nets := []swarm.NetworkAttachmentConfig{}
	for _, network := range networks {
		nets = append(nets, swarm.NetworkAttachmentConfig{
			Target:  namespace + "_" + network,
			Aliases: []string{name},
		})
	}
	return nets
}

func getSwarmServicesSpecForStack(services []swarm.Service, stack string) Services {
	specs := Services{}

	for _, service := range services {
		if service.Spec.Labels["com.docker.stack.namespace"] == stack {
			specs[service.Spec.Name] = service
		}
	}

	return specs
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"

	"github.com/pachyderm/pachyderm/src/client"
	healthclient "github.com/pachyderm/pachyderm/src/client/health"
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/discovery"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pkg/shard"
	"github.com/pachyderm/pachyderm/src/client/pkg/uuid"
	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/health"
	pfs_persist "github.com/pachyderm/pachyderm/src/server/pfs/db"
	"github.com/pachyderm/pachyderm/src/server/pfs/drive"
	pfs_server "github.com/pachyderm/pachyderm/src/server/pfs/server"
	cache_pb "github.com/pachyderm/pachyderm/src/server/pkg/cache/groupcachepb"
	cache_server "github.com/pachyderm/pachyderm/src/server/pkg/cache/server"
	"github.com/pachyderm/pachyderm/src/server/pkg/metrics"
	"github.com/pachyderm/pachyderm/src/server/pkg/netutil"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps"
	"github.com/pachyderm/pachyderm/src/server/pps/persist"
	persist_server "github.com/pachyderm/pachyderm/src/server/pps/persist/server"
	pps_server "github.com/pachyderm/pachyderm/src/server/pps/server"

	flag "github.com/spf13/pflag"
	"go.pedge.io/env"
	"go.pedge.io/lion"
	"go.pedge.io/lion/proto"
	"go.pedge.io/proto/server"
	"google.golang.org/grpc"
	"k8s.io/kubernetes/pkg/api"
	kube_client "k8s.io/kubernetes/pkg/client/restclient"
	kube "k8s.io/kubernetes/pkg/client/unversioned"
)

var readinessCheck bool

func init() {
	flag.BoolVar(&readinessCheck, "readiness-check", false, "Set to true when checking if local pod is ready")
	flag.Parse()
}

type appEnv struct {
	Port               uint16 `env:"PORT,default=650"`
	NumShards          uint64 `env:"NUM_SHARDS,default=32"`
	StorageRoot        string `env:"PACH_ROOT,required"`
	StorageBackend     string `env:"STORAGE_BACKEND,default="`
	DatabaseAddress    string `env:"RETHINK_PORT_28015_TCP_ADDR,required"`
	PPSDatabaseName    string `env:"DATABASE_NAME,default=pachyderm_pps"`
	PFSDatabaseName    string `env:"DATABASE_NAME,default=pachyderm_pfs"`
	KubeAddress        string `env:"KUBERNETES_PORT_443_TCP_ADDR,required"`
	EtcdAddress        string `env:"ETCD_PORT_2379_TCP_ADDR,required"`
	Namespace          string `env:"NAMESPACE,default=default"`
	Metrics            bool   `env:"METRICS,default=true"`
	Init               bool   `env:"INIT,default=false"`
	BlockCacheBytes    int64  `env:"BLOCK_CACHE_BYTES,default=1073741824"` //default = 1 gigabyte
	JobShimImage       string `env:"JOB_SHIM_IMAGE,default="`
	JobImagePullPolicy string `env:"JOB_IMAGE_PULL_POLICY,default="`
}

func main() {
	env.Main(do, &appEnv{})
}

func do(appEnvObj interface{}) error {
	go func() {
		lion.Println(http.ListenAndServe(":651", nil))
	}()
	appEnv := appEnvObj.(*appEnv)
	etcdClient := getEtcdClient(appEnv)
	rethinkAddress := fmt.Sprintf("%s:28015", appEnv.DatabaseAddress)
	if appEnv.Init {
		if err := setClusterID(etcdClient); err != nil {
			return fmt.Errorf("error connecting to etcd, if this error persists it likely indicates that kubernetes services are not working correctly. See https://github.com/pachyderm/pachyderm/blob/master/SETUP.md#pachd-or-pachd-init-crash-loop-with-error-connecting-to-etcd for more info")
		}
		if err := persist_server.InitDBs(rethinkAddress, appEnv.PPSDatabaseName); err != nil {
			return err
		}
		return pfs_persist.InitDB(rethinkAddress, appEnv.PFSDatabaseName)
	}
	if readinessCheck {
		c, err := client.NewFromAddress("127.0.0.1:650")
		if err != nil {
			return err
		}

		// We want to use a PPS API instead of a PFS API because PFS APIs
		// typically talk to every node, but the point of the readiness probe
		// is that it checks to see if this particular node is functioning,
		// and removing it from the service if it's not.  So if we use a PFS
		// API such as ListRepo for readiness probe, then the failure of any
		// node will result in the failures of all readiness probes, causing
		// all nodes to be removed from the pachd service.
		_, err = c.ListPipeline()
		if err != nil {
			return err
		}
		os.Exit(0)

		return nil
	}

	clusterID, err := getClusterID(etcdClient)
	if err != nil {
		return err
	}
	kubeClient, err := getKubeClient(appEnv)
	if err != nil {
		return err
	}
	if appEnv.Metrics {
		err := metrics.InitializeReporter(clusterID, kubeClient, rethinkAddress, appEnv.PFSDatabaseName, appEnv.PPSDatabaseName)
		if err != nil {
			return err
		}
	}
	rethinkAPIServer, err := getRethinkAPIServer(appEnv)
	if err != nil {
		return err
	}
	address, err := netutil.ExternalIP()
	if err != nil {
		return err
	}
	address = fmt.Sprintf("%s:%d", address, appEnv.Port)
	sharder := shard.NewSharder(
		etcdClient,
		appEnv.NumShards,
		appEnv.Namespace,
	)
	go func() {
		if err := sharder.AssignRoles(address, nil); err != nil {
			protolion.Printf("error from sharder.AssignRoles: %s", sanitizeErr(err))
		}
	}()
	driver, err := getPFSDriver(address, appEnv)
	//	driver, err := drive.NewDriver(address)
	if err != nil {
		return err
	}
	router := shard.NewRouter(
		sharder,
		grpcutil.NewDialer(
			grpc.WithInsecure(),
		),
		address,
	)
	cacheServer := cache_server.NewCacheServer(router, appEnv.NumShards)
	go func() {
		if err := sharder.RegisterFrontends(nil, address, []shard.Frontend{cacheServer}); err != nil {
			protolion.Printf("error from sharder.RegisterFrontend %s", sanitizeErr(err))
		}
	}()
	apiServer := pfs_server.NewAPIServer(driver)
	ppsAPIServer := pps_server.NewAPIServer(
		ppsserver.NewHasher(appEnv.NumShards, appEnv.NumShards),
		address,
		kubeClient,
		getNamespace(),
		appEnv.JobShimImage,
		appEnv.JobImagePullPolicy,
	)
	go func() {
		if err := sharder.Register(nil, address, []shard.Server{ppsAPIServer, cacheServer}); err != nil {
			protolion.Printf("error from sharder.Register %s", sanitizeErr(err))
		}
	}()
	blockAPIServer, err := pfs_server.NewBlockAPIServer(appEnv.StorageRoot, appEnv.BlockCacheBytes, appEnv.StorageBackend)
	if err != nil {
		return err
	}
	healthServer := health.NewHealthServer()
	return protoserver.Serve(
		func(s *grpc.Server) {
			pfsclient.RegisterAPIServer(s, apiServer)
			pfsclient.RegisterBlockAPIServer(s, blockAPIServer)
			ppsclient.RegisterAPIServer(s, ppsAPIServer)
			ppsserver.RegisterInternalJobAPIServer(s, ppsAPIServer)
			persist.RegisterAPIServer(s, rethinkAPIServer)
			cache_pb.RegisterGroupCacheServer(s, cacheServer)
			healthclient.RegisterHealthServer(s, healthServer)
		},
		protoserver.ServeOptions{
			Version: version.Version,
		},
		protoserver.ServeEnv{
			GRPCPort: appEnv.Port,
		},
	)
}

func getEtcdClient(env *appEnv) discovery.Client {
	return discovery.NewEtcdClient(fmt.Sprintf("http://%s:2379", env.EtcdAddress))
}

const clusterIDKey = "cluster-id"

func setClusterID(client discovery.Client) error {
	return client.Set(clusterIDKey, uuid.NewWithoutDashes(), 0)
}

func getClusterID(client discovery.Client) (string, error) {
	id, err := client.Get(clusterIDKey)
	if err != nil {
		return "", err
	}
	if id == "" {
		return "", fmt.Errorf("clusterID not yet set")
	}
	return id, nil
}

func getKubeClient(env *appEnv) (*kube.Client, error) {
	kubeClient, err := kube.NewInCluster()
	if err != nil {
		protolion.Errorf("falling back to insecure kube client due to error from NewInCluster: %s", sanitizeErr(err))
	} else {
		return kubeClient, err
	}
	config := &kube_client.Config{
		Host:     fmt.Sprintf("%s:443", env.KubeAddress),
		Insecure: true,
	}
	return kube.New(config)
}

func getPFSDriver(address string, env *appEnv) (drive.Driver, error) {
	rethinkAddress := fmt.Sprintf("%s:28015", env.DatabaseAddress)
	return pfs_persist.NewDriver(address, rethinkAddress, env.PFSDatabaseName)
}

func getRethinkAPIServer(env *appEnv) (persist.APIServer, error) {
	if err := persist_server.CheckDBs(fmt.Sprintf("%s:28015", env.DatabaseAddress), env.PPSDatabaseName); err != nil {
		return nil, err
	}
	return persist_server.NewRethinkAPIServer(fmt.Sprintf("%s:28015", env.DatabaseAddress), env.PPSDatabaseName)
}

// getNamespace returns the kubernetes namespace that this pachd pod runs in
func getNamespace() string {
	namespace := os.Getenv("PACHD_POD_NAMESPACE")
	if namespace != "" {
		return namespace
	}
	return api.NamespaceDefault
}

func sanitizeErr(err error) error {
	if err == nil {
		return nil
	}

	return errors.New(grpc.ErrorDesc(err))
}
<commit_msg>Log error, but dont quit if metrics initialization fails<commit_after>package main

import (
	"errors"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"

	"github.com/pachyderm/pachyderm/src/client"
	healthclient "github.com/pachyderm/pachyderm/src/client/health"
	pfsclient "github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/discovery"
	"github.com/pachyderm/pachyderm/src/client/pkg/grpcutil"
	"github.com/pachyderm/pachyderm/src/client/pkg/shard"
	"github.com/pachyderm/pachyderm/src/client/pkg/uuid"
	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/client/version"
	"github.com/pachyderm/pachyderm/src/server/health"
	pfs_persist "github.com/pachyderm/pachyderm/src/server/pfs/db"
	"github.com/pachyderm/pachyderm/src/server/pfs/drive"
	pfs_server "github.com/pachyderm/pachyderm/src/server/pfs/server"
	cache_pb "github.com/pachyderm/pachyderm/src/server/pkg/cache/groupcachepb"
	cache_server "github.com/pachyderm/pachyderm/src/server/pkg/cache/server"
	"github.com/pachyderm/pachyderm/src/server/pkg/metrics"
	"github.com/pachyderm/pachyderm/src/server/pkg/netutil"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps"
	"github.com/pachyderm/pachyderm/src/server/pps/persist"
	persist_server "github.com/pachyderm/pachyderm/src/server/pps/persist/server"
	pps_server "github.com/pachyderm/pachyderm/src/server/pps/server"

	flag "github.com/spf13/pflag"
	"go.pedge.io/env"
	"go.pedge.io/lion"
	"go.pedge.io/lion/proto"
	"go.pedge.io/proto/server"
	"google.golang.org/grpc"
	"k8s.io/kubernetes/pkg/api"
	kube_client "k8s.io/kubernetes/pkg/client/restclient"
	kube "k8s.io/kubernetes/pkg/client/unversioned"
)

var readinessCheck bool

func init() {
	flag.BoolVar(&readinessCheck, "readiness-check", false, "Set to true when checking if local pod is ready")
	flag.Parse()
}

type appEnv struct {
	Port               uint16 `env:"PORT,default=650"`
	NumShards          uint64 `env:"NUM_SHARDS,default=32"`
	StorageRoot        string `env:"PACH_ROOT,required"`
	StorageBackend     string `env:"STORAGE_BACKEND,default="`
	DatabaseAddress    string `env:"RETHINK_PORT_28015_TCP_ADDR,required"`
	PPSDatabaseName    string `env:"DATABASE_NAME,default=pachyderm_pps"`
	PFSDatabaseName    string `env:"DATABASE_NAME,default=pachyderm_pfs"`
	KubeAddress        string `env:"KUBERNETES_PORT_443_TCP_ADDR,required"`
	EtcdAddress        string `env:"ETCD_PORT_2379_TCP_ADDR,required"`
	Namespace          string `env:"NAMESPACE,default=default"`
	Metrics            bool   `env:"METRICS,default=true"`
	Init               bool   `env:"INIT,default=false"`
	BlockCacheBytes    int64  `env:"BLOCK_CACHE_BYTES,default=1073741824"` //default = 1 gigabyte
	JobShimImage       string `env:"JOB_SHIM_IMAGE,default="`
	JobImagePullPolicy string `env:"JOB_IMAGE_PULL_POLICY,default="`
}

func main() {
	env.Main(do, &appEnv{})
}

func do(appEnvObj interface{}) error {
	go func() {
		lion.Println(http.ListenAndServe(":651", nil))
	}()
	appEnv := appEnvObj.(*appEnv)
	etcdClient := getEtcdClient(appEnv)
	rethinkAddress := fmt.Sprintf("%s:28015", appEnv.DatabaseAddress)
	if appEnv.Init {
		if err := setClusterID(etcdClient); err != nil {
			return fmt.Errorf("error connecting to etcd, if this error persists it likely indicates that kubernetes services are not working correctly. See https://github.com/pachyderm/pachyderm/blob/master/SETUP.md#pachd-or-pachd-init-crash-loop-with-error-connecting-to-etcd for more info")
		}
		if err := persist_server.InitDBs(rethinkAddress, appEnv.PPSDatabaseName); err != nil {
			return err
		}
		return pfs_persist.InitDB(rethinkAddress, appEnv.PFSDatabaseName)
	}
	if readinessCheck {
		c, err := client.NewFromAddress("127.0.0.1:650")
		if err != nil {
			return err
		}

		// We want to use a PPS API instead of a PFS API because PFS APIs
		// typically talk to every node, but the point of the readiness probe
		// is that it checks to see if this particular node is functioning,
		// and removing it from the service if it's not.  So if we use a PFS
		// API such as ListRepo for readiness probe, then the failure of any
		// node will result in the failures of all readiness probes, causing
		// all nodes to be removed from the pachd service.
		_, err = c.ListPipeline()
		if err != nil {
			return err
		}

		os.Exit(0)

		return nil
	}

	clusterID, err := getClusterID(etcdClient)
	if err != nil {
		return err
	}
	kubeClient, err := getKubeClient(appEnv)
	if err != nil {
		return err
	}
	if appEnv.Metrics {
		err := metrics.InitializeReporter(clusterID, kubeClient, rethinkAddress, appEnv.PFSDatabaseName, appEnv.PPSDatabaseName)
		protolion.Printf("error initializing metrics reporting: %s", err)
	}
	rethinkAPIServer, err := getRethinkAPIServer(appEnv)
	if err != nil {
		return err
	}
	address, err := netutil.ExternalIP()
	if err != nil {
		return err
	}
	address = fmt.Sprintf("%s:%d", address, appEnv.Port)
	sharder := shard.NewSharder(
		etcdClient,
		appEnv.NumShards,
		appEnv.Namespace,
	)
	go func() {
		if err := sharder.AssignRoles(address, nil); err != nil {
			protolion.Printf("error from sharder.AssignRoles: %s", sanitizeErr(err))
		}
	}()
	driver, err := getPFSDriver(address, appEnv)
	//	driver, err := drive.NewDriver(address)
	if err != nil {
		return err
	}
	router := shard.NewRouter(
		sharder,
		grpcutil.NewDialer(
			grpc.WithInsecure(),
		),
		address,
	)
	cacheServer := cache_server.NewCacheServer(router, appEnv.NumShards)
	go func() {
		if err := sharder.RegisterFrontends(nil, address, []shard.Frontend{cacheServer}); err != nil {
			protolion.Printf("error from sharder.RegisterFrontend %s", sanitizeErr(err))
		}
	}()
	apiServer := pfs_server.NewAPIServer(driver)
	ppsAPIServer := pps_server.NewAPIServer(
		ppsserver.NewHasher(appEnv.NumShards, appEnv.NumShards),
		address,
		kubeClient,
		getNamespace(),
		appEnv.JobShimImage,
		appEnv.JobImagePullPolicy,
	)
	go func() {
		if err := sharder.Register(nil, address, []shard.Server{ppsAPIServer, cacheServer}); err != nil {
			protolion.Printf("error from sharder.Register %s", sanitizeErr(err))
		}
	}()
	blockAPIServer, err := pfs_server.NewBlockAPIServer(appEnv.StorageRoot, appEnv.BlockCacheBytes, appEnv.StorageBackend)
	if err != nil {
		return err
	}
	healthServer := health.NewHealthServer()
	return protoserver.Serve(
		func(s *grpc.Server) {
			pfsclient.RegisterAPIServer(s, apiServer)
			pfsclient.RegisterBlockAPIServer(s, blockAPIServer)
			ppsclient.RegisterAPIServer(s, ppsAPIServer)
			ppsserver.RegisterInternalJobAPIServer(s, ppsAPIServer)
			persist.RegisterAPIServer(s, rethinkAPIServer)
			cache_pb.RegisterGroupCacheServer(s, cacheServer)
			healthclient.RegisterHealthServer(s, healthServer)
		},
		protoserver.ServeOptions{
			Version: version.Version,
		},
		protoserver.ServeEnv{
			GRPCPort: appEnv.Port,
		},
	)
}

func getEtcdClient(env *appEnv) discovery.Client {
	return discovery.NewEtcdClient(fmt.Sprintf("http://%s:2379", env.EtcdAddress))
}

const clusterIDKey = "cluster-id"

func setClusterID(client discovery.Client) error {
	return client.Set(clusterIDKey, uuid.NewWithoutDashes(), 0)
}

func getClusterID(client discovery.Client) (string, error) {
	id, err := client.Get(clusterIDKey)
	if err != nil {
		return "", err
	}
	if id == "" {
		return "", fmt.Errorf("clusterID not yet set")
	}
	return id, nil
}

func getKubeClient(env *appEnv) (*kube.Client, error) {
	kubeClient, err := kube.NewInCluster()
	if err != nil {
		protolion.Errorf("falling back to insecure kube client due to error from NewInCluster: %s", sanitizeErr(err))
	} else {
		return kubeClient, err
	}
	config := &kube_client.Config{
		Host:     fmt.Sprintf("%s:443", env.KubeAddress),
		Insecure: true,
	}
	return kube.New(config)
}

func getPFSDriver(address string, env *appEnv) (drive.Driver, error) {
	rethinkAddress := fmt.Sprintf("%s:28015", env.DatabaseAddress)
	return pfs_persist.NewDriver(address, rethinkAddress, env.PFSDatabaseName)
}

func getRethinkAPIServer(env *appEnv) (persist.APIServer, error) {
	if err := persist_server.CheckDBs(fmt.Sprintf("%s:28015", env.DatabaseAddress), env.PPSDatabaseName); err != nil {
		return nil, err
	}
	return persist_server.NewRethinkAPIServer(fmt.Sprintf("%s:28015", env.DatabaseAddress), env.PPSDatabaseName)
}

// getNamespace returns the kubernetes namespace that this pachd pod runs in
func getNamespace() string {
	namespace := os.Getenv("PACHD_POD_NAMESPACE")
	if namespace != "" {
		return namespace
	}
	return api.NamespaceDefault
}

func sanitizeErr(err error) error {
	if err == nil {
		return nil
	}

	return errors.New(grpc.ErrorDesc(err))
}
<|endoftext|>


<commit_before>package gosockjs

import (
	"bytes"
	"code.google.com/p/gorilla/mux"
	"errors"
	"io"
	"net/http"
	"sync"
)

func xhrProlog(w http.ResponseWriter, req *http.Request) bool {
	w.Header().Set("Access-Control-Allow-Credentials", "true")
	h := req.Header.Get("Access-Control-Request-Headers")
	if h != "" {
		w.Header().Set("Access-Control-Allow-Headers", h)
	}
	if req.Method == "OPTIONS" {
		writeCacheAndExpires(w, req)
		writeOptionsAccess(w, req, "POST")
		w.WriteHeader(http.StatusNoContent)
		return true
	}
	writeCorsHeaders(w, req)
	writeNoCache(w, req)
	return false
}

type xhrTransport struct {
	receiver *xhrReceiver
	s        *session
	lock     sync.RWMutex
}

func (t *xhrTransport) writeFrame(w io.Writer, frame []byte) error {
	_, err := w.Write(append(frame, '\n'))
	return err
}

func (t *xhrTransport) sendFrame(frame []byte) error {
	t.lock.RLock()
	defer t.lock.RUnlock()
	if t.receiver != nil {
		return t.writeFrame(t.receiver, frame)
	}
	return errors.New("No receiver")
}

func (t *xhrTransport) closeTransport() {
	t.lock.RLock()
	defer t.lock.RUnlock()
	if t.receiver != nil {
		t.receiver.Close()
	}
}

func (t *xhrTransport) setReceiver(r *xhrReceiver) error {
	t.lock.Lock()
	if t.receiver != nil {
		defer t.lock.Unlock()
		// Nyet.
		t.writeFrame(r, closeFrame(2010, "Another connection still open"))
		return errors.New("Another connection still open")
	}
	t.receiver = r
	r.t = t
	t.lock.Unlock()
	t.s.newReceiver()
	return nil
}

func (t *xhrTransport) clearReceiver() {
	t.lock.Lock()
	t.receiver = nil
	defer t.lock.Unlock()
	t.s.disconnectReceiver()
}

type xhrReceiver struct {
	t         *xhrTransport
	w         io.WriteCloser
	byteCount chan int
	nwritten  int
	opts      xhrOptions
}

func (r *xhrReceiver) Write(data []byte) (int, error) {
	if r.nwritten >= r.opts.maxBytes() {
		return 0, errors.New("Capacity exceeded")
	}
	if len(data) == 0 {
		return 0, nil
	}
	n, err := r.w.Write(data)
	if r.byteCount != nil && n > 0 {
		r.byteCount <- n
		r.nwritten += n
	}
	return n, err
}

func (r *xhrReceiver) Close() error {
	return r.w.Close()
}

// differentiate between XhrPolling and XhrStreaming
type xhrOptions interface {
	maxBytes() int
	writePrelude(w io.Writer) error
	streaming() bool
}

type xhrPollingOptions struct {
	r *Router
}

func (o xhrPollingOptions) maxBytes() int {
	return 1
}

func (o xhrPollingOptions) writePrelude(io.Writer) error {
	return nil
}

func (o xhrPollingOptions) streaming() bool {
	return false
}

type xhrStreamingOptions struct {
	r *Router
}

func (o xhrStreamingOptions) maxBytes() int {
	return 4096 // r.maxBytes
}

func (o xhrStreamingOptions) writePrelude(w io.Writer) error {
	prelude := make([]byte, 2049)
	for i, _ := range prelude {
		prelude[i] = 'h'
	}
	prelude[2048] = '\n'
	_, err := w.Write(prelude)
	return err
}

func (o xhrStreamingOptions) streaming() bool {
	return true
}

// The handlers
func xhrHandlerBase(opts xhrOptions, r *Router, w http.ResponseWriter, req *http.Request) {
	if xhrProlog(w, req) {
		return
	}
	w.Header().Set("Content-type", "application/javascript; charset=UTF-8")
	// For CORS, if the server sent Access-Control-Request-Headers, we
	// echo it back.
	acrh := req.Header.Get("Access-Control-Request-Headers")
	if acrh != "" {
		w.Header().Set("Access-Control-Allow-Headers", acrh)
	}
	sessionId := mux.Vars(req)["sessionid"]

	w.WriteHeader(http.StatusOK)
	opts.writePrelude(w)
	hijackAndContinue(w, func(w io.WriteCloser, done chan struct{}) {
		defer w.Close()
		var trans *xhrTransport
		// Find the session
		s, _ := r.GetOrCreateSession(sessionId)
		s.sessionLock.Lock()
		// TODO: encapsulate this logic
		var sessionUnlocked bool
		defer func() {
			if !sessionUnlocked {
				s.sessionLock.Unlock()
			}
		}()
		if s.trans != nil {
			if s.closed {
				s.trans.writeFrame(w, closeFrame(3000, "Go away!"))
				return
			}
			var ok bool
			trans, ok = s.trans.(*xhrTransport)
			if !ok {
				s.trans.writeFrame(w, closeFrame(1001, "Another kind of connection is using this session"))
				return
			}
		} else {
			trans = new(xhrTransport)
			s.trans = trans
			trans.s = s
			trans.writeFrame(w, openFrame())
			conn := &Conn{s}
			go r.handler(conn)
			if !opts.streaming() {
				return
			}
		}
		s.sessionLock.Unlock()
		sessionUnlocked = true
		byteCount := make(chan int)
		var leavingVoluntarily bool
		go func() {
			nwritten := 0
			for {
				select {
				case nb := <-byteCount:
					nwritten += nb
					if nwritten >= opts.maxBytes() {
						leavingVoluntarily = true
						w.Close()
						return
					}
				case <-done:
					w.Close()
					return
				}
			}
		}()
		err := trans.setReceiver(&xhrReceiver{trans, w, byteCount, 0, opts})
		if err != nil {
			return
		}
		defer trans.clearReceiver()
		// The session may already have closed from underneath us!
		// If so, die now
		if s.closed {
			return
		}
		<-done
		// If the session isn't closed, and we're not closing voluntarily, then
		// assume the client closed us and close the session.
		if !leavingVoluntarily && !s.closed {
			trans.clearReceiver()
			s.Close()
			r.RemoveSession(sessionId, s)
		}
	})
}

func xhrSendHandler(r *Router, w http.ResponseWriter, req *http.Request) {
	if xhrProlog(w, req) {
		return
	}
	w.Header().Set("Content-type", "text/plain; charset=UTF-8")
	sessionId := mux.Vars(req)["sessionid"]
	// Find the session
	s := r.GetSession(sessionId)
	if s == nil {
		http.NotFoundHandler().ServeHTTP(w, req)
		return
	}
	// Synchronization? What if an xhr request is still creating this?
	buf := bytes.NewBuffer(nil)
	io.Copy(buf, req.Body)
	req.Body.Close()
	if buf.Len() == 0 {
		http.Error(w, EmptyPayload.Error(), http.StatusInternalServerError)
		return
	}
	err := s.fromClient(message(buf.Bytes()))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-length", "0")
	w.WriteHeader(http.StatusNoContent)
}

func xhrHandler(r *Router, w http.ResponseWriter, req *http.Request) {
	xhrHandlerBase(xhrPollingOptions{r}, r, w, req)
}

func xhrStreamingHandler(r *Router, w http.ResponseWriter, req *http.Request) {
	xhrHandlerBase(xhrStreamingOptions{r}, r, w, req)
}
<commit_msg>minor tweak<commit_after>package gosockjs

import (
	"bytes"
	"code.google.com/p/gorilla/mux"
	"errors"
	"io"
	"net/http"
	"sync"
)

func xhrProlog(w http.ResponseWriter, req *http.Request) bool {
	w.Header().Set("Access-Control-Allow-Credentials", "true")
	h := req.Header.Get("Access-Control-Request-Headers")
	if h != "" {
		w.Header().Set("Access-Control-Allow-Headers", h)
	}
	if req.Method == "OPTIONS" {
		writeCacheAndExpires(w, req)
		writeOptionsAccess(w, req, "POST")
		w.WriteHeader(http.StatusNoContent)
		return true
	}
	writeCorsHeaders(w, req)
	writeNoCache(w, req)
	return false
}

type xhrTransport struct {
	receiver *xhrReceiver
	s        *session
	lock     sync.RWMutex
}

func (t *xhrTransport) writeFrame(w io.Writer, frame []byte) error {
	_, err := w.Write(append(frame, '\n'))
	return err
}

func (t *xhrTransport) sendFrame(frame []byte) error {
	t.lock.RLock()
	defer t.lock.RUnlock()
	if t.receiver != nil {
		return t.writeFrame(t.receiver, frame)
	}
	return errors.New("No receiver")
}

func (t *xhrTransport) closeTransport() {
	t.lock.RLock()
	defer t.lock.RUnlock()
	if t.receiver != nil {
		t.receiver.Close()
	}
}

func (t *xhrTransport) setReceiver(r *xhrReceiver) error {
	t.lock.Lock()
	if t.receiver != nil {
		defer t.lock.Unlock()
		// Nyet.
		t.writeFrame(r, closeFrame(2010, "Another connection still open"))
		return errors.New("Another connection still open")
	}
	t.receiver = r
	r.t = t
	t.lock.Unlock()
	t.s.newReceiver()
	return nil
}

func (t *xhrTransport) clearReceiver() {
	t.lock.Lock()
	t.receiver = nil
	defer t.lock.Unlock()
	t.s.disconnectReceiver()
}

type xhrReceiver struct {
	t         *xhrTransport
	w         io.WriteCloser
	byteCount chan int
	nwritten  int
	opts      xhrOptions
}

func (r *xhrReceiver) Write(data []byte) (int, error) {
	if r.nwritten >= r.opts.maxBytes() {
		return 0, errors.New("Capacity exceeded")
	}
	if len(data) == 0 {
		return 0, nil
	}
	n, err := r.w.Write(data)
	if r.byteCount != nil && n > 0 {
		r.byteCount <- n
		r.nwritten += n
	}
	return n, err
}

func (r *xhrReceiver) Close() error {
	return r.w.Close()
}

// differentiate between XhrPolling and XhrStreaming
type xhrOptions interface {
	maxBytes() int
	writePrelude(w io.Writer) error
	streaming() bool
}

type xhrPollingOptions struct {
	r *Router
}

func (o xhrPollingOptions) maxBytes() int {
	return 1
}

func (o xhrPollingOptions) writePrelude(io.Writer) error {
	return nil
}

func (o xhrPollingOptions) streaming() bool {
	return false
}

type xhrStreamingOptions struct {
	r *Router
}

func (o xhrStreamingOptions) maxBytes() int {
	return 4096 // r.maxBytes
}

func (o xhrStreamingOptions) writePrelude(w io.Writer) error {
	prelude := make([]byte, 2049)
	for i, _ := range prelude {
		prelude[i] = 'h'
	}
	prelude[2048] = '\n'
	_, err := w.Write(prelude)
	return err
}

func (o xhrStreamingOptions) streaming() bool {
	return true
}

// The handlers
func xhrHandlerBase(opts xhrOptions, r *Router, w http.ResponseWriter, req *http.Request) {
	if xhrProlog(w, req) {
		return
	}
	w.Header().Set("Content-type", "application/javascript; charset=UTF-8")
	// For CORS, if the server sent Access-Control-Request-Headers, we
	// echo it back.
	acrh := req.Header.Get("Access-Control-Request-Headers")
	if acrh != "" {
		w.Header().Set("Access-Control-Allow-Headers", acrh)
	}
	sessionId := mux.Vars(req)["sessionid"]

	w.WriteHeader(http.StatusOK)
	opts.writePrelude(w)
	hijackAndContinue(w, func(w io.WriteCloser, done chan struct{}) {
		defer w.Close()
		var trans *xhrTransport
		// Find the session
		s, _ := r.GetOrCreateSession(sessionId)
		s.sessionLock.Lock()
		// TODO: encapsulate this logic
		var sessionUnlocked bool
		defer func() {
			if !sessionUnlocked {
				s.sessionLock.Unlock()
			}
		}()
		if s.trans != nil {
			if s.closed {
				s.trans.writeFrame(w, closeFrame(3000, "Go away!"))
				return
			}
			var ok bool
			trans, ok = s.trans.(*xhrTransport)
			if !ok {
				s.trans.writeFrame(w, closeFrame(1001, "Another kind of connection is using this session"))
				return
			}
		} else {
			trans = new(xhrTransport)
			s.trans = trans
			trans.s = s
			trans.writeFrame(w, openFrame())
			conn := &Conn{s}
			go r.handler(conn)
			if !opts.streaming() {
				return
			}
		}
		s.sessionLock.Unlock()
		sessionUnlocked = true
		byteCount := make(chan int)
		var leavingVoluntarily bool
		loopDone := make(chan bool)
		go func() {
			defer close(loopDone)
			nwritten := 0
			for {
				select {
				case nb := <-byteCount:
					nwritten += nb
					if nwritten >= opts.maxBytes() {
						leavingVoluntarily = true
						w.Close()
						return
					}
				case <-done:
					w.Close()
					return
				}
			}
		}()
		err := trans.setReceiver(&xhrReceiver{trans, w, byteCount, 0, opts})
		if err != nil {
			return
		}
		defer trans.clearReceiver()
		// The session may already have closed from underneath us!
		// If so, die now
		if s.closed {
			return
		}
		<-loopDone
		// If the session isn't closed, and we're not closing voluntarily, then
		// assume the client closed us and close the session.
		if !leavingVoluntarily && !s.closed {
			trans.clearReceiver()
			s.Close()
			r.RemoveSession(sessionId, s)
		}
	})
}

func xhrSendHandler(r *Router, w http.ResponseWriter, req *http.Request) {
	if xhrProlog(w, req) {
		return
	}
	w.Header().Set("Content-type", "text/plain; charset=UTF-8")
	sessionId := mux.Vars(req)["sessionid"]
	// Find the session
	s := r.GetSession(sessionId)
	if s == nil {
		http.NotFoundHandler().ServeHTTP(w, req)
		return
	}
	// Synchronization? What if an xhr request is still creating this?
	buf := bytes.NewBuffer(nil)
	io.Copy(buf, req.Body)
	req.Body.Close()
	if buf.Len() == 0 {
		http.Error(w, EmptyPayload.Error(), http.StatusInternalServerError)
		return
	}
	err := s.fromClient(message(buf.Bytes()))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-length", "0")
	w.WriteHeader(http.StatusNoContent)
}

func xhrHandler(r *Router, w http.ResponseWriter, req *http.Request) {
	xhrHandlerBase(xhrPollingOptions{r}, r, w, req)
}

func xhrStreamingHandler(r *Router, w http.ResponseWriter, req *http.Request) {
	xhrHandlerBase(xhrStreamingOptions{r}, r, w, req)
}
<|endoftext|>


<commit_before>package hukabuk

import (
	"appengine"
	"appengine/datastore"

	"github.com/GoogleCloudPlatform/go-endpoints/endpoints"
)

type Game struct {
	// Id is just the Key Id, but we keep it around for responses.
	Id      int64    `json:"id" datastore:"-"`
	Deck    Deck     `json:"-"`
	Players []string `json:"players"`
}

func GetGame(c appengine.Context, gameId int64, game *Game) error {
	key := datastore.NewKey(c, "Game", "", gameId, nil)
	err := datastore.Get(c, key, game)
	if err == nil {
		game.Id = gameId
		return nil
	} else {
		// TODO(djh): Distinguish between a "good" error (key does not exist)
		//            and a "bad" one (request failure).
		c.Debugf("GetGame err: %v", err)
		return endpoints.NewNotFoundError("Game not found.")
	}
}

type GetGamesResponse struct {
	// H/T to: http://stackoverflow.com/a/21152548/1068170
	Games []Game `json:"games"`
}

func GetGames(c appengine.Context, u *userLocal, resp *GetGamesResponse) error {
	resp.Games = make([]Game, 0)
	return nil
}

func StartGame(c appengine.Context, u *userLocal, game *Game) error {
	key := datastore.NewIncompleteKey(c, "Game", nil)
	game.Players = append(game.Players, u.GooglePlusID)
	game.Deck = Deck{}
	game.Deck.Shuffle()

	newKey, err := datastore.Put(c, key, game)
	if err != nil {
		c.Infof("error failure: %v", err)
		return endpoints.NewBadRequestError("Creating game failed.")
	}
	game.Id = newKey.IntID()
	c.Infof("newKey.IntID(): %v, key.IntID(): %v", game.Id, key.IntID())
	return nil
}
<commit_msg>Implementing datastore query for GetGames.<commit_after>package hukabuk

import (
	"appengine"
	"appengine/datastore"

	"github.com/GoogleCloudPlatform/go-endpoints/endpoints"
)

type Game struct {
	// Id is just the Key Id, but we keep it around for responses.
	Id      int64    `json:"id" datastore:"-"`
	Deck    Deck     `json:"-"`
	Players []string `json:"players"`
}

func GetGame(c appengine.Context, gameId int64, game *Game) error {
	key := datastore.NewKey(c, "Game", "", gameId, nil)
	err := datastore.Get(c, key, game)
	if err == nil {
		game.Id = gameId
		return nil
	} else {
		// TODO(djh): Distinguish between a "good" error (key does not exist)
		//            and a "bad" one (request failure).
		c.Debugf("GetGame err: %v", err)
		return endpoints.NewNotFoundError("Game not found.")
	}
}

type GetGamesResponse struct {
	// H/T to: http://stackoverflow.com/a/21152548/1068170
	Games []Game `json:"games"`
}

func GetGames(c appengine.Context, u *userLocal, resp *GetGamesResponse) error {
	q := datastore.NewQuery("Game").Filter("Players =", u.GooglePlusID)
	keys, err := q.GetAll(c, &resp.Games)
	if err == nil {
		for i, key := range keys {
			resp.Games[i].Id = key.IntID()
		}
		return nil
	} else {
		return err
	}
}

func StartGame(c appengine.Context, u *userLocal, game *Game) error {
	key := datastore.NewIncompleteKey(c, "Game", nil)
	game.Players = append(game.Players, u.GooglePlusID)
	game.Deck = Deck{}
	game.Deck.Shuffle()

	newKey, err := datastore.Put(c, key, game)
	if err != nil {
		c.Infof("error failure: %v", err)
		return endpoints.NewBadRequestError("Creating game failed.")
	}
	// key.IntID() will not be populated.
	game.Id = newKey.IntID()
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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
//
//     https://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 (
	"testing"
)

func TestMain(t *testing.T) {
	main()
}
<commit_msg>iam/quickstart: use testutil to test main (#865)<commit_after>// Copyright 2019 Google LLC
//
// 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
//
//     https://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 (
	"testing"
	"time"

	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
)

func TestMain(t *testing.T) {
	testutil.SystemTest(t)

	r := testutil.BuildMain(t)
	if stdout, stderr, err := r.Run(nil, 10*time.Second); err != nil {
		t.Errorf("error running main: %v\n\nstdout:\n----\n%v\n----\nstderr:\n----\n%v\n----", err, string(stdout), string(stderr))
	}
	r.Cleanup()
}
<|endoftext|>


<commit_before>package linq

// Zip applies a specified function to the corresponding elements of two
// collections, producing a collection of the results.
//
// The method steps through the two input collections, applying function
// resultSelector to corresponding elements of the two collections. The method
// returns a collection of the values that are returned by resultSelector. If
// the input collections do not have the same number of elements, the method
// combines elements until it reaches the end of one of the collections. For
// example, if one collection has three elements and the other one has four, the
// result collection has only three elements.
func (q Query) Zip(
	q2 Query,
	resultSelector func(interface{}, interface{}) interface{},
) Query {

	return Query{
		Iterate: func() Iterator {
			next1 := q.Iterate()
			next2 := q2.Iterate()

			return func() (item interface{}, ok bool) {
				item1, ok1 := next1()
				item2, ok2 := next2()

				if ok1 && ok2 {
					return resultSelector(item1, item2), true
				}

				return nil, false
			}
		},
	}
}

// ZipT is the typed version of Zip.
//
//   - resultSelectorFn is of type "func(TFirst,TSecond)TResult"
//
// NOTE: Zip has better performance than ZipT.
func (q Query) ZipT(q2 Query, resultSelectorFn interface{}) Query {
	resultSelectorFunc, ok := resultSelectorFn.(func(interface{}, interface{}) interface{})
	if !ok {
		resultSelectorGenericFunc, err := newGenericFunc(
			"ZipT", "resultSelectorFn", resultSelectorFn,
			simpleParamValidator(newElemTypeSlice(new(genericType), new(genericType)), newElemTypeSlice(new(genericType))),
		)
		if err != nil {
			panic(err)
		}

		resultSelectorFunc = func(item1 interface{}, item2 interface{}) interface{} {
			return resultSelectorGenericFunc.Call(item1, item2)
		}
	}
	return Query{
		Iterate: func() Iterator {
			next1 := q.Iterate()
			next2 := q2.Iterate()

			return func() (item interface{}, ok bool) {
				item1, ok1 := next1()
				item2, ok2 := next2()

				if ok1 && ok2 {
					return resultSelectorFunc(item1, item2), true
				}

				return nil, false
			}
		},
	}
}
<commit_msg>Fixed ZipT to reuse regular method (#51)<commit_after>package linq

// Zip applies a specified function to the corresponding elements of two
// collections, producing a collection of the results.
//
// The method steps through the two input collections, applying function
// resultSelector to corresponding elements of the two collections. The method
// returns a collection of the values that are returned by resultSelector. If
// the input collections do not have the same number of elements, the method
// combines elements until it reaches the end of one of the collections. For
// example, if one collection has three elements and the other one has four, the
// result collection has only three elements.
func (q Query) Zip(
	q2 Query,
	resultSelector func(interface{}, interface{}) interface{},
) Query {

	return Query{
		Iterate: func() Iterator {
			next1 := q.Iterate()
			next2 := q2.Iterate()

			return func() (item interface{}, ok bool) {
				item1, ok1 := next1()
				item2, ok2 := next2()

				if ok1 && ok2 {
					return resultSelector(item1, item2), true
				}

				return nil, false
			}
		},
	}
}

// ZipT is the typed version of Zip.
//
//   - resultSelectorFn is of type "func(TFirst,TSecond)TResult"
//
// NOTE: Zip has better performance than ZipT.
func (q Query) ZipT(q2 Query, resultSelectorFn interface{}) Query {
	resultSelectorGenericFunc, err := newGenericFunc(
		"ZipT", "resultSelectorFn", resultSelectorFn,
		simpleParamValidator(newElemTypeSlice(new(genericType), new(genericType)), newElemTypeSlice(new(genericType))),
	)
	if err != nil {
		panic(err)
	}

	resultSelectorFunc := func(item1 interface{}, item2 interface{}) interface{} {
		return resultSelectorGenericFunc.Call(item1, item2)
	}

	return q.Zip(q2, resultSelectorFunc)
}
<|endoftext|>


<commit_before>package application

import (
	"encoding/json"
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/fiorix/freegeoip"
	"github.com/rs/cors"
	"github.com/thoas/stats"
	"github.com/tylerb/graceful"
	"net/http"
	"net/url"
	"runtime"
	"strconv"
	"time"
)

func Run(config string) error {
	runtime.GOMAXPROCS(runtime.NumCPU())

	app, err := NewFromConfigPath(config)

	if err != nil {
		return err
	}

	mux := http.NewServeMux()

	middleware := stats.New()

	mux.HandleFunc("/stats", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")

		stats := middleware.Data()

		b, _ := json.Marshal(stats)

		w.Write(b)
	})

	databasePath, err := app.Jq.String("database_path")

	if err != nil {
		databasePath = DatabaseURL
	}

	db, err := openDB(databasePath, UpdateInterval, RetryInterval)

	if err != nil {
		return err
	}

	mux.Handle("/json/", freegeoip.NewHandler(db, &freegeoip.JSONEncoder{}))

	allowedOrigins, _ := app.Jq.ArrayOfStrings("allowed_origins")
	allowedMethods, _ := app.Jq.ArrayOfStrings("allowed_methods")

	n := negroni.New(negroni.NewRecovery(), negroni.NewLogger())
	n.UseHandler(mux)
	n.Use(cors.New(cors.Options{
		AllowedOrigins: allowedOrigins,
		AllowedMethods: allowedMethods,
	}))
	graceful.Run(fmt.Sprintf(":%s", strconv.Itoa(app.Port())), 10*time.Second, n)

	return nil
}

func openDB(dsn string, updateIntvl time.Duration, maxRetryIntvl time.Duration) (db *freegeoip.DB, err error) {
	u, err := url.Parse(dsn)
	if err != nil || len(u.Scheme) == 0 {
		db, err = freegeoip.Open(dsn)
	} else {
		db, err = freegeoip.OpenURL(dsn, updateIntvl, maxRetryIntvl)
	}
	return
}
<commit_msg>Move important part to the top<commit_after>package application

import (
	"encoding/json"
	"fmt"
	"github.com/codegangsta/negroni"
	"github.com/fiorix/freegeoip"
	"github.com/rs/cors"
	"github.com/thoas/stats"
	"github.com/tylerb/graceful"
	"net/http"
	"net/url"
	"runtime"
	"strconv"
	"time"
)

func Run(config string) error {
	runtime.GOMAXPROCS(runtime.NumCPU())

	app, err := NewFromConfigPath(config)

	if err != nil {
		return err
	}

	databasePath, err := app.Jq.String("database_path")

	if err != nil {
		databasePath = DatabaseURL
	}

	db, err := openDB(databasePath, UpdateInterval, RetryInterval)

	if err != nil {
		return err
	}

	mux := http.NewServeMux()

	middleware := stats.New()

	mux.HandleFunc("/stats", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")

		stats := middleware.Data()

		b, _ := json.Marshal(stats)

		w.Write(b)
	})

	mux.Handle("/json/", freegeoip.NewHandler(db, &freegeoip.JSONEncoder{}))

	allowedOrigins, _ := app.Jq.ArrayOfStrings("allowed_origins")
	allowedMethods, _ := app.Jq.ArrayOfStrings("allowed_methods")

	n := negroni.New(negroni.NewRecovery(), negroni.NewLogger())
	n.UseHandler(mux)
	n.Use(cors.New(cors.Options{
		AllowedOrigins: allowedOrigins,
		AllowedMethods: allowedMethods,
	}))
	graceful.Run(fmt.Sprintf(":%s", strconv.Itoa(app.Port())), 10*time.Second, n)

	return nil
}

func openDB(dsn string, updateIntvl time.Duration, maxRetryIntvl time.Duration) (db *freegeoip.DB, err error) {
	u, err := url.Parse(dsn)
	if err != nil || len(u.Scheme) == 0 {
		db, err = freegeoip.Open(dsn)
	} else {
		db, err = freegeoip.OpenURL(dsn, updateIntvl, maxRetryIntvl)
	}
	return
}
<|endoftext|>


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

// Package baidubce define a set of core data structure, and implements a set of core functions.
package baidubce

import (
	"fmt"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"baidubce/util"
)

const (
	// ExpirationPeriodInSeconds 1800s is the default expiration period.
	ExpirationPeriodInSeconds = 1800
)

// Region map of baidubce
var Region = map[string]string{
	"bj": "bj.bcebos.com",
	"gz": "gz.bcebos.com",
}

// Credentials struct for baidubce.
type Credentials struct {
	AccessKeyID     string
	SecretAccessKey string
}

// NewCredentials returns an instance of type `Credentials`.
func NewCredentials(AccessKeyID, secretAccessKey string) *Credentials {
	return &Credentials{AccessKeyID, secretAccessKey}
}

// DefaultCredentials provided a default `Credentials` instance.
var DefaultCredentials = Credentials{
	os.Getenv("BAIDU_BCE_AK"),
	os.Getenv("BAIDU_BCE_SK"),
}

// Config contains options for baidubce api.
type Config struct {
	Credentials
	Endpoint   string
	APIVersion string
}

// DefaultConfig provided a default `Config` instance.
var DefaultConfig = Config{DefaultCredentials, "", "v1"}

// SignOption contains options for signature of baidubce api.
type SignOption struct {
	Timestamp                 string
	ExpirationPeriodInSeconds int
	Headers                   map[string]string
	HeadersToSign             []string
	headersToSignSpecified    bool
}

// NewSignOption is the instance factory for `SignOption`.
func NewSignOption(timestamp string, expirationPeriodInSeconds int,
	headers map[string]string, headersToSign []string) *SignOption {

	return &SignOption{timestamp, expirationPeriodInSeconds,
		headers, headersToSign, len(headersToSign) > 0}
}

func CheckSignOption(option *SignOption) *SignOption {
	if option == nil {
		return &SignOption{}
	}

	return option
}

func (option *SignOption) AddHeadersToSign(headers ...string) {
	if option.HeadersToSign == nil {
		option.HeadersToSign = []string{}
	}

	option.HeadersToSign = append(option.HeadersToSign, headers...)
}

func (option *SignOption) AddHeader(key, value string) {
	if option.Headers == nil {
		option.Headers = make(map[string]string)
	}

	option.Headers[key] = value
}

func (option *SignOption) AddHeaders(headers map[string]string) {
	if headers == nil {
		return
	}

	if option.Headers == nil {
		option.Headers = make(map[string]string)
	}

	for key, value := range headers {
		option.Headers[key] = value
	}
}

func (option *SignOption) init() {
	if option.Timestamp == "" {
		option.Timestamp = util.TimeToUTCString(time.Now())
	}

	if option.ExpirationPeriodInSeconds <= 0 {
		option.ExpirationPeriodInSeconds = ExpirationPeriodInSeconds
	}

	if option.Headers == nil {
		option.Headers = make(map[string]string, 3)
	} else {
		util.MapKeyToLower(option.Headers)
	}

	option.headersToSignSpecified = len(option.HeadersToSign) > 0
	util.SliceToLower(option.HeadersToSign)

	if !util.Contains(option.HeadersToSign, "host", true) {
		option.HeadersToSign = append(option.HeadersToSign, "host")
	}

	if util.Contains(option.HeadersToSign, "date", true) {
		if !util.MapContains(option.Headers, generateHeaderValidCompareFunc("date")) {
			option.Headers["date"] = time.Now().Format(time.RFC1123)
		} else {
			option.Headers["date"] = util.TimeStringToRFC1123(util.GetMapValue(option.Headers, "date", true))
		}
	} else if util.Contains(option.HeadersToSign, "x-bce-date", true) {
		if !util.MapContains(option.Headers, generateHeaderValidCompareFunc("x-bce-date")) {
			option.Headers["x-bce-date"] = option.Timestamp
		}
	} else {
		option.HeadersToSign = append(option.HeadersToSign, "x-bce-date")
		option.Headers["x-bce-date"] = option.Timestamp
	}
}

func (option *SignOption) signedHeadersToString() string {
	var result string
	length := len(option.HeadersToSign)
	if option.headersToSignSpecified && length > 0 {
		headers := make([]string, 0, length)
		headers = append(headers, option.HeadersToSign...)
		sort.Strings(headers)
		result = strings.Join(headers, ";")
	}

	return result
}

// GenerateAuthorization returns authorization code of baidubce api.
func GenerateAuthorization(credentials Credentials, req Request, option *SignOption) string {
	if option == nil {
		option = &SignOption{}
	}
	option.init()

	authorization := "bce-auth-v1/" + credentials.AccessKeyID
	authorization += "/" + option.Timestamp
	authorization += "/" + strconv.Itoa(option.ExpirationPeriodInSeconds)
	signature := sign(credentials, req, option)
	authorization += "/" + option.signedHeadersToString() + "/" + signature

	req.addHeader("Authorization", authorization)

	return authorization
}

// Client is the base client struct for all products of baidubce.
type Client struct {
	Config
}

func (c *Client) GetUriPath(uriPath string) string {
	if strings.Index(uriPath, "/") == 0 {
		uriPath = uriPath[1:]
	}

	if c.APIVersion != "" {
		return fmt.Sprintf("/%s/%s", c.APIVersion, uriPath)
	}

	return fmt.Sprintf("/%s", uriPath)
}

// SendRequest sends a http request to the endpoint of baidubce api.
func (c *Client) SendRequest(req *Request, option *SignOption, autoReadAllBytesFromResponseBody bool) (*Response, *Error) {
	GenerateAuthorization(c.Credentials, *req, option)
	httpClient := http.Client{}
	res, err := httpClient.Do(req.raw())

	if err != nil {
		return nil, NewError(err)
	}

	bceResponse, err := NewResponse(res, autoReadAllBytesFromResponseBody)

	if err != nil {
		return nil, NewError(err)
	}

	if res.StatusCode >= 400 {
		return bceResponse, NewErrorFromJSON(bceResponse.Body)
	}

	return bceResponse, nil
}

func generateHeaderValidCompareFunc(headerKey string) func(string, string) bool {
	return func(key, value string) bool {
		return strings.ToLower(key) == strings.ToLower(headerKey) && value != ""
	}
}

// sign returns signed signature.
func sign(credentials Credentials, req Request, option *SignOption) string {
	signingKey := getSigningKey(credentials, option)
	req.prepareHeaders(option)
	canonicalRequest := req.canonical(option)
	signature := util.HmacSha256Hex(signingKey, canonicalRequest)

	return signature
}

func getSigningKey(credentials Credentials, option *SignOption) string {
	var authStringPrefix = fmt.Sprintf("bce-auth-v1/%s", credentials.AccessKeyID)
	authStringPrefix += "/" + option.Timestamp
	authStringPrefix += "/" + strconv.Itoa(option.ExpirationPeriodInSeconds)

	return util.HmacSha256Hex(credentials.SecretAccessKey, authStringPrefix)
}
<commit_msg>fix bug of generating authorization<commit_after>/**
 * Copyright (c) 2015 Guoyao Wu, All Rights Reserved
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations under the License.
 *
 * @file core.go
 * @author guoyao
 */

// Package baidubce define a set of core data structure, and implements a set of core functions.
package baidubce

import (
	"fmt"
	"net/http"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"baidubce/util"
)

const (
	// ExpirationPeriodInSeconds 1800s is the default expiration period.
	ExpirationPeriodInSeconds = 1800
)

// Region map of baidubce
var Region = map[string]string{
	"bj": "bj.bcebos.com",
	"gz": "gz.bcebos.com",
}

// Credentials struct for baidubce.
type Credentials struct {
	AccessKeyID     string
	SecretAccessKey string
}

// NewCredentials returns an instance of type `Credentials`.
func NewCredentials(AccessKeyID, secretAccessKey string) *Credentials {
	return &Credentials{AccessKeyID, secretAccessKey}
}

// DefaultCredentials provided a default `Credentials` instance.
var DefaultCredentials = Credentials{
	os.Getenv("BAIDU_BCE_AK"),
	os.Getenv("BAIDU_BCE_SK"),
}

// Config contains options for baidubce api.
type Config struct {
	Credentials
	Endpoint   string
	APIVersion string
}

// DefaultConfig provided a default `Config` instance.
var DefaultConfig = Config{DefaultCredentials, "", "v1"}

// SignOption contains options for signature of baidubce api.
type SignOption struct {
	Timestamp                 string
	ExpirationPeriodInSeconds int
	Headers                   map[string]string
	HeadersToSign             []string
	headersToSignSpecified    bool
}

// NewSignOption is the instance factory for `SignOption`.
func NewSignOption(timestamp string, expirationPeriodInSeconds int,
	headers map[string]string, headersToSign []string) *SignOption {

	return &SignOption{timestamp, expirationPeriodInSeconds,
		headers, headersToSign, len(headersToSign) > 0}
}

func CheckSignOption(option *SignOption) *SignOption {
	if option == nil {
		return &SignOption{}
	}

	return option
}

func (option *SignOption) AddHeadersToSign(headers ...string) {
	if option.HeadersToSign == nil {
		option.HeadersToSign = []string{}
	}

	option.HeadersToSign = append(option.HeadersToSign, headers...)
}

func (option *SignOption) AddHeader(key, value string) {
	if option.Headers == nil {
		option.Headers = make(map[string]string)
	}

	option.Headers[key] = value
}

func (option *SignOption) AddHeaders(headers map[string]string) {
	if headers == nil {
		return
	}

	if option.Headers == nil {
		option.Headers = make(map[string]string)
	}

	for key, value := range headers {
		option.Headers[key] = value
	}
}

func (option *SignOption) init() {
	if option.Timestamp == "" {
		option.Timestamp = util.TimeToUTCString(time.Now())
	}

	if option.ExpirationPeriodInSeconds <= 0 {
		option.ExpirationPeriodInSeconds = ExpirationPeriodInSeconds
	}

	if option.Headers == nil {
		option.Headers = make(map[string]string, 3)
	} else {
		util.MapKeyToLower(option.Headers)
	}

	option.headersToSignSpecified = len(option.HeadersToSign) > 0
	util.SliceToLower(option.HeadersToSign)

	if !util.Contains(option.HeadersToSign, "host", true) {
		option.HeadersToSign = append(option.HeadersToSign, "host")
	}

	if !option.headersToSignSpecified {
		option.HeadersToSign = append(option.HeadersToSign, "x-bce-date")
		option.Headers["x-bce-date"] = option.Timestamp
	} else if util.Contains(option.HeadersToSign, "date", true) {
		if !util.MapContains(option.Headers, generateHeaderValidCompareFunc("date")) {
			option.Headers["date"] = time.Now().Format(time.RFC1123)
		} else {
			option.Headers["date"] = util.TimeStringToRFC1123(util.GetMapValue(option.Headers, "date", true))
		}
	} else if util.Contains(option.HeadersToSign, "x-bce-date", true) {
		if !util.MapContains(option.Headers, generateHeaderValidCompareFunc("x-bce-date")) {
			option.Headers["x-bce-date"] = option.Timestamp
		}
	}
}

func (option *SignOption) signedHeadersToString() string {
	var result string
	length := len(option.HeadersToSign)
	if option.headersToSignSpecified && length > 0 {
		headers := make([]string, 0, length)
		headers = append(headers, option.HeadersToSign...)
		sort.Strings(headers)
		result = strings.Join(headers, ";")
	}

	return result
}

// GenerateAuthorization returns authorization code of baidubce api.
func GenerateAuthorization(credentials Credentials, req Request, option *SignOption) string {
	if option == nil {
		option = &SignOption{}
	}
	option.init()

	authorization := "bce-auth-v1/" + credentials.AccessKeyID
	authorization += "/" + option.Timestamp
	authorization += "/" + strconv.Itoa(option.ExpirationPeriodInSeconds)
	signature := sign(credentials, req, option)
	authorization += "/" + option.signedHeadersToString() + "/" + signature

	req.addHeader("Authorization", authorization)

	return authorization
}

// Client is the base client struct for all products of baidubce.
type Client struct {
	Config
}

func (c *Client) GetUriPath(uriPath string) string {
	if strings.Index(uriPath, "/") == 0 {
		uriPath = uriPath[1:]
	}

	if c.APIVersion != "" {
		return fmt.Sprintf("/%s/%s", c.APIVersion, uriPath)
	}

	return fmt.Sprintf("/%s", uriPath)
}

// SendRequest sends a http request to the endpoint of baidubce api.
func (c *Client) SendRequest(req *Request, option *SignOption, autoReadAllBytesFromResponseBody bool) (*Response, *Error) {
	GenerateAuthorization(c.Credentials, *req, option)
	httpClient := http.Client{}
	res, err := httpClient.Do(req.raw())

	if err != nil {
		return nil, NewError(err)
	}

	bceResponse, err := NewResponse(res, autoReadAllBytesFromResponseBody)

	if err != nil {
		return nil, NewError(err)
	}

	if res.StatusCode >= 400 {
		return bceResponse, NewErrorFromJSON(bceResponse.Body)
	}

	return bceResponse, nil
}

func generateHeaderValidCompareFunc(headerKey string) func(string, string) bool {
	return func(key, value string) bool {
		return strings.ToLower(key) == strings.ToLower(headerKey) && value != ""
	}
}

// sign returns signed signature.
func sign(credentials Credentials, req Request, option *SignOption) string {
	signingKey := getSigningKey(credentials, option)
	req.prepareHeaders(option)
	canonicalRequest := req.canonical(option)
	signature := util.HmacSha256Hex(signingKey, canonicalRequest)

	return signature
}

func getSigningKey(credentials Credentials, option *SignOption) string {
	var authStringPrefix = fmt.Sprintf("bce-auth-v1/%s", credentials.AccessKeyID)
	authStringPrefix += "/" + option.Timestamp
	authStringPrefix += "/" + strconv.Itoa(option.ExpirationPeriodInSeconds)

	return util.HmacSha256Hex(credentials.SecretAccessKey, authStringPrefix)
}
<|endoftext|>


<commit_before>package etcd

import (
	"context"
	"testing"
	"time"

	"encoding/json"

	"github.com/coreos/etcd/client"
	"github.com/stretchr/testify/assert"
)

func TestClientStruct(t *testing.T) {
	cfg := &Config{
		Clusters: []string{"http://localhost:2379"},
	}
	// local etcd server without tls and basic auth
	cl, err := NewClient(cfg)
	assert.Nil(t, err)

	t.Run("Key", func(t *testing.T) {
		assert := assert.New(t)
		key := "/confl/test1/key"
		values := []string{
			"test1",
			"test2",
			"test3",
		}
		for _, value := range values {
			_, err := cl.client.Set(context.Background(), key, value, &client.SetOptions{})
			assert.Nil(err)
			v, err := cl.Key(key)
			assert.Nil(err)
			assert.Equal(value, v)
		}
	})

	t.Run("watchNext", func(t *testing.T) {
		assert := assert.New(t)
		key := "/confl/test1/watchnext"
		values := []string{
			"test1",
			"test2",
			"test3",
		}
		valueCh := make(chan string)
		go func() {
			for {
				resp, err := cl.watchNext(key)
				value := <-valueCh
				assert.Nil(err)
				assert.Equal(value, resp.Node.Value)
			}
		}()
		time.Sleep(time.Second)
		for _, value := range values {
			_, err := cl.client.Set(context.Background(), key, value, &client.SetOptions{})
			assert.Nil(err)
			valueCh <- value
		}
	})

	t.Run("WatchKey", func(t *testing.T) {
		assert := assert.New(t)
		type config struct {
			Name string `json:"name"`
			Age  int    `json:"age"`
		}
		key := "/confl/test1/watchkey"
		values := []config{
			config{Name: "confl1", Age: 1},
			config{Name: "confl2", Age: 2},
			config{Name: "confl3", Age: 3},
		}

		changeCh := make(chan struct{})
		valueCh := make(chan string)
		doneCh := make(chan struct{})
		go cl.WatchKey(key, changeCh)
		go func() {
			for range changeCh {
				v := <-valueCh
				value, err := cl.Key(key)
				assert.Nil(err)
				assert.Equal(v, value)
				doneCh <- struct{}{}
			}
		}()
		for _, value := range values {
			data, err := json.Marshal(value)
			assert.Nil(err)
			v := string(data)
			_, err = cl.client.Set(context.Background(), key, v, &client.SetOptions{})
			valueCh <- v
			assert.Nil(err)
			<-doneCh
		}
	})
}
<commit_msg>fix test case<commit_after>package etcd

import (
	"context"
	"testing"

	"encoding/json"

	"github.com/coreos/etcd/client"
	"github.com/stretchr/testify/assert"
)

func TestClientStruct(t *testing.T) {
	cfg := &Config{
		Clusters: []string{"http://localhost:2379"},
	}
	// local etcd server without tls and basic auth
	cl, err := NewClient(cfg)
	assert.Nil(t, err)

	t.Run("Key", func(t *testing.T) {
		assert := assert.New(t)
		key := "/confl/test1/key"
		values := []string{
			"test1",
			"test2",
			"test3",
		}
		for _, value := range values {
			_, err := cl.client.Set(context.Background(), key, value, &client.SetOptions{})
			assert.Nil(err)
			v, err := cl.Key(key)
			assert.Nil(err)
			assert.Equal(value, v)
		}
	})

	t.Run("watchNext", func(t *testing.T) {
		assert := assert.New(t)
		key := "/confl/test1/watchnext"
		values := []string{
			"test1",
			"test2",
			"test3",
		}
		valueCh := make(chan string)
		respCh := make(chan *client.Response)
		go func() {
			for {
				go func() {
					resp, err := cl.watchNext(key)
					if assert.Nil(err) {
						respCh <- resp
					}
				}()
				value := <-valueCh
				resp := <-respCh
				assert.Nil(err)
				assert.Equal(value, resp.Node.Value)
			}
		}()
		for _, value := range values {
			_, err := cl.client.Set(context.Background(), key, value, &client.SetOptions{})
			assert.Nil(err)
			valueCh <- value
		}
	})

	t.Run("WatchKey", func(t *testing.T) {
		assert := assert.New(t)
		type config struct {
			Name string `json:"name"`
			Age  int    `json:"age"`
		}
		key := "/confl/test1/watchkey"
		values := []config{
			config{Name: "confl1", Age: 1},
			config{Name: "confl2", Age: 2},
			config{Name: "confl3", Age: 3},
		}

		changeCh := make(chan struct{})
		valueCh := make(chan string)
		doneCh := make(chan struct{})
		go cl.WatchKey(key, changeCh)
		go func() {
			for range changeCh {
				v := <-valueCh
				value, err := cl.Key(key)
				assert.Nil(err)
				assert.Equal(v, value)
				doneCh <- struct{}{}
			}
		}()
		for _, value := range values {
			data, err := json.Marshal(value)
			assert.Nil(err)
			v := string(data)
			_, err = cl.client.Set(context.Background(), key, v, &client.SetOptions{})
			valueCh <- v
			assert.Nil(err)
			<-doneCh
		}
	})
}
<|endoftext|>