<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package leveldb

import (
	"time"

	"github.com/syndtr/goleveldb/leveldb/memdb"
	"github.com/syndtr/goleveldb/leveldb/opt"
)

func (d *DB) doWriteJournal(b *Batch) error {
	err := d.journal.journal.Append(b.encode())
	if err == nil && b.sync {
		err = d.journal.writer.Sync()
	}
	return err
}

func (d *DB) writeJournal() {
	for b := range d.jch {
		if b == nil {
			break
		}

		// write journal
		d.jack <- d.doWriteJournal(b)
	}

	close(d.jch)
	close(d.jack)
	d.ewg.Done()
}

func (d *DB) flush() (m *memdb.DB, err error) {
	s := d.s

	delayed, cwait := false, false
	for {
		v := s.version()
		mem := d.getMem()
		switch {
		case v.tLen(0) >= kL0_SlowdownWritesTrigger && !delayed:
			delayed = true
			time.Sleep(time.Millisecond)
		case mem.cur.Size() <= s.o.GetWriteBuffer():
			// still room
			return mem.cur, nil
		case mem.froze != nil:
			if cwait {
				err = d.geterr()
				if err != nil {
					return
				}
				d.cch <- cSched
			} else {
				cwait = true
				d.cch <- cWait
			}
		case v.tLen(0) >= kL0_StopWritesTrigger:
			d.cch <- cSched
		default:
			// create new memdb and journal
			m, err = d.newMem()
			if err != nil {
				return
			}

			// schedule compaction
			select {
			case d.cch <- cSched:
			default:
			}
		}
		v.release()
	}

	return
}

// Write apply the specified batch to the database.
func (d *DB) Write(b *Batch, wo *opt.WriteOptions) (err error) {
	err = d.wok()
	if err != nil || b == nil || b.len() == 0 {
		return
	}

	b.init(wo.HasFlag(opt.WFSync))

	select {
	case d.wqueue <- b:
		return <-d.wack
	case d.wlock <- struct{}{}:
	}

	merged := 0
	defer func() {
		<-d.wlock
		for i := 0; i < merged; i++ {
			d.wack <- err
		}
	}()

	mem, err := d.flush()
	if err != nil {
		return
	}

	// calculate maximum size of the batch
	m := 1 << 20
	if x := b.size(); x <= 128<<10 {
		m = x + (128 << 10)
	}

	// merge with other batch
drain:
	for b.size() <= m && !b.sync {
		select {
		case nb := <-d.wqueue:
			b.append(nb)
			merged++
		default:
			break drain
		}
	}

	// set batch first seq number relative from last seq
	b.seq = d.seq + 1

	// write journal concurrently if it is large enough
	if b.size() >= (128 << 10) {
		d.jch <- b
		b.memReplay(mem)
		err = <-d.jack
		if err != nil {
			b.revertMemReplay(mem)
			return
		}
	} else {
		err = d.doWriteJournal(b)
		if err != nil {
			return
		}
		b.memReplay(mem)
	}

	// set last seq number
	d.addSeq(uint64(b.len()))

	return
}

// Put set the database entry for "key" to "value".
func (d *DB) Put(key, value []byte, wo *opt.WriteOptions) error {
	b := new(Batch)
	b.Put(key, value)
	return d.Write(b, wo)
}

// Delete remove the database entry (if any) for "key". It is not an error
// if "key" did not exist in the database.
func (d *DB) Delete(key []byte, wo *opt.WriteOptions) error {
	b := new(Batch)
	b.Delete(key)
	return d.Write(b, wo)
}
<commit_msg>leveldb: version: Fix missing v.release() call on DB.flush()<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package leveldb

import (
	"time"

	"github.com/syndtr/goleveldb/leveldb/memdb"
	"github.com/syndtr/goleveldb/leveldb/opt"
)

func (d *DB) doWriteJournal(b *Batch) error {
	err := d.journal.journal.Append(b.encode())
	if err == nil && b.sync {
		err = d.journal.writer.Sync()
	}
	return err
}

func (d *DB) writeJournal() {
	for b := range d.jch {
		if b == nil {
			break
		}

		// write journal
		d.jack <- d.doWriteJournal(b)
	}

	close(d.jch)
	close(d.jack)
	d.ewg.Done()
}

func (d *DB) flush() (m *memdb.DB, err error) {
	s := d.s

	delayed, cwait := false, false
	for {
		v := s.version()
		mem := d.getMem()
		switch {
		case v.tLen(0) >= kL0_SlowdownWritesTrigger && !delayed:
			delayed = true
			time.Sleep(time.Millisecond)
		case mem.cur.Size() <= s.o.GetWriteBuffer():
			// still room
			v.release()
			return mem.cur, nil
		case mem.froze != nil:
			if cwait {
				err = d.geterr()
				if err != nil {
					return
				}
				d.cch <- cSched
			} else {
				cwait = true
				d.cch <- cWait
			}
		case v.tLen(0) >= kL0_StopWritesTrigger:
			d.cch <- cSched
		default:
			// create new memdb and journal
			m, err = d.newMem()
			if err != nil {
				return
			}

			// schedule compaction
			select {
			case d.cch <- cSched:
			default:
			}
		}
		v.release()
	}

	return
}

// Write apply the specified batch to the database.
func (d *DB) Write(b *Batch, wo *opt.WriteOptions) (err error) {
	err = d.wok()
	if err != nil || b == nil || b.len() == 0 {
		return
	}

	b.init(wo.HasFlag(opt.WFSync))

	select {
	case d.wqueue <- b:
		return <-d.wack
	case d.wlock <- struct{}{}:
	}

	merged := 0
	defer func() {
		<-d.wlock
		for i := 0; i < merged; i++ {
			d.wack <- err
		}
	}()

	mem, err := d.flush()
	if err != nil {
		return
	}

	// calculate maximum size of the batch
	m := 1 << 20
	if x := b.size(); x <= 128<<10 {
		m = x + (128 << 10)
	}

	// merge with other batch
drain:
	for b.size() <= m && !b.sync {
		select {
		case nb := <-d.wqueue:
			b.append(nb)
			merged++
		default:
			break drain
		}
	}

	// set batch first seq number relative from last seq
	b.seq = d.seq + 1

	// write journal concurrently if it is large enough
	if b.size() >= (128 << 10) {
		d.jch <- b
		b.memReplay(mem)
		err = <-d.jack
		if err != nil {
			b.revertMemReplay(mem)
			return
		}
	} else {
		err = d.doWriteJournal(b)
		if err != nil {
			return
		}
		b.memReplay(mem)
	}

	// set last seq number
	d.addSeq(uint64(b.len()))

	return
}

// Put set the database entry for "key" to "value".
func (d *DB) Put(key, value []byte, wo *opt.WriteOptions) error {
	b := new(Batch)
	b.Put(key, value)
	return d.Write(b, wo)
}

// Delete remove the database entry (if any) for "key". It is not an error
// if "key" did not exist in the database.
func (d *DB) Delete(key []byte, wo *opt.WriteOptions) error {
	b := new(Batch)
	b.Delete(key)
	return d.Write(b, wo)
}
<|endoftext|>


<commit_before>package lfs

import (
	"fmt"
	"github.com/cheggaaa/pb"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
)

// Uploadable describes a file that can be uploaded.
type Uploadable struct {
	OIDPath  string
	Filename string
	CB       CopyCallback
	Size     int64
}

// NewUploadable builds the Uploadable from the given information.
func NewUploadable(oid, filename string, index, totalFiles int) (*Uploadable, *WrappedError) {
	path, err := LocalMediaPath(oid)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	if err := ensureFile(filename, path); err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	fi, err := os.Stat(filename)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	cb, file, cbErr := CopyCallbackFile("push", filename, index, totalFiles)
	if cbErr != nil {
		fmt.Fprintln(os.Stderr, cbErr.Error())
	}

	if file != nil {
		defer file.Close()
	}

	return &Uploadable{path, filename, cb, fi.Size()}, nil
}

// UploadQueue provides a queue that will allow concurrent uploads.
type UploadQueue struct {
	uploadc  chan *Uploadable
	errorc   chan *WrappedError
	errors   []*WrappedError
	wg       sync.WaitGroup
	workers  int
	files    int
	finished int64
	size     int64
}

// NewUploadQueue builds an UploadQueue, allowing `workers` concurrent uploads.
func NewUploadQueue(workers, files int) *UploadQueue {
	return &UploadQueue{
		uploadc: make(chan *Uploadable, files),
		errorc:  make(chan *WrappedError),
		workers: workers,
		files:   files,
	}
}

// Upload adds an Uploadable to the upload queue.
func (q *UploadQueue) Upload(u *Uploadable) {
	q.wg.Add(1)
	q.size += u.Size
	q.uploadc <- u
}

// Process starts the upload queue and displays a progress bar.
func (q *UploadQueue) Process() {
	bar := pb.New64(q.size)
	bar.SetUnits(pb.U_BYTES)
	bar.ShowBar = false
	bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.finished, q.files))
	bar.Start()

	go func() {
		for err := range q.errorc {
			q.errors = append(q.errors, err)
		}
	}()

	for i := 0; i < q.workers; i++ {
		go func(n int) {
			for upload := range q.uploadc {
				cb := func(total, read int64, current int) error {
					bar.Add(current)
					if upload.CB != nil {
						return upload.CB(total, read, current)
					}
					return nil
				}

				err := Upload(upload.OIDPath, upload.Filename, cb)
				if err != nil {
					q.errorc <- err
				}
				f := atomic.AddInt64(&q.finished, 1)
				bar.Prefix(fmt.Sprintf("(%d of %d files) ", f, q.files))
				q.wg.Done()
			}
		}(i)
	}

	close(q.uploadc)
	q.wg.Wait()
	close(q.errorc)

	bar.Finish()
}

// Errors returns any errors encountered during uploading.
func (q *UploadQueue) Errors() []*WrappedError {
	return q.errors
}

// ensureFile makes sure that the cleanPath exists before pushing it.  If it
// does not exist, it attempts to clean it by reading the file at smudgePath.
func ensureFile(smudgePath, cleanPath string) error {
	if _, err := os.Stat(cleanPath); err == nil {
		return nil
	}

	expectedOid := filepath.Base(cleanPath)
	localPath := filepath.Join(LocalWorkingDir, smudgePath)
	file, err := os.Open(localPath)
	if err != nil {
		return err
	}

	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return err
	}

	cleaned, err := PointerClean(file, stat.Size(), nil)
	if err != nil {
		return err
	}

	cleaned.Close()

	if expectedOid != cleaned.Oid {
		return fmt.Errorf("Expected %s to have an OID of %s, got %s", smudgePath, expectedOid, cleaned.Oid)
	}

	return nil
}
<commit_msg>Syncrhonize uploads until one is successful, then allow concurrent<commit_after>package lfs

import (
	"fmt"
	"github.com/cheggaaa/pb"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
)

var (
	clientAuthorized = 0
)

// Uploadable describes a file that can be uploaded.
type Uploadable struct {
	OIDPath  string
	Filename string
	CB       CopyCallback
	Size     int64
}

// NewUploadable builds the Uploadable from the given information.
func NewUploadable(oid, filename string, index, totalFiles int) (*Uploadable, *WrappedError) {
	path, err := LocalMediaPath(oid)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	if err := ensureFile(filename, path); err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	fi, err := os.Stat(filename)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	cb, file, cbErr := CopyCallbackFile("push", filename, index, totalFiles)
	if cbErr != nil {
		fmt.Fprintln(os.Stderr, cbErr.Error())
	}

	if file != nil {
		defer file.Close()
	}

	return &Uploadable{path, filename, cb, fi.Size()}, nil
}

// UploadQueue provides a queue that will allow concurrent uploads.
type UploadQueue struct {
	uploadc  chan *Uploadable
	errorc   chan *WrappedError
	errors   []*WrappedError
	wg       sync.WaitGroup
	workers  int
	files    int
	finished int64
	size     int64
	authCond *sync.Cond
}

// NewUploadQueue builds an UploadQueue, allowing `workers` concurrent uploads.
func NewUploadQueue(workers, files int) *UploadQueue {
	return &UploadQueue{
		uploadc:  make(chan *Uploadable, files),
		errorc:   make(chan *WrappedError),
		workers:  workers,
		files:    files,
		authCond: sync.NewCond(&sync.Mutex{}),
	}
}

// Upload adds an Uploadable to the upload queue.
func (q *UploadQueue) Upload(u *Uploadable) {
	q.wg.Add(1)
	q.size += u.Size
	q.uploadc <- u
}

// Process starts the upload queue and displays a progress bar.
func (q *UploadQueue) Process() {
	bar := pb.New64(q.size)
	bar.SetUnits(pb.U_BYTES)
	bar.ShowBar = false
	bar.Prefix(fmt.Sprintf("(%d of %d files) ", q.finished, q.files))
	bar.Start()

	go func() {
		for err := range q.errorc {
			q.errors = append(q.errors, err)
		}
	}()

	workersReady := make(chan int, q.workers)

	for i := 0; i < q.workers; i++ {
		go func(n int) {
			workersReady <- 1

			for upload := range q.uploadc {
				q.authCond.L.Lock()
				if clientAuthorized == 0 {
					q.authCond.Wait()
				}
				q.authCond.L.Unlock()

				cb := func(total, read int64, current int) error {
					bar.Add(current)
					if upload.CB != nil {
						return upload.CB(total, read, current)
					}
					return nil
				}

				err := Upload(upload.OIDPath, upload.Filename, cb)
				if err != nil {
					q.errorc <- err
					// This one failed, send a Signal() to wake up the next one
					q.authCond.Signal()
				} else {
					// This one succeeded, set the auth condition to 1 and Broadcast()
					clientAuthorized = 1
					q.authCond.Broadcast()
				}

				f := atomic.AddInt64(&q.finished, 1)
				bar.Prefix(fmt.Sprintf("(%d of %d files) ", f, q.files))
				q.wg.Done()
			}
		}(i)
	}

	close(q.uploadc)
	<-workersReady
	q.authCond.Signal() // Signal the first goroutine to run
	q.wg.Wait()
	close(q.errorc)

	bar.Finish()
}

// Errors returns any errors encountered during uploading.
func (q *UploadQueue) Errors() []*WrappedError {
	return q.errors
}

// ensureFile makes sure that the cleanPath exists before pushing it.  If it
// does not exist, it attempts to clean it by reading the file at smudgePath.
func ensureFile(smudgePath, cleanPath string) error {
	if _, err := os.Stat(cleanPath); err == nil {
		return nil
	}

	expectedOid := filepath.Base(cleanPath)
	localPath := filepath.Join(LocalWorkingDir, smudgePath)
	file, err := os.Open(localPath)
	if err != nil {
		return err
	}

	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return err
	}

	cleaned, err := PointerClean(file, stat.Size(), nil)
	if err != nil {
		return err
	}

	cleaned.Close()

	if expectedOid != cleaned.Oid {
		return fmt.Errorf("Expected %s to have an OID of %s, got %s", smudgePath, expectedOid, cleaned.Oid)
	}

	return nil
}
<|endoftext|>


<commit_before>package lfs

import (
	"fmt"
	"os"
	"path/filepath"
)

// Uploadable describes a file that can be uploaded.
type Uploadable struct {
	oid      string
	OidPath  string
	Filename string
	size     int64
	object   *objectResource
}

// NewUploadable builds the Uploadable from the given information.
func NewUploadable(oid, filename string, index, totalFiles int) (*Uploadable, *WrappedError) {
	path, err := LocalMediaPath(oid)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	if err := ensureFile(filename, path); err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	fi, err := os.Stat(filename)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	return &Uploadable{oid: oid, OidPath: path, Filename: filename, size: fi.Size()}, nil
}

func (u *Uploadable) Check() (*objectResource, *WrappedError) {
	return UploadCheck(u.OidPath)
}

func (u *Uploadable) Transfer(cb CopyCallback) *WrappedError {
	wcb := func(total, read int64, current int) error {
		cb(total, read, current)
		return nil
	}

	return UploadObject(u.object, wcb)
}

func (u *Uploadable) Object() *objectResource {
	return u.object
}

func (u *Uploadable) Oid() string {
	return u.oid
}

func (u *Uploadable) Size() int64 {
	return u.size
}

func (u *Uploadable) Name() string {
	return u.Filename
}

func (u *Uploadable) SetObject(o *objectResource) {
	u.object = o
}

// NewUploadQueue builds an UploadQueue, allowing `workers` concurrent uploads.
func NewUploadQueue(workers, files int) *TransferQueue {
	q := newTransferQueue(workers, files)
	q.transferKind = "upload"
	return q
}

// ensureFile makes sure that the cleanPath exists before pushing it.  If it
// does not exist, it attempts to clean it by reading the file at smudgePath.
func ensureFile(smudgePath, cleanPath string) error {
	if _, err := os.Stat(cleanPath); err == nil {
		return nil
	}

	expectedOid := filepath.Base(cleanPath)
	localPath := filepath.Join(LocalWorkingDir, smudgePath)
	file, err := os.Open(localPath)
	if err != nil {
		return err
	}

	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return err
	}

	cleaned, err := PointerClean(file, stat.Size(), nil)
	if err != nil {
		return err
	}

	cleaned.Close()

	if expectedOid != cleaned.Oid {
		return fmt.Errorf("Expected %s to have an OID of %s, got %s", smudgePath, expectedOid, cleaned.Oid)
	}

	return nil
}
<commit_msg>アー アアアア アーアー<commit_after>package lfs

import (
	"fmt"
	"os"
	"path/filepath"
)

// Uploadable describes a file that can be uploaded.
type Uploadable struct {
	oid      string
	OidPath  string
	Filename string
	size     int64
	object   *objectResource
}

// NewUploadable builds the Uploadable from the given information.
func NewUploadable(oid, filename string, index, totalFiles int) (*Uploadable, *WrappedError) {
	path, err := LocalMediaPath(oid)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	if err := ensureFile(filename, path); err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	fi, err := os.Stat(path)
	if err != nil {
		return nil, Errorf(err, "Error uploading file %s (%s)", filename, oid)
	}

	return &Uploadable{oid: oid, OidPath: path, Filename: filename, size: fi.Size()}, nil
}

func (u *Uploadable) Check() (*objectResource, *WrappedError) {
	return UploadCheck(u.OidPath)
}

func (u *Uploadable) Transfer(cb CopyCallback) *WrappedError {
	wcb := func(total, read int64, current int) error {
		cb(total, read, current)
		return nil
	}

	return UploadObject(u.object, wcb)
}

func (u *Uploadable) Object() *objectResource {
	return u.object
}

func (u *Uploadable) Oid() string {
	return u.oid
}

func (u *Uploadable) Size() int64 {
	return u.size
}

func (u *Uploadable) Name() string {
	return u.Filename
}

func (u *Uploadable) SetObject(o *objectResource) {
	u.object = o
}

// NewUploadQueue builds an UploadQueue, allowing `workers` concurrent uploads.
func NewUploadQueue(workers, files int) *TransferQueue {
	q := newTransferQueue(workers, files)
	q.transferKind = "upload"
	return q
}

// ensureFile makes sure that the cleanPath exists before pushing it.  If it
// does not exist, it attempts to clean it by reading the file at smudgePath.
func ensureFile(smudgePath, cleanPath string) error {
	if _, err := os.Stat(cleanPath); err == nil {
		return nil
	}

	expectedOid := filepath.Base(cleanPath)
	localPath := filepath.Join(LocalWorkingDir, smudgePath)
	file, err := os.Open(localPath)
	if err != nil {
		return err
	}

	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		return err
	}

	cleaned, err := PointerClean(file, stat.Size(), nil)
	if err != nil {
		return err
	}

	cleaned.Close()

	if expectedOid != cleaned.Oid {
		return fmt.Errorf("Expected %s to have an OID of %s, got %s", smudgePath, expectedOid, cleaned.Oid)
	}

	return nil
}
<|endoftext|>


<commit_before>// +build darwin

package load

import (
	"context"
	"os/exec"
	"strconv"
	"strings"

	"github.com/shirou/gopsutil/internal/common"
)

func Avg() (*AvgStat, error) {
	return AvgWithContext(context.Background())
}

func AvgWithContext(ctx context.Context) (*AvgStat, error) {
	values, err := common.DoSysctrlWithContext(ctx, "vm.loadavg")
	if err != nil {
		return nil, err
	}

	load1, err := strconv.ParseFloat(values[0], 64)
	if err != nil {
		return nil, err
	}
	load5, err := strconv.ParseFloat(values[1], 64)
	if err != nil {
		return nil, err
	}
	load15, err := strconv.ParseFloat(values[2], 64)
	if err != nil {
		return nil, err
	}

	ret := &AvgStat{
		Load1:  float64(load1),
		Load5:  float64(load5),
		Load15: float64(load15),
	}

	return ret, nil
}

// Misc returnes miscellaneous host-wide statistics.
// darwin use ps command to get process running/blocked count.
// Almost same as FreeBSD implementation, but state is different.
// U means 'Uninterruptible Sleep'.
func Misc() (*MiscStat, error) {
	return MiscWithContext(context.Background())
}

func MiscWithContext(ctx context.Context) (*MiscStat, error) {
	bin, err := exec.LookPath("ps")
	if err != nil {
		return nil, err
	}
	out, err := invoke.CommandWithContext(ctx, bin, "axo", "state")
	if err != nil {
		return nil, err
	}
	lines := strings.Split(string(out), "\n")

	ret := MiscStat{}
	for _, l := range lines {
		if strings.Contains(l, "R") {
			ret.ProcsRunning++
		} else if strings.Contains(l, "U") {
			// uninterruptible sleep == blocked
			ret.ProcsBlocked++
		}
	}

	return &ret, nil
}
<commit_msg>[load][darwin] Remove calls to sysctl binary in load/load_darwin.go #639<commit_after>// +build darwin

package load

import (
	"context"
	"os/exec"
	"strings"
	"unsafe"

	"golang.org/x/sys/unix"
)

func Avg() (*AvgStat, error) {
	return AvgWithContext(context.Background())
}

func AvgWithContext(ctx context.Context) (*AvgStat, error) {
	// This SysctlRaw method borrowed from
	// https://github.com/prometheus/node_exporter/blob/master/collector/loadavg_freebsd.go
	// this implementation is common with BSDs
	type loadavg struct {
		load  [3]uint32
		scale int
	}
	b, err := unix.SysctlRaw("vm.loadavg")
	if err != nil {
		return nil, err
	}
	load := *(*loadavg)(unsafe.Pointer((&b[0])))
	scale := float64(load.scale)
	ret := &AvgStat{
		Load1:  float64(load.load[0]) / scale,
		Load5:  float64(load.load[1]) / scale,
		Load15: float64(load.load[2]) / scale,
	}

	return ret, nil
}

// Misc returnes miscellaneous host-wide statistics.
// darwin use ps command to get process running/blocked count.
// Almost same as FreeBSD implementation, but state is different.
// U means 'Uninterruptible Sleep'.
func Misc() (*MiscStat, error) {
	return MiscWithContext(context.Background())
}

func MiscWithContext(ctx context.Context) (*MiscStat, error) {
	bin, err := exec.LookPath("ps")
	if err != nil {
		return nil, err
	}
	out, err := invoke.CommandWithContext(ctx, bin, "axo", "state")
	if err != nil {
		return nil, err
	}
	lines := strings.Split(string(out), "\n")

	ret := MiscStat{}
	for _, l := range lines {
		if strings.Contains(l, "R") {
			ret.ProcsRunning++
		} else if strings.Contains(l, "U") {
			// uninterruptible sleep == blocked
			ret.ProcsBlocked++
		}
	}

	return &ret, nil
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/pion/rtcp"
	"github.com/pion/webrtc/v2"
)

const (
	rtcpPLIInterval = time.Second * 3
)

type TrackEventType uint32

const (
	TrackEventTypeAdd = iota + 1
	TrackEventTypeRemove
)

type TrackEvent struct {
	ClientID string
	Track    *webrtc.Track
	Type     TrackEventType
}

type TrackMetadata struct {
	Mid      string `json:"mid"`
	UserID   string `json:"userId"`
	StreamID string `json:"streamId"`
	Kind     string `json:"kind"`
}

type TrackInfo struct {
	RTPTransceiver *webrtc.RTPTransceiver
	RTPSender      *webrtc.RTPSender
	TrackMetadata  TrackMetadata
}

type trackListener struct {
	log              Logger
	clientID         string
	peerConnection   *webrtc.PeerConnection
	localTracks      []*webrtc.Track
	localTracksMu    sync.RWMutex
	trackInfoByTrack map[*webrtc.Track]TrackInfo
	onTrackEvent     func(TrackEvent)
	mu               sync.RWMutex
}

func newTrackListener(
	loggerFactory LoggerFactory,
	clientID string,
	peerConnection *webrtc.PeerConnection,
	onTrackEvent func(TrackEvent),
) *trackListener {
	p := &trackListener{
		log:              loggerFactory.GetLogger("peer"),
		clientID:         clientID,
		peerConnection:   peerConnection,
		trackInfoByTrack: map[*webrtc.Track]TrackInfo{},
		onTrackEvent:     onTrackEvent,
	}

	p.log.Printf("[%s] Setting PeerConnection.OnTrack listener", clientID)
	peerConnection.OnTrack(p.handleTrack)

	return p
}

func (p *trackListener) ClientID() string {
	return p.clientID
}

// GetTracksMetadata gets metadata of the sending tracks with updated Mid
func (p *trackListener) GetTracksMetadata() (metadata []TrackMetadata) {
	p.mu.RLock()
	defer p.mu.RUnlock()

	metadata = make([]TrackMetadata, 0)

	for _, trackInfo := range p.trackInfoByTrack {
		m := trackInfo.TrackMetadata
		m.Mid = trackInfo.RTPTransceiver.Mid()
		metadata = append(metadata, m)
	}
	return
}

func (p *trackListener) AddTrack(sourcePC *webrtc.PeerConnection, sourceClientID string, track *webrtc.Track) error {
	p.localTracksMu.Lock()
	defer p.localTracksMu.Unlock()

	p.log.Printf("[%s] peer.AddTrack: add sendonly transceiver for track: %s", p.clientID, track.ID())
	rtpSender, err := p.peerConnection.AddTrack(track)

	go func() {
		for {
			rtcps, err := rtpSender.ReadRTCP()
			if err != nil {
				p.log.Printf("[%s] Error reading rtcp for sender track: %d: %s",
					p.clientID,
					track.SSRC(),
					err,
				)
				break
			}
			for _, pkt := range rtcps {
				switch pkt.(type) {
				case *rtcp.PictureLossIndication:
					err := sourcePC.WriteRTCP(
						[]rtcp.Packet{
							&rtcp.PictureLossIndication{
								MediaSSRC: track.SSRC(),
							},
						},
					)
					if err != nil {
						p.log.Printf("[%s] Error sending rtcp PLI for local track: %d: %s",
							p.clientID,
							track.SSRC(),
							err,
						)
					}
					// default:
					//   p.log.Printf("[%s] Got rtcp pkt for track: %d (%T)", p.clientID, track.SSRC(), pkt)
				}
			}
		}
	}()
	// t, err := p.peerConnection.AddTransceiverFromTrack(
	// 	track,
	// 	webrtc.RtpTransceiverInit{
	// 		Direction: webrtc.RTPTransceiverDirectionSendonly,
	// 	},
	// )

	var transceiver *webrtc.RTPTransceiver
	for _, tr := range p.peerConnection.GetTransceivers() {
		if tr.Sender() == rtpSender {
			transceiver = tr
			break
		}
	}

	if err != nil {
		return fmt.Errorf("[%s] peer.AddTrack: error adding track: %s: %s", p.clientID, track.ID(), err)
	}

	p.trackInfoByTrack[track] = TrackInfo{
		RTPSender:      rtpSender,
		RTPTransceiver: transceiver,
		TrackMetadata: TrackMetadata{
			Mid:      "",
			Kind:     track.Kind().String(),
			UserID:   sourceClientID,
			StreamID: track.Label(),
		},
	}
	return nil
}

func (p *trackListener) RemoveTrack(track *webrtc.Track) error {
	p.localTracksMu.Lock()
	defer p.localTracksMu.Unlock()
	p.log.Printf("[%s] peer.RemoveTrack: %s", p.clientID, track.ID())
	trackInfo, ok := p.trackInfoByTrack[track]
	if !ok {
		return fmt.Errorf("[%s] peer.RemoveTrack: cannot find sender for track: %s", p.clientID, track.ID())
	}
	delete(p.trackInfoByTrack, track)
	return p.peerConnection.RemoveTrack(trackInfo.RTPSender)
}

func (p *trackListener) handleTrack(remoteTrack *webrtc.Track, receiver *webrtc.RTPReceiver) {
	p.log.Printf("[%s] peer.handleTrack (id: %s, label: %s, type: %s, ssrc: %d)",
		p.clientID, remoteTrack.ID(), remoteTrack.Label(), remoteTrack.Kind(), remoteTrack.SSRC())
	localTrack, err := p.startCopyingTrack(remoteTrack, receiver)
	if err != nil {
		p.log.Printf("Error copying remote track: %s", err)
		return
	}
	p.localTracksMu.Lock()
	p.localTracks = append(p.localTracks, localTrack)
	p.localTracksMu.Unlock()

	p.log.Printf("[%s] peer.handleTrack add track to list of local tracks: %s", p.clientID, localTrack.ID())

	p.sendTrackEvent(TrackEvent{p.clientID, localTrack, TrackEventTypeAdd})
}

func (p *trackListener) sendTrackEvent(t TrackEvent) {
	go p.onTrackEvent(t)
}

func (p *trackListener) Tracks() []*webrtc.Track {
	return p.localTracks
}

func (p *trackListener) startCopyingTrack(remoteTrack *webrtc.Track, receiver *webrtc.RTPReceiver) (*webrtc.Track, error) {
	remoteTrackID := remoteTrack.ID()
	if remoteTrackID == "" {
		remoteTrackID = NewUUIDBase62()
	}
	// this is the media stream ID we add the p.clientID in the string to know
	// which user the video came from and the remoteTrack.Label() so we can
	// associate audio/video tracks from the same MediaStream
	remoteTrackLabel := remoteTrack.Label()
	if remoteTrackLabel == "" {
		remoteTrackLabel = NewUUIDBase62()
	}
	localTrackLabel := "sfu_" + p.clientID + "_" + remoteTrackLabel

	localTrackID := "sfu_" + remoteTrackID
	p.log.Printf("[%s] peer.startCopyingTrack: (id: %s, label: %s) to (id: %s, label: %s), ssrc: %d",
		p.clientID, remoteTrack.ID(), remoteTrack.Label(), localTrackID, localTrackLabel, remoteTrack.SSRC())

	ssrc := remoteTrack.SSRC()
	// Create a local track, all our SFU clients will be fed via this track
	localTrack, err := p.peerConnection.NewTrack(remoteTrack.PayloadType(), ssrc, localTrackID, localTrackLabel)
	if err != nil {
		err = fmt.Errorf("[%s] peer.startCopyingTrack: error creating new track, trackID: %s, error: %s", p.clientID, remoteTrack.ID(), err)
		return nil, err
	}

	// Send a PLI on an interval so that the publisher is pushing a keyframe every rtcpPLIInterval
	// This can be less wasteful by processing incoming RTCP events, then we would emit a NACK/PLI when a viewer requests it

	go func() {
		defer p.sendTrackEvent(TrackEvent{p.clientID, localTrack, TrackEventTypeRemove})
		rtpBuf := make([]byte, 1400)
		for {
			i, err := remoteTrack.Read(rtpBuf)
			if err != nil {
				p.log.Printf(
					"[%s] Error reading from remote track: %s: %s",
					p.clientID,
					remoteTrack.ID(),
					err,
				)
				return
			}

			// ErrClosedPipe means we don't have any subscribers, this is ok if no peers have connected yet
			if _, err = localTrack.Write(rtpBuf[:i]); err != nil && err != io.ErrClosedPipe {
				p.log.Printf(
					"[%s] Error writing to local track: %s: %s",
					p.clientID,
					localTrackID,
					err,
				)
				return
			}
		}
	}()

	return localTrack, nil
}
<commit_msg>Reuse existing PLI packet<commit_after>package server

import (
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/pion/rtcp"
	"github.com/pion/webrtc/v2"
)

const (
	rtcpPLIInterval = time.Second * 3
)

type TrackEventType uint32

const (
	TrackEventTypeAdd = iota + 1
	TrackEventTypeRemove
)

type TrackEvent struct {
	ClientID string
	Track    *webrtc.Track
	Type     TrackEventType
}

type TrackMetadata struct {
	Mid      string `json:"mid"`
	UserID   string `json:"userId"`
	StreamID string `json:"streamId"`
	Kind     string `json:"kind"`
}

type TrackInfo struct {
	RTPTransceiver *webrtc.RTPTransceiver
	RTPSender      *webrtc.RTPSender
	TrackMetadata  TrackMetadata
}

type trackListener struct {
	log              Logger
	clientID         string
	peerConnection   *webrtc.PeerConnection
	localTracks      []*webrtc.Track
	localTracksMu    sync.RWMutex
	trackInfoByTrack map[*webrtc.Track]TrackInfo
	onTrackEvent     func(TrackEvent)
	mu               sync.RWMutex
}

func newTrackListener(
	loggerFactory LoggerFactory,
	clientID string,
	peerConnection *webrtc.PeerConnection,
	onTrackEvent func(TrackEvent),
) *trackListener {
	p := &trackListener{
		log:              loggerFactory.GetLogger("peer"),
		clientID:         clientID,
		peerConnection:   peerConnection,
		trackInfoByTrack: map[*webrtc.Track]TrackInfo{},
		onTrackEvent:     onTrackEvent,
	}

	p.log.Printf("[%s] Setting PeerConnection.OnTrack listener", clientID)
	peerConnection.OnTrack(p.handleTrack)

	return p
}

func (p *trackListener) ClientID() string {
	return p.clientID
}

// GetTracksMetadata gets metadata of the sending tracks with updated Mid
func (p *trackListener) GetTracksMetadata() (metadata []TrackMetadata) {
	p.mu.RLock()
	defer p.mu.RUnlock()

	metadata = make([]TrackMetadata, 0)

	for _, trackInfo := range p.trackInfoByTrack {
		m := trackInfo.TrackMetadata
		m.Mid = trackInfo.RTPTransceiver.Mid()
		metadata = append(metadata, m)
	}
	return
}

func (p *trackListener) AddTrack(sourcePC *webrtc.PeerConnection, sourceClientID string, track *webrtc.Track) error {
	p.localTracksMu.Lock()
	defer p.localTracksMu.Unlock()

	p.log.Printf("[%s] peer.AddTrack: add sendonly transceiver for track: %s", p.clientID, track.ID())
	rtpSender, err := p.peerConnection.AddTrack(track)

	go func() {
		for {
			rtcps, err := rtpSender.ReadRTCP()
			if err != nil {
				p.log.Printf("[%s] Error reading rtcp for sender track: %d: %s",
					p.clientID,
					track.SSRC(),
					err,
				)
				break
			}
			for _, pkt := range rtcps {
				switch pktType := pkt.(type) {
				case *rtcp.PictureLossIndication:
					err := sourcePC.WriteRTCP(
						[]rtcp.Packet{
							pktType,
							// &rtcp.PictureLossIndication{
							// 	MediaSSRC: track.SSRC(),
							// },
						},
					)
					if err != nil {
						p.log.Printf("[%s] Error sending rtcp PLI for local track: %d: %s",
							p.clientID,
							track.SSRC(),
							err,
						)
					}
					// default:
					//   p.log.Printf("[%s] Got rtcp pkt for track: %d (%T)", p.clientID, track.SSRC(), pkt)
				}
			}
		}
	}()
	// t, err := p.peerConnection.AddTransceiverFromTrack(
	// 	track,
	// 	webrtc.RtpTransceiverInit{
	// 		Direction: webrtc.RTPTransceiverDirectionSendonly,
	// 	},
	// )

	var transceiver *webrtc.RTPTransceiver
	for _, tr := range p.peerConnection.GetTransceivers() {
		if tr.Sender() == rtpSender {
			transceiver = tr
			break
		}
	}

	if err != nil {
		return fmt.Errorf("[%s] peer.AddTrack: error adding track: %s: %s", p.clientID, track.ID(), err)
	}

	p.trackInfoByTrack[track] = TrackInfo{
		RTPSender:      rtpSender,
		RTPTransceiver: transceiver,
		TrackMetadata: TrackMetadata{
			Mid:      "",
			Kind:     track.Kind().String(),
			UserID:   sourceClientID,
			StreamID: track.Label(),
		},
	}
	return nil
}

func (p *trackListener) RemoveTrack(track *webrtc.Track) error {
	p.localTracksMu.Lock()
	defer p.localTracksMu.Unlock()
	p.log.Printf("[%s] peer.RemoveTrack: %s", p.clientID, track.ID())
	trackInfo, ok := p.trackInfoByTrack[track]
	if !ok {
		return fmt.Errorf("[%s] peer.RemoveTrack: cannot find sender for track: %s", p.clientID, track.ID())
	}
	delete(p.trackInfoByTrack, track)
	return p.peerConnection.RemoveTrack(trackInfo.RTPSender)
}

func (p *trackListener) handleTrack(remoteTrack *webrtc.Track, receiver *webrtc.RTPReceiver) {
	p.log.Printf("[%s] peer.handleTrack (id: %s, label: %s, type: %s, ssrc: %d)",
		p.clientID, remoteTrack.ID(), remoteTrack.Label(), remoteTrack.Kind(), remoteTrack.SSRC())
	localTrack, err := p.startCopyingTrack(remoteTrack, receiver)
	if err != nil {
		p.log.Printf("Error copying remote track: %s", err)
		return
	}
	p.localTracksMu.Lock()
	p.localTracks = append(p.localTracks, localTrack)
	p.localTracksMu.Unlock()

	p.log.Printf("[%s] peer.handleTrack add track to list of local tracks: %s", p.clientID, localTrack.ID())

	p.sendTrackEvent(TrackEvent{p.clientID, localTrack, TrackEventTypeAdd})
}

func (p *trackListener) sendTrackEvent(t TrackEvent) {
	go p.onTrackEvent(t)
}

func (p *trackListener) Tracks() []*webrtc.Track {
	return p.localTracks
}

func (p *trackListener) startCopyingTrack(remoteTrack *webrtc.Track, receiver *webrtc.RTPReceiver) (*webrtc.Track, error) {
	remoteTrackID := remoteTrack.ID()
	if remoteTrackID == "" {
		remoteTrackID = NewUUIDBase62()
	}
	// this is the media stream ID we add the p.clientID in the string to know
	// which user the video came from and the remoteTrack.Label() so we can
	// associate audio/video tracks from the same MediaStream
	remoteTrackLabel := remoteTrack.Label()
	if remoteTrackLabel == "" {
		remoteTrackLabel = NewUUIDBase62()
	}
	localTrackLabel := "sfu_" + p.clientID + "_" + remoteTrackLabel

	localTrackID := "sfu_" + remoteTrackID
	p.log.Printf("[%s] peer.startCopyingTrack: (id: %s, label: %s) to (id: %s, label: %s), ssrc: %d",
		p.clientID, remoteTrack.ID(), remoteTrack.Label(), localTrackID, localTrackLabel, remoteTrack.SSRC())

	ssrc := remoteTrack.SSRC()
	// Create a local track, all our SFU clients will be fed via this track
	localTrack, err := p.peerConnection.NewTrack(remoteTrack.PayloadType(), ssrc, localTrackID, localTrackLabel)
	if err != nil {
		err = fmt.Errorf("[%s] peer.startCopyingTrack: error creating new track, trackID: %s, error: %s", p.clientID, remoteTrack.ID(), err)
		return nil, err
	}

	// Send a PLI on an interval so that the publisher is pushing a keyframe every rtcpPLIInterval
	// This can be less wasteful by processing incoming RTCP events, then we would emit a NACK/PLI when a viewer requests it

	go func() {
		defer p.sendTrackEvent(TrackEvent{p.clientID, localTrack, TrackEventTypeRemove})
		rtpBuf := make([]byte, 1400)
		for {
			i, err := remoteTrack.Read(rtpBuf)
			if err != nil {
				p.log.Printf(
					"[%s] Error reading from remote track: %s: %s",
					p.clientID,
					remoteTrack.ID(),
					err,
				)
				return
			}

			// ErrClosedPipe means we don't have any subscribers, this is ok if no peers have connected yet
			if _, err = localTrack.Write(rtpBuf[:i]); err != nil && err != io.ErrClosedPipe {
				p.log.Printf(
					"[%s] Error writing to local track: %s: %s",
					p.clientID,
					localTrackID,
					err,
				)
				return
			}
		}
	}()

	return localTrack, nil
}
<|endoftext|>


<commit_before>// Copyright 2013, Cong Ding. 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.
//
// author: Cong Ding <dinggnu@gmail.com>
//
package logging

import (
	"fmt"
	"strings"
)

// pre-defined formats
const (
	BasicFormat = "%s [%s] %s - %s\n name,levelname,asctime,message"
	RichFormat  = "%s [%s] %s - %d - %s:%s:%d - %s\n name, levelname, asctime, thread, filename, funcName, lineno, message"
)

func (logger *logging) genLog(level Level, message string) string {
	format := strings.Split(logger.format, "\n")
	if len(format) != 2 {
		return "logging format error"
	}
	args := strings.Split(format[1], ",")
	fs := make([]interface{}, len(args))
	l := new(log)
	l.message = message
	l.level = level
	for k, v := range args {
		fs[k] = fields[strings.TrimSpace(v)](logger, l)
	}
	return fmt.Sprintf(format[0], fs...)
}
<commit_msg>add seqid to richformat<commit_after>// Copyright 2013, Cong Ding. 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.
//
// author: Cong Ding <dinggnu@gmail.com>
//
package logging

import (
	"fmt"
	"strings"
)

// pre-defined formats
const (
	BasicFormat = "%s [%s] %s - %s\n name,levelname,asctime,message"
	RichFormat  = "%s [%s] %d %s - %d - %s:%s:%d - %s\n name, levelname, seqid, asctime, thread, filename, funcName, lineno, message"
)

func (logger *logging) genLog(level Level, message string) string {
	format := strings.Split(logger.format, "\n")
	if len(format) != 2 {
		return "logging format error"
	}
	args := strings.Split(format[1], ",")
	fs := make([]interface{}, len(args))
	l := new(log)
	l.message = message
	l.level = level
	for k, v := range args {
		fs[k] = fields[strings.TrimSpace(v)](logger, l)
	}
	return fmt.Sprintf(format[0], fs...)
}
<|endoftext|>


<commit_before>package vagrant

import (
	"bytes"
	"text/template"
)

var (
	vagrantFile = `# -*- mode: ruby -*-
# vi: set ft=ruby :

VAGRANTFILE_API_VERSION = "2"

$script = <<SCRIPT
echo I am provisioning...
date > /etc/vagrant_provisioned_at
wget "https://s3.amazonaws.com/kodingdev-softlayer/softlayer"
chmod +x softlayer
./softlayer -data {{ .ProvisionData }}
SCRIPT

Vagrant.configure(VAGRANTFILE_API_VERSION) do |config|
  config.vm.box = "{{ .Box }}"
  config.vm.hostname = "{{ .Hostname }}"

  config.vm.provider "virtualbox" do |vb|
    # Use VBoxManage to customize the VM. For example to change memory:
    vb.customize ["modifyvm", :id, "--memory", "{{ .Memory }}", "--cpus", "{{ .Cpus }}"]
  end

  config.vm.network :forwarded_port, guest: 56789, host: 56790, auto_correct: true
  config.vm.provision "shell", inline: $script
end
`

	vagrantTemplate = template.Must(template.New("vagrant").Parse(vagrantFile))
)

func createTemplate(opts *VagrantCreateOptions) (string, error) {
	buf := new(bytes.Buffer)
	if err := vagrantTemplate.Execute(buf, opts); err != nil {
		return "", err
	}

	return buf.String(), nil
}
<commit_msg>vagrant: change script to provisionklient<commit_after>package vagrant

import (
	"bytes"
	"text/template"
)

var (
	vagrantFile = `# -*- mode: ruby -*-
# vi: set ft=ruby :

VAGRANTFILE_API_VERSION = "2"

$script = <<SCRIPT
echo I am provisioning...
date > /etc/vagrant_provisioned_at
wget "https://s3.amazonaws.com/kodingdev-provision/provisionklient"
chmod +x provisionklient
./provisionklient -data {{ .ProvisionData }}
SCRIPT

Vagrant.configure(VAGRANTFILE_API_VERSION) do |config|
  config.vm.box = "{{ .Box }}"
  config.vm.hostname = "{{ .Hostname }}"

  config.vm.provider "virtualbox" do |vb|
    # Use VBoxManage to customize the VM. For example to change memory:
    vb.customize ["modifyvm", :id, "--memory", "{{ .Memory }}", "--cpus", "{{ .Cpus }}"]
  end

  config.vm.network :forwarded_port, guest: 56789, host: 56790, auto_correct: true
  config.vm.provision "shell", inline: $script
end
`

	vagrantTemplate = template.Must(template.New("vagrant").Parse(vagrantFile))
)

func createTemplate(opts *VagrantCreateOptions) (string, error) {
	buf := new(bytes.Buffer)
	if err := vagrantTemplate.Execute(buf, opts); err != nil {
		return "", err
	}

	return buf.String(), nil
}
<|endoftext|>


<commit_before>package graph

import (
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func getTestGraph() *Graph {
	g := NewGraph()

	a := NewNode("A")
	b := NewNode("B")
	c := NewNode("C")

	g.Add(a)
	g.Add(b)
	g.Add(c)

	g.AddArc(a, b)
	g.AddArc(a, c)
	g.AddArc(b, c)
	g.AddArc(c, b)

	return g
}

func TestCreate(t *testing.T) {
	Convey("A graph should be created without a panic", t, func() {
		So(func() { getTestGraph() }, ShouldNotPanic)
	})
}

func TestPrint(t *testing.T) {
	t.Log(getTestGraph().String())
}

func TestAdjacent(t *testing.T) {
	g := getTestGraph()

	Convey("A should be adjacent to B and C should not be adjacent to A", t, func() {
		So(g.Adjacent(g.Get("A"), g.Get("B")), ShouldBeTrue)
		So(g.Adjacent(g.Get("C"), g.Get("A")), ShouldBeFalse)
	})
}

func TestNeighbors(t *testing.T) {
	g := getTestGraph()

	Convey("A should point to B and C", t, func() {
		So(g.Neighbors(g.Get("A")), ShouldResemble, []*Node{g.Get("B"), g.Get("C")})
	})
}

func TestPointingTo(t *testing.T) {
	g := getTestGraph()

	Convey("A and B should point to C", t, func() {
		So(g.PointingTo(g.Get("C")), ShouldResemble, []*Node{g.Get("A"), g.Get("B")})
	})
}

func TestRemove(t *testing.T) {
	g := getTestGraph()

	Convey("Removing B should work", t, func() {
		So(g.Get("B"), ShouldNotBeNil)

		So(func() { g.Remove(g.Get("B")) }, ShouldNotPanic)

		So(g.Get("B"), ShouldBeNil)
	})
}

func TestAddArc(t *testing.T) {
	g := getTestGraph()

	Convey("Adding an arc between B and A should make them adjacent", t, func() {
		So(func() { g.AddArc(g.Get("B"), g.Get("A")) }, ShouldNotPanic)
		So(g.Adjacent(g.Get("B"), g.Get("A")), ShouldBeTrue)
	})
}

func TestRemoveArc(t *testing.T) {
	g := getTestGraph()

	Convey("Removing an arc between A and B should make them non adjacent", t, func() {
		So(func() { g.RemoveArc(g.Get("A"), g.Get("B")) }, ShouldNotPanic)
		So(g.Adjacent(g.Get("A"), g.Get("B")), ShouldBeFalse)
	})
}

func BenchmarkAdd(b *testing.B) {
	g := NewGraph()
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		g.Add(NewNode(string(i)))
	}
}

func BenchmarkRemove(b *testing.B) {
	g := NewGraph()

	for i := 0; i < b.N; i++ {
		g.Add(NewNode(string(i)))
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		g.Remove(g.Get(string(i)))
	}
}

func BenchmarkAdjacent(b *testing.B) {

}
<commit_msg>adds unit tests for graph.Import and graph.Export.<commit_after>package graph

import (
	. "github.com/smartystreets/goconvey/convey"
	"os"
	"testing"
)

func getTestGraph() *Graph {
	g := NewGraph()

	a := NewNode("A")
	b := NewNode("B")
	c := NewNode("C")

	g.Add(a)
	g.Add(b)
	g.Add(c)

	g.AddArc(a, b)
	g.AddArc(a, c)
	g.AddArc(b, c)
	g.AddArc(c, b)

	return g
}

func TestCreate(t *testing.T) {
	Convey("A graph should be created without a panic", t, func() {
		So(func() { getTestGraph() }, ShouldNotPanic)
	})
}

func TestPrint(t *testing.T) {
	t.Log(getTestGraph().String())
}

func TestAdjacent(t *testing.T) {
	g := getTestGraph()

	Convey("A should be adjacent to B and C should not be adjacent to A", t, func() {
		So(g.Adjacent(g.Get("A"), g.Get("B")), ShouldBeTrue)
		So(g.Adjacent(g.Get("C"), g.Get("A")), ShouldBeFalse)
	})
}

func TestNeighbors(t *testing.T) {
	g := getTestGraph()

	Convey("A should point to B and C", t, func() {
		So(g.Neighbors(g.Get("A")), ShouldResemble, []*Node{g.Get("B"), g.Get("C")})
	})
}

func TestPointingTo(t *testing.T) {
	g := getTestGraph()

	Convey("A and B should point to C", t, func() {
		So(g.PointingTo(g.Get("C")), ShouldResemble, []*Node{g.Get("A"), g.Get("B")})
	})
}

func TestRemove(t *testing.T) {
	g := getTestGraph()

	Convey("Removing B should work", t, func() {
		So(g.Get("B"), ShouldNotBeNil)

		So(func() { g.Remove(g.Get("B")) }, ShouldNotPanic)

		So(g.Get("B"), ShouldBeNil)
	})
}

func TestAddArc(t *testing.T) {
	g := getTestGraph()

	Convey("Adding an arc between B and A should make them adjacent", t, func() {
		So(func() { g.AddArc(g.Get("B"), g.Get("A")) }, ShouldNotPanic)
		So(g.Adjacent(g.Get("B"), g.Get("A")), ShouldBeTrue)
	})
}

func TestRemoveArc(t *testing.T) {
	g := getTestGraph()

	Convey("Removing an arc between A and B should make them non adjacent", t, func() {
		So(func() { g.RemoveArc(g.Get("A"), g.Get("B")) }, ShouldNotPanic)
		So(g.Adjacent(g.Get("A"), g.Get("B")), ShouldBeFalse)
	})
}

func TestImportExport(t *testing.T) {
	g := getTestGraph()

	file, err := os.Create(".test-graph.gob")

	if err != nil {
		t.Error("error opening graph file: ", err)
	}

	Convey("Exporting should work", t, func() {
		So(func() { g.Export(file) }, ShouldNotPanic)
	})

	var importedGraph *Graph
	var importError error

	in, err := os.Open(".test-graph.gob")

	if err != nil {
		t.Error("error opening graph file: ", err)
	}

	Convey("Importing should work", t, func() {
		So(func() {
			importError, importedGraph = Import(in)
		}, ShouldNotPanic)
		So(importError, ShouldBeNil)

		So(g.String(), ShouldEqual, importedGraph.String())
	})
}

func BenchmarkAdd(b *testing.B) {
	g := NewGraph()
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		g.Add(NewNode(string(i)))
	}
}

func BenchmarkRemove(b *testing.B) {
	g := NewGraph()

	for i := 0; i < b.N; i++ {
		g.Add(NewNode(string(i)))
	}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		g.Remove(g.Get(string(i)))
	}
}

func BenchmarkAdjacent(b *testing.B) {

}
<|endoftext|>


<commit_before>// Copyright (c) 2014 Ludovic Fauvet
// Licensed under the MIT license

package main

import (
	"errors"
	"fmt"
	"github.com/garyburd/redigo/redis"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	healthCheckThreads = 10
	userAgent          = "Mirrorbits/" + VERSION + " PING CHECK"
	clientTimeout      = time.Duration(20 * time.Second)
	clientDeadline     = time.Duration(40 * time.Second)
	redirectError      = errors.New("Redirect not allowed")
)

type Monitor struct {
	redis           *redisobj
	cache           *Cache
	mirrors         map[string]*MMirror
	mapLock         sync.Mutex
	httpClient      http.Client
	healthCheckChan chan string
	syncChan        chan string
	stop            chan bool
	wg              sync.WaitGroup
}

type MMirror struct {
	Mirror
	checking   bool
	scanning   bool
	lastUpdate int64
	lastScan   int64 //FIXME this could be removed since the mirrors are updated after each events no?
}

func NewMonitor(r *redisobj, c *Cache) *Monitor {
	monitor := new(Monitor)
	monitor.redis = r
	monitor.cache = c
	monitor.mirrors = make(map[string]*MMirror)
	monitor.healthCheckChan = make(chan string, healthCheckThreads*5)
	monitor.syncChan = make(chan string)
	monitor.stop = make(chan bool)

	transport := http.Transport{
		DisableKeepAlives:   true,
		MaxIdleConnsPerHost: 0,
		Dial: func(network, addr string) (net.Conn, error) {
			deadline := time.Now().Add(clientDeadline)
			c, err := net.DialTimeout(network, addr, clientTimeout)
			if err != nil {
				return nil, err
			}
			c.SetDeadline(deadline)
			return c, nil
		},
	}

	monitor.httpClient = http.Client{
		CheckRedirect: checkRedirect,
		Transport:     &transport,
	}
	return monitor
}

func (m *Monitor) Terminate() {
	close(m.stop)
	m.wg.Wait()
}

// Return an error if the endpoint is an unauthorized redirect
func checkRedirect(req *http.Request, via []*http.Request) error {
	if GetConfig().DisallowRedirects {
		return redirectError
	}
	return nil
}

// Main monitor loop
func (m *Monitor) monitorLoop() {
	m.wg.Add(1)
	m.syncSource()

	mirrorUpdateEvent := make(chan string, 10)
	m.cache.SubscribeEvent(MIRROR_UPDATE, mirrorUpdateEvent)

	for {
		ids, err := m.mirrorsID()
		if err == nil {
			m.syncMirrors(ids...)
			break
		}
		select {
		case <-m.stop:
			m.wg.Done()
			return
		default:
		}
		time.Sleep(500 * time.Millisecond)
	}

	for i := 0; i < healthCheckThreads; i++ {
		go m.healthCheckLoop()
	}

	for i := 0; i < GetConfig().ConcurrentSync; i++ {
		go m.syncLoop()
	}

	// Setup recurrent tasks
	sourceSyncTicker := time.NewTicker(5 * time.Minute)

	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case id := <-mirrorUpdateEvent:
			m.syncMirrors(id)
		case <-sourceSyncTicker.C:
			m.syncSource()
		case <-time.After(1 * time.Second):
			m.mapLock.Lock()
			for k, v := range m.mirrors {
				if v.lastUpdate+60 < time.Now().UTC().Unix() &&
					m.mirrors[k].checking == false {
					select {
					case m.healthCheckChan <- k:
						m.mirrors[k].checking = true
					default:
					}
				}
				if v.lastScan+60*30 < time.Now().UTC().Unix() &&
					m.mirrors[k].scanning == false {
					select {
					case m.syncChan <- k:
						m.mirrors[k].scanning = true
					default:
					}
				}
			}
			m.mapLock.Unlock()
		}
	}
}

// Returns a list of all mirrors ID
func (m *Monitor) mirrorsID() ([]string, error) {
	rconn := m.redis.pool.Get()
	defer rconn.Close()

	return redis.Strings(rconn.Do("LRANGE", "MIRRORS", "0", "-1"))
}

// Sync the remote mirror struct with the local dataset
// TODO needs improvements
func (m *Monitor) syncMirrors(mirrorsIDs ...string) ([]Mirror, error) {

	mirrors := make([]Mirror, 0, len(mirrorsIDs))

	for _, id := range mirrorsIDs {
		mirror, err := m.cache.GetMirror(id)
		if err != nil && err != redis.ErrNil {
			log.Error("Fetching mirror %s failed: %s", id, err.Error())
			continue
		} else if err == redis.ErrNil {
			// Mirror has been deleted
			m.mapLock.Lock()
			delete(m.mirrors, id)
			m.mapLock.Unlock()
			continue
		}
		if mirror.Enabled {
			mirrors = append(mirrors, mirror)
		}
	}

	m.mapLock.Lock()

	// Prepare the list of mirrors
	for _, e := range mirrors {
		var isChecking bool = false
		var isScanning bool = false
		var lastUpdate int64 = 0
		var lastScan int64 = e.LastSync
		if _, ok := m.mirrors[e.ID]; ok {
			isChecking = m.mirrors[e.ID].checking
			isScanning = m.mirrors[e.ID].scanning
			lastUpdate = m.mirrors[e.ID].lastUpdate
			if m.mirrors[e.ID].lastScan > e.LastSync {
				lastScan = m.mirrors[e.ID].lastScan
			}
		}
		m.mirrors[e.ID] = &MMirror{
			Mirror:     e,
			checking:   isChecking,
			scanning:   isScanning,
			lastUpdate: lastUpdate,
			lastScan:   lastScan,
		}
	}
	m.mapLock.Unlock()

	log.Debug("Mirror list updated")
	return mirrors, nil
}

// Main health check loop
// TODO merge with the monitorLoop?
func (m *Monitor) healthCheckLoop() {
	m.wg.Add(1)
	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case k := <-m.healthCheckChan:
			m.mapLock.Lock()
			mirror := m.mirrors[k]
			m.mapLock.Unlock()

			m.healthCheck(mirror.Mirror)

			m.mapLock.Lock()
			if _, ok := m.mirrors[k]; ok {
				m.mirrors[k].lastUpdate = time.Now().UTC().Unix()
				m.mirrors[k].checking = false
			}
			m.mapLock.Unlock()
		}
	}
}

// Main sync loop
// TODO merge with the monitorLoop?
func (m *Monitor) syncLoop() {
	m.wg.Add(1)
	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case k := <-m.syncChan:
			m.mapLock.Lock()
			mirror := m.mirrors[k]
			m.mapLock.Unlock()

			log.Debug("Scanning %s", k)

			var err error = NoSyncMethod

			// First try to scan with rsync
			if mirror.RsyncURL != "" {
				err = Scan(m.redis).ScanRsync(mirror.RsyncURL, k, m.stop)
			}
			// If it failed or rsync wasn't supported
			// fallback to FTP
			if err != nil && mirror.FtpURL != "" {
				err = Scan(m.redis).ScanFTP(mirror.FtpURL, k, m.stop)
			}

			if err != nil {
				log.Error("%-30.30s Scan failed: %s", k, err.Error())
			}

			select {
			case m.healthCheckChan <- k:
			default:
			}

			m.mapLock.Lock()
			if _, ok := m.mirrors[k]; ok {
				m.mirrors[k].lastScan = time.Now().UTC().Unix()
				m.mirrors[k].scanning = false
			}
			m.mapLock.Unlock()
		}
	}
}

// Do an actual health check against a given mirror
func (m *Monitor) healthCheck(mirror Mirror) error {
	file, size, err := m.getRandomFile(mirror.ID)
	if err != nil {
		if err != redis.ErrNil {
			log.Warning("%-30.30s Error: Cannot obtain a random file: %s", mirror.ID, err)
		}
		return err
	}

	req, err := http.NewRequest("HEAD", strings.TrimRight(mirror.HttpURL, "/")+file, nil)
	req.Header.Set("User-Agent", userAgent)
	req.Close = true

	start := time.Now()
	resp, err := m.httpClient.Do(req)
	elapsed := time.Since(start)

	if err != nil {
		if opErr, ok := err.(*net.OpError); ok {
			log.Debug("Op: %s | Net: %s | Addr: %s | Err: %s | Temporary: %t", opErr.Op, opErr.Net, opErr.Addr, opErr.Error(), opErr.Temporary())
		}
		markMirrorDown(m.redis, mirror.ID, "Unreachable")
		log.Error("%-30.30s Error: %s (%dms)", mirror.ID, err.Error(), elapsed/time.Millisecond)
		return err
	}
	resp.Body.Close()

	contentLength := resp.Header.Get("Content-Length")

	if resp.StatusCode == 404 {
		markMirrorDown(m.redis, mirror.ID, fmt.Sprintf("File not found %s (error 404)", file))
		if GetConfig().DisableOnMissingFile {
			disableMirror(m.redis, mirror.ID)
		}
		log.Error("%-30.30s Error: File %s not found (error 404)", mirror.ID, file)
	} else if resp.StatusCode != 200 {
		markMirrorDown(m.redis, mirror.ID, fmt.Sprintf("Got status code %d", resp.StatusCode))
		log.Warning("%-30.30s Down! Status: %d", mirror.ID, resp.StatusCode)
	} else {
		markMirrorUp(m.redis, mirror.ID)
		rsize, err := strconv.ParseInt(contentLength, 10, 64)
		if err == nil && rsize != size {
			log.Warning("%-30.30s File size mismatch! [%s] (%dms)", mirror.ID, file, elapsed/time.Millisecond)
		} else {
			log.Notice("%-30.30s Up! (%dms)", mirror.ID, elapsed/time.Millisecond)
		}
	}
	return nil
}

// Get a random filename known to be served by the given mirror
func (m *Monitor) getRandomFile(identifier string) (file string, size int64, err error) {
	sinterKey := fmt.Sprintf("HANDLEDFILES_%s", identifier)

	rconn := m.redis.pool.Get()
	defer rconn.Close()

	file, err = redis.String(rconn.Do("SRANDMEMBER", sinterKey))
	if err != nil {
		return
	}

	size, err = redis.Int64(rconn.Do("HGET", fmt.Sprintf("FILE_%s", file), "size"))
	if err != nil {
		return
	}

	return
}

// Trigger a sync of the local repository
func (m *Monitor) syncSource() {
	err := Scan(m.redis).ScanSource(m.stop)
	if err != nil {
		log.Error("Scanning source failed: %s", err.Error())
	}
}
<commit_msg>monitor: remove a blocking sleep<commit_after>// Copyright (c) 2014 Ludovic Fauvet
// Licensed under the MIT license

package main

import (
	"errors"
	"fmt"
	"github.com/garyburd/redigo/redis"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	healthCheckThreads = 10
	userAgent          = "Mirrorbits/" + VERSION + " PING CHECK"
	clientTimeout      = time.Duration(20 * time.Second)
	clientDeadline     = time.Duration(40 * time.Second)
	redirectError      = errors.New("Redirect not allowed")
)

type Monitor struct {
	redis           *redisobj
	cache           *Cache
	mirrors         map[string]*MMirror
	mapLock         sync.Mutex
	httpClient      http.Client
	healthCheckChan chan string
	syncChan        chan string
	stop            chan bool
	wg              sync.WaitGroup
}

type MMirror struct {
	Mirror
	checking   bool
	scanning   bool
	lastUpdate int64
	lastScan   int64 //FIXME this could be removed since the mirrors are updated after each events no?
}

func NewMonitor(r *redisobj, c *Cache) *Monitor {
	monitor := new(Monitor)
	monitor.redis = r
	monitor.cache = c
	monitor.mirrors = make(map[string]*MMirror)
	monitor.healthCheckChan = make(chan string, healthCheckThreads*5)
	monitor.syncChan = make(chan string)
	monitor.stop = make(chan bool)

	transport := http.Transport{
		DisableKeepAlives:   true,
		MaxIdleConnsPerHost: 0,
		Dial: func(network, addr string) (net.Conn, error) {
			deadline := time.Now().Add(clientDeadline)
			c, err := net.DialTimeout(network, addr, clientTimeout)
			if err != nil {
				return nil, err
			}
			c.SetDeadline(deadline)
			return c, nil
		},
	}

	monitor.httpClient = http.Client{
		CheckRedirect: checkRedirect,
		Transport:     &transport,
	}
	return monitor
}

func (m *Monitor) Terminate() {
	close(m.stop)
	m.wg.Wait()
}

// Return an error if the endpoint is an unauthorized redirect
func checkRedirect(req *http.Request, via []*http.Request) error {
	if GetConfig().DisallowRedirects {
		return redirectError
	}
	return nil
}

// Main monitor loop
func (m *Monitor) monitorLoop() {
	m.wg.Add(1)
	m.syncSource()

	mirrorUpdateEvent := make(chan string, 10)
	m.cache.SubscribeEvent(MIRROR_UPDATE, mirrorUpdateEvent)

	for {
		ids, err := m.mirrorsID()
		if err == nil {
			m.syncMirrors(ids...)
			break
		}
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case <-time.After(500 * time.Millisecond):
		}
	}

	for i := 0; i < healthCheckThreads; i++ {
		go m.healthCheckLoop()
	}

	for i := 0; i < GetConfig().ConcurrentSync; i++ {
		go m.syncLoop()
	}

	// Setup recurrent tasks
	sourceSyncTicker := time.NewTicker(5 * time.Minute)

	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case id := <-mirrorUpdateEvent:
			m.syncMirrors(id)
		case <-sourceSyncTicker.C:
			m.syncSource()
		case <-time.After(1 * time.Second):
			m.mapLock.Lock()
			for k, v := range m.mirrors {
				if v.lastUpdate+60 < time.Now().UTC().Unix() &&
					m.mirrors[k].checking == false {
					select {
					case m.healthCheckChan <- k:
						m.mirrors[k].checking = true
					default:
					}
				}
				if v.lastScan+60*30 < time.Now().UTC().Unix() &&
					m.mirrors[k].scanning == false {
					select {
					case m.syncChan <- k:
						m.mirrors[k].scanning = true
					default:
					}
				}
			}
			m.mapLock.Unlock()
		}
	}
}

// Returns a list of all mirrors ID
func (m *Monitor) mirrorsID() ([]string, error) {
	rconn := m.redis.pool.Get()
	defer rconn.Close()

	return redis.Strings(rconn.Do("LRANGE", "MIRRORS", "0", "-1"))
}

// Sync the remote mirror struct with the local dataset
// TODO needs improvements
func (m *Monitor) syncMirrors(mirrorsIDs ...string) ([]Mirror, error) {

	mirrors := make([]Mirror, 0, len(mirrorsIDs))

	for _, id := range mirrorsIDs {
		mirror, err := m.cache.GetMirror(id)
		if err != nil && err != redis.ErrNil {
			log.Error("Fetching mirror %s failed: %s", id, err.Error())
			continue
		} else if err == redis.ErrNil {
			// Mirror has been deleted
			m.mapLock.Lock()
			delete(m.mirrors, id)
			m.mapLock.Unlock()
			continue
		}
		if mirror.Enabled {
			mirrors = append(mirrors, mirror)
		}
	}

	m.mapLock.Lock()

	// Prepare the list of mirrors
	for _, e := range mirrors {
		var isChecking bool = false
		var isScanning bool = false
		var lastUpdate int64 = 0
		var lastScan int64 = e.LastSync
		if _, ok := m.mirrors[e.ID]; ok {
			isChecking = m.mirrors[e.ID].checking
			isScanning = m.mirrors[e.ID].scanning
			lastUpdate = m.mirrors[e.ID].lastUpdate
			if m.mirrors[e.ID].lastScan > e.LastSync {
				lastScan = m.mirrors[e.ID].lastScan
			}
		}
		m.mirrors[e.ID] = &MMirror{
			Mirror:     e,
			checking:   isChecking,
			scanning:   isScanning,
			lastUpdate: lastUpdate,
			lastScan:   lastScan,
		}
	}
	m.mapLock.Unlock()

	log.Debug("Mirror list updated")
	return mirrors, nil
}

// Main health check loop
// TODO merge with the monitorLoop?
func (m *Monitor) healthCheckLoop() {
	m.wg.Add(1)
	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case k := <-m.healthCheckChan:
			m.mapLock.Lock()
			mirror := m.mirrors[k]
			m.mapLock.Unlock()

			m.healthCheck(mirror.Mirror)

			m.mapLock.Lock()
			if _, ok := m.mirrors[k]; ok {
				m.mirrors[k].lastUpdate = time.Now().UTC().Unix()
				m.mirrors[k].checking = false
			}
			m.mapLock.Unlock()
		}
	}
}

// Main sync loop
// TODO merge with the monitorLoop?
func (m *Monitor) syncLoop() {
	m.wg.Add(1)
	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case k := <-m.syncChan:
			m.mapLock.Lock()
			mirror := m.mirrors[k]
			m.mapLock.Unlock()

			log.Debug("Scanning %s", k)

			var err error = NoSyncMethod

			// First try to scan with rsync
			if mirror.RsyncURL != "" {
				err = Scan(m.redis).ScanRsync(mirror.RsyncURL, k, m.stop)
			}
			// If it failed or rsync wasn't supported
			// fallback to FTP
			if err != nil && mirror.FtpURL != "" {
				err = Scan(m.redis).ScanFTP(mirror.FtpURL, k, m.stop)
			}

			if err != nil {
				log.Error("%-30.30s Scan failed: %s", k, err.Error())
			}

			select {
			case m.healthCheckChan <- k:
			default:
			}

			m.mapLock.Lock()
			if _, ok := m.mirrors[k]; ok {
				m.mirrors[k].lastScan = time.Now().UTC().Unix()
				m.mirrors[k].scanning = false
			}
			m.mapLock.Unlock()
		}
	}
}

// Do an actual health check against a given mirror
func (m *Monitor) healthCheck(mirror Mirror) error {
	file, size, err := m.getRandomFile(mirror.ID)
	if err != nil {
		if err != redis.ErrNil {
			log.Warning("%-30.30s Error: Cannot obtain a random file: %s", mirror.ID, err)
		}
		return err
	}

	req, err := http.NewRequest("HEAD", strings.TrimRight(mirror.HttpURL, "/")+file, nil)
	req.Header.Set("User-Agent", userAgent)
	req.Close = true

	start := time.Now()
	resp, err := m.httpClient.Do(req)
	elapsed := time.Since(start)

	if err != nil {
		if opErr, ok := err.(*net.OpError); ok {
			log.Debug("Op: %s | Net: %s | Addr: %s | Err: %s | Temporary: %t", opErr.Op, opErr.Net, opErr.Addr, opErr.Error(), opErr.Temporary())
		}
		markMirrorDown(m.redis, mirror.ID, "Unreachable")
		log.Error("%-30.30s Error: %s (%dms)", mirror.ID, err.Error(), elapsed/time.Millisecond)
		return err
	}
	resp.Body.Close()

	contentLength := resp.Header.Get("Content-Length")

	if resp.StatusCode == 404 {
		markMirrorDown(m.redis, mirror.ID, fmt.Sprintf("File not found %s (error 404)", file))
		if GetConfig().DisableOnMissingFile {
			disableMirror(m.redis, mirror.ID)
		}
		log.Error("%-30.30s Error: File %s not found (error 404)", mirror.ID, file)
	} else if resp.StatusCode != 200 {
		markMirrorDown(m.redis, mirror.ID, fmt.Sprintf("Got status code %d", resp.StatusCode))
		log.Warning("%-30.30s Down! Status: %d", mirror.ID, resp.StatusCode)
	} else {
		markMirrorUp(m.redis, mirror.ID)
		rsize, err := strconv.ParseInt(contentLength, 10, 64)
		if err == nil && rsize != size {
			log.Warning("%-30.30s File size mismatch! [%s] (%dms)", mirror.ID, file, elapsed/time.Millisecond)
		} else {
			log.Notice("%-30.30s Up! (%dms)", mirror.ID, elapsed/time.Millisecond)
		}
	}
	return nil
}

// Get a random filename known to be served by the given mirror
func (m *Monitor) getRandomFile(identifier string) (file string, size int64, err error) {
	sinterKey := fmt.Sprintf("HANDLEDFILES_%s", identifier)

	rconn := m.redis.pool.Get()
	defer rconn.Close()

	file, err = redis.String(rconn.Do("SRANDMEMBER", sinterKey))
	if err != nil {
		return
	}

	size, err = redis.Int64(rconn.Do("HGET", fmt.Sprintf("FILE_%s", file), "size"))
	if err != nil {
		return
	}

	return
}

// Trigger a sync of the local repository
func (m *Monitor) syncSource() {
	err := Scan(m.redis).ScanSource(m.stop)
	if err != nil {
		log.Error("Scanning source failed: %s", err.Error())
	}
}
<|endoftext|>


<commit_before>package dbr

import (
	"fmt"
	"strings"
)

type Query struct {
	*Session

	SelectSql      string
	FromSql        string
	WhereFragments []string
	Params         []interface{}
	OrderSql       string
	LimitCount     uint64
	LimitValid     bool
	OffsetCount    uint64
	OffsetValid    bool
}

func (q *Query) From(fromSql string) *Query {
	q.FromSql = fromSql
	return q
}

func (q *Query) Where(fromSql string, args ...interface{}) *Query {
	q.WhereFragments = append(q.WhereFragments, fromSql)
	for _, a := range args {
		q.Params = append(q.Params, a)
	}
	return q
}

func (q *Query) Limit(limit uint64) *Query {
	q.LimitCount = limit
	q.LimitValid = true
	return q
}

func (q *Query) Offset(offset uint64) *Query {
	q.OffsetCount = offset
	q.OffsetValid = true
	return q
}

// Assumes page/perPage are valid. Page and perPage must be >= 1
func (q *Query) Paginate(page, perPage uint64) *Query {
	q.Limit(perPage)
	q.Offset((page - 1) * perPage)
	return q
}

func (q *Query) Order(ord string) *Query {
	q.OrderSql = ord
	return q
}

func (q *Query) OrderDir(ord string, isAsc bool) *Query {
	if isAsc {
		q.OrderSql = ord + " ASC"
	} else {
		q.OrderSql = ord + " DESC"
	}
	return q
}

func (q *Query) ToSql() (string, []interface{}) {
	var whereStr string
	if len(q.WhereFragments) >= 0 {
		whereStr = " WHERE (" + strings.Join(q.WhereFragments, ") AND (") + ") "
	}

	var limitStr string
	if q.LimitValid {
		limitStr = fmt.Sprintf(" LIMIT %d", q.LimitCount)
	}

	var offsetStr string
	if q.OffsetValid {
		offsetStr = fmt.Sprintf(" OFFSET %d", q.OffsetCount)
	}

	var orderStr string
	if q.OrderSql != "" {
		orderStr = " ORDER BY " + q.OrderSql
	}

	return fmt.Sprintf("SELECT * FROM %s%s%s%s%s", q.FromSql, whereStr, orderStr, limitStr, offsetStr), q.Params
}

func (q *Query) SelectAll(dest interface{}) (int, error) {
	sql, params := q.ToSql()
	return q.Session.SelectAll(dest, sql, params...)
}

func (q *Query) SelectOne(dest interface{}) (bool, error) {
	sql, params := q.ToSql()
	return q.Session.SelectOne(dest, sql, params...)
}
<commit_msg>private subdomains api<commit_after>package dbr

import (
	"fmt"
	"strings"
)

type Query struct {
	*Session

	SelectSql      string
	FromSql        string
	WhereFragments []string
	Params         []interface{}
	OrderSql       string
	LimitCount     uint64
	LimitValid     bool
	OffsetCount    uint64
	OffsetValid    bool
}

func (q *Query) From(fromSql string) *Query {
	q.FromSql = fromSql
	return q
}

func (q *Query) Where(fromSql string, args ...interface{}) *Query {
	q.WhereFragments = append(q.WhereFragments, fromSql)
	for _, a := range args {
		q.Params = append(q.Params, a)
	}
	return q
}

func (q *Query) Limit(limit uint64) *Query {
	q.LimitCount = limit
	q.LimitValid = true
	return q
}

func (q *Query) Offset(offset uint64) *Query {
	q.OffsetCount = offset
	q.OffsetValid = true
	return q
}

// Assumes page/perPage are valid. Page and perPage must be >= 1
func (q *Query) Paginate(page, perPage uint64) *Query {
	q.Limit(perPage)
	q.Offset((page - 1) * perPage)
	return q
}

func (q *Query) Order(ord string) *Query {
	q.OrderSql = ord
	return q
}

func (q *Query) OrderDir(ord string, isAsc bool) *Query {
	if isAsc {
		q.OrderSql = ord + " ASC"
	} else {
		q.OrderSql = ord + " DESC"
	}
	return q
}

func (q *Query) ToSql() (string, []interface{}) {
	var whereStr string
	if len(q.WhereFragments) > 0 {
		whereStr = " WHERE (" + strings.Join(q.WhereFragments, ") AND (") + ") "
	}

	var limitStr string
	if q.LimitValid {
		limitStr = fmt.Sprintf(" LIMIT %d", q.LimitCount)
	}

	var offsetStr string
	if q.OffsetValid {
		offsetStr = fmt.Sprintf(" OFFSET %d", q.OffsetCount)
	}

	var orderStr string
	if q.OrderSql != "" {
		orderStr = " ORDER BY " + q.OrderSql
	}

	return fmt.Sprintf("SELECT * FROM %s%s%s%s%s", q.FromSql, whereStr, orderStr, limitStr, offsetStr), q.Params
}

func (q *Query) SelectAll(dest interface{}) (int, error) {
	sql, params := q.ToSql()
	return q.Session.SelectAll(dest, sql, params...)
}

func (q *Query) SelectOne(dest interface{}) (bool, error) {
	sql, params := q.ToSql()
	return q.Session.SelectOne(dest, sql, params...)
}
<|endoftext|>


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

package iotest

import (
	"io";
	"log";
	"os";
)

type writeLogger struct {
	prefix	string;
	w	io.Writer;
}

func (l *writeLogger) Write(p []byte) (n int, err os.Error) {
	n, err = l.w.Write(p);
	if err != nil {
		log.Stdoutf("%s %x: %v", l.prefix, p[0:n], err)
	} else {
		log.Stdoutf("%s %x", l.prefix, p[0:n])
	}
	return;
}

// NewWriteLogger returns a writer that behaves like w except
// that it logs (using log.Stdout) each write to standard output,
// printing the prefix and the hexadecimal data written.
func NewWriteLogger(prefix string, w io.Writer) io.Writer {
	return &writeLogger{prefix, w}
}

type readLogger struct {
	prefix	string;
	r	io.Reader;
}

func (l *readLogger) Read(p []byte) (n int, err os.Error) {
	n, err = l.r.Read(p);
	if err != nil {
		log.Stdoutf("%s %x: %v", l.prefix, p[0:n], err)
	} else {
		log.Stdoutf("%s %x", l.prefix, p[0:n])
	}
	return;
}

// NewReadLogger returns a writer that behaves like w except
// that it logs (using log.Stdout) each write to standard output,
// printing the prefix and the hexadecimal data written.
func NewReadLogger(prefix string, r io.Reader) io.Reader {
	return &readLogger{prefix, r}
}
<commit_msg>Fix doc striing for NewReadLogger.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package iotest

import (
	"io";
	"log";
	"os";
)

type writeLogger struct {
	prefix	string;
	w	io.Writer;
}

func (l *writeLogger) Write(p []byte) (n int, err os.Error) {
	n, err = l.w.Write(p);
	if err != nil {
		log.Stdoutf("%s %x: %v", l.prefix, p[0:n], err)
	} else {
		log.Stdoutf("%s %x", l.prefix, p[0:n])
	}
	return;
}

// NewWriteLogger returns a writer that behaves like w except
// that it logs (using log.Stdout) each write to standard output,
// printing the prefix and the hexadecimal data written.
func NewWriteLogger(prefix string, w io.Writer) io.Writer {
	return &writeLogger{prefix, w}
}

type readLogger struct {
	prefix	string;
	r	io.Reader;
}

func (l *readLogger) Read(p []byte) (n int, err os.Error) {
	n, err = l.r.Read(p);
	if err != nil {
		log.Stdoutf("%s %x: %v", l.prefix, p[0:n], err)
	} else {
		log.Stdoutf("%s %x", l.prefix, p[0:n])
	}
	return;
}

// NewReadLogger returns a reader that behaves like r except
// that it logs (using log.Stdout) each read to standard output,
// printing the prefix and the hexadecimal data written.
func NewReadLogger(prefix string, r io.Reader) io.Reader {
	return &readLogger{prefix, r}
}
<|endoftext|>


<commit_before>package menu

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"sort"
	"strings"

	"github.com/fatih/color"
	"github.com/miquella/ask"
	"golang.org/x/crypto/ssh"

	"github.com/miquella/vaulted/lib"
)

type SSHKeyMenu struct {
	*Menu
}

func (m *SSHKeyMenu) Help() {
	menuColor.Set()
	defer color.Unset()

	fmt.Println("a,add    - Add")
	fmt.Println("D,delete - Delete")
	fmt.Println("v        - Vault Signing URL")
	fmt.Println("u,users  - Vault User Principals")
	fmt.Println("?,help   - Help")
	fmt.Println("b,back   - Back")
	fmt.Println("q,quit   - Quit")
}

func (m *SSHKeyMenu) Handler() error {
	for {
		var err error
		m.Printer()
		input, err := interaction.ReadMenu("Edit ssh keys: [a,D,v,u,b]: ")
		if err != nil {
			return err
		}
		switch input {
		case "a", "add", "key", "keys":
			err = m.AddSSHKey()
		case "D", "delete", "remove":
			var key string
			key, err = interaction.ReadValue("Key: ")
			if err == nil {
				if _, exists := m.Vault.SSHKeys[key]; exists {
					delete(m.Vault.SSHKeys, key)
				} else {
					color.Red("Key '%s' not found", key)
				}
			}
		case "v":
			signingUrl, err := interaction.ReadValue("Vault signing URL: ")
			if err != nil {
				return err
			}
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.VaultSigningUrl = signingUrl
		case "u", "users":
			userPrincipals, err := interaction.ReadValue("Vault user principals (comma separated): ")
			if err != nil {
				return err
			}
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			if userPrincipals != "" {
				m.Vault.SSHOptions.ValidPrincipals = strings.Split(userPrincipals, ",")
			} else {
				m.Vault.SSHOptions.ValidPrincipals = []string{}
			}
		case "b", "back":
			return nil
		case "q", "quit", "exit":
			var confirm string
			confirm, err = interaction.ReadValue("Are you sure you wish to save and exit the vault? (y/n): ")
			if err == nil {
				if confirm == "y" {
					return ErrSaveAndExit
				}
			}
		case "?", "help":
			m.Help()
		default:
			color.Red("Command not recognized")
		}

		if err != nil {
			return err
		}
	}
}

func (m *SSHKeyMenu) AddSSHKey() error {
	var err error

	homeDir := ""
	user, err := user.Current()
	if err == nil {
		homeDir = user.HomeDir
	} else {
		homeDir = os.Getenv("HOME")
	}

	defaultFilename := ""
	filename := ""
	if homeDir != "" {
		defaultFilename = filepath.Join(homeDir, ".ssh", "id_rsa")
		filename, err = interaction.ReadValue(fmt.Sprintf("Key file (default: %s): ", defaultFilename))
		if err != nil {
			return err
		}
		if filename == "" {
			filename = defaultFilename
		}
		if !filepath.IsAbs(filename) {
			filename = filepath.Join(filepath.Join(homeDir, ".ssh"), filename)
		}
	} else {
		filename, err = interaction.ReadValue("Key file: ")
		if err != nil {
			return err
		}
	}

	decryptedBlock, err := loadAndDecryptKey(filename)
	if err != nil {
		color.Red("%v", err)
		return nil
	}

	comment := loadPublicKeyComment(filename + ".pub")
	var name string
	if comment != "" {
		name, err = interaction.ReadValue(fmt.Sprintf("Name (default: %s): ", comment))
		if err != nil {
			return err
		}
		if name == "" {
			name = comment
		}
	} else {
		name, err = interaction.ReadValue("Name: ")
		if err != nil {
			return err
		}
		if name == "" {
			name = filename
		}
	}

	if m.Vault.SSHKeys == nil {
		m.Vault.SSHKeys = make(map[string]string)
	}
	m.Vault.SSHKeys[name] = string(pem.EncodeToMemory(decryptedBlock))

	return nil
}

func loadAndDecryptKey(filename string) (*pem.Block, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(data)
	if block == nil {
		return nil, err
	}

	if x509.IsEncryptedPEMBlock(block) {
		var passphrase string
		var decryptedBytes []byte
		for i := 0; i < 3; i++ {
			passphrase, err = ask.HiddenAsk("Passphrase: ")
			if err != nil {
				return nil, err
			}

			decryptedBytes, err = x509.DecryptPEMBlock(block, []byte(passphrase))
			if err == nil {
				break
			}
			if err != x509.IncorrectPasswordError {
				return nil, err
			}
		}

		if err != nil {
			return nil, err
		}

		return &pem.Block{
			Type:  block.Type,
			Bytes: decryptedBytes,
		}, nil
	}
	return block, nil
}

func loadPublicKeyComment(filename string) string {
	f, err := os.Open(filename)
	if err != nil {
		return ""
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	if err != nil {
		return ""
	}

	_, comment, _, _, err := ssh.ParseAuthorizedKey(data)
	if err != nil {
		return ""
	}
	return comment
}

func (m *SSHKeyMenu) Printer() {
	color.Cyan("\nSSH Keys:")
	if len(m.Vault.SSHKeys) > 0 {
		keys := []string{}
		for key := range m.Vault.SSHKeys {
			keys = append(keys, key)
		}
		sort.Strings(keys)

		for _, key := range keys {
			green.Printf("  %s\n", key)
		}
	} else {
		fmt.Println("  [Empty]")
	}

	if m.Vault.SSHOptions != nil {
		if m.Vault.SSHOptions.VaultSigningUrl != "" {
			color.Cyan("\nHashiCorp Vault:")
			green.Printf("  Signing URL: ")
			fmt.Printf("%s\n", m.Vault.SSHOptions.VaultSigningUrl)

			if len(m.Vault.SSHOptions.ValidPrincipals) > 0 {
				green.Printf("  User Principals: ")
				fmt.Printf("%s\n", m.Vault.SSHOptions.ValidPrincipals)
			}
		}
	}
}
<commit_msg>add generate key to ssh edit menu<commit_after>package menu

import (
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"sort"
	"strings"

	"github.com/fatih/color"
	"github.com/miquella/ask"
	"golang.org/x/crypto/ssh"

	"github.com/miquella/vaulted/lib"
)

type SSHKeyMenu struct {
	*Menu
}

func (m *SSHKeyMenu) Help() {
	menuColor.Set()
	defer color.Unset()

	fmt.Println("a,add      - Add")
	fmt.Println("D,delete   - Delete")
	fmt.Println("g,generate - Generate Key")
	fmt.Println("v          - Vault Signing URL")
	fmt.Println("u,users    - Vault User Principals")
	fmt.Println("?,help     - Help")
	fmt.Println("b,back     - Back")
	fmt.Println("q,quit     - Quit")
}

func (m *SSHKeyMenu) Handler() error {
	for {
		var err error
		m.Printer()
		input, err := interaction.ReadMenu("Edit ssh keys: [a,D,g,v,u,b]: ")
		if err != nil {
			return err
		}
		switch input {
		case "a", "add", "key", "keys":
			err = m.AddSSHKey()
		case "D", "delete", "remove":
			var key string
			key, err = interaction.ReadValue("Key: ")
			if err == nil {
				if _, exists := m.Vault.SSHKeys[key]; exists {
					delete(m.Vault.SSHKeys, key)
				} else {
					color.Red("Key '%s' not found", key)
				}
			}
		case "g", "generate":
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.GenerateRSAKey = !m.Vault.SSHOptions.GenerateRSAKey
		case "v":
			signingUrl, err := interaction.ReadValue("Vault signing URL: ")
			if err != nil {
				return err
			}
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			m.Vault.SSHOptions.VaultSigningUrl = signingUrl
		case "u", "users":
			userPrincipals, err := interaction.ReadValue("Vault user principals (comma separated): ")
			if err != nil {
				return err
			}
			if m.Vault.SSHOptions == nil {
				m.Vault.SSHOptions = &vaulted.SSHOptions{}
			}
			if userPrincipals != "" {
				m.Vault.SSHOptions.ValidPrincipals = strings.Split(userPrincipals, ",")
			} else {
				m.Vault.SSHOptions.ValidPrincipals = []string{}
			}
		case "b", "back":
			return nil
		case "q", "quit", "exit":
			var confirm string
			confirm, err = interaction.ReadValue("Are you sure you wish to save and exit the vault? (y/n): ")
			if err == nil {
				if confirm == "y" {
					return ErrSaveAndExit
				}
			}
		case "?", "help":
			m.Help()
		default:
			color.Red("Command not recognized")
		}

		if err != nil {
			return err
		}
	}
}

func (m *SSHKeyMenu) AddSSHKey() error {
	var err error

	homeDir := ""
	user, err := user.Current()
	if err == nil {
		homeDir = user.HomeDir
	} else {
		homeDir = os.Getenv("HOME")
	}

	defaultFilename := ""
	filename := ""
	if homeDir != "" {
		defaultFilename = filepath.Join(homeDir, ".ssh", "id_rsa")
		filename, err = interaction.ReadValue(fmt.Sprintf("Key file (default: %s): ", defaultFilename))
		if err != nil {
			return err
		}
		if filename == "" {
			filename = defaultFilename
		}
		if !filepath.IsAbs(filename) {
			filename = filepath.Join(filepath.Join(homeDir, ".ssh"), filename)
		}
	} else {
		filename, err = interaction.ReadValue("Key file: ")
		if err != nil {
			return err
		}
	}

	decryptedBlock, err := loadAndDecryptKey(filename)
	if err != nil {
		color.Red("%v", err)
		return nil
	}

	comment := loadPublicKeyComment(filename + ".pub")
	var name string
	if comment != "" {
		name, err = interaction.ReadValue(fmt.Sprintf("Name (default: %s): ", comment))
		if err != nil {
			return err
		}
		if name == "" {
			name = comment
		}
	} else {
		name, err = interaction.ReadValue("Name: ")
		if err != nil {
			return err
		}
		if name == "" {
			name = filename
		}
	}

	if m.Vault.SSHKeys == nil {
		m.Vault.SSHKeys = make(map[string]string)
	}
	m.Vault.SSHKeys[name] = string(pem.EncodeToMemory(decryptedBlock))

	return nil
}

func loadAndDecryptKey(filename string) (*pem.Block, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	block, _ := pem.Decode(data)
	if block == nil {
		return nil, err
	}

	if x509.IsEncryptedPEMBlock(block) {
		var passphrase string
		var decryptedBytes []byte
		for i := 0; i < 3; i++ {
			passphrase, err = ask.HiddenAsk("Passphrase: ")
			if err != nil {
				return nil, err
			}

			decryptedBytes, err = x509.DecryptPEMBlock(block, []byte(passphrase))
			if err == nil {
				break
			}
			if err != x509.IncorrectPasswordError {
				return nil, err
			}
		}

		if err != nil {
			return nil, err
		}

		return &pem.Block{
			Type:  block.Type,
			Bytes: decryptedBytes,
		}, nil
	}
	return block, nil
}

func loadPublicKeyComment(filename string) string {
	f, err := os.Open(filename)
	if err != nil {
		return ""
	}
	defer f.Close()

	data, err := ioutil.ReadAll(f)
	if err != nil {
		return ""
	}

	_, comment, _, _, err := ssh.ParseAuthorizedKey(data)
	if err != nil {
		return ""
	}
	return comment
}

func (m *SSHKeyMenu) Printer() {
	color.Cyan("\nSSH Keys:")
	if len(m.Vault.SSHKeys) > 0 {
		keys := []string{}
		for key := range m.Vault.SSHKeys {
			keys = append(keys, key)
		}
		sort.Strings(keys)

		for _, key := range keys {
			green.Printf("  %s\n", key)
		}
	} else {
		fmt.Println("  [Empty]")
	}

	if m.Vault.SSHOptions != nil {
		color.Green("\n  Generate SSH Key: %t\n", m.Vault.SSHOptions.GenerateRSAKey)

		if m.Vault.SSHOptions.VaultSigningUrl != "" {
			color.Cyan("\nHashiCorp Vault:")
			green.Printf("  Signing URL: ")
			fmt.Printf("%s\n", m.Vault.SSHOptions.VaultSigningUrl)

			if len(m.Vault.SSHOptions.ValidPrincipals) > 0 {
				green.Printf("  User Principals: ")
				fmt.Printf("%s\n", m.Vault.SSHOptions.ValidPrincipals)
			}
		}
	}
}
<|endoftext|>


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

package merkle

import (
	"errors"
	"fmt"
	"math/big"
	"sort"

	"github.com/golang/glog"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/storage"
)

var (
	// ErrSubtreeOverrun indicates that a subtree exceeds the maximum tree depth.
	ErrSubtreeOverrun = errors.New("subtree with prefix exceeds maximum tree size")
	smtOne            = big.NewInt(1)
	smtZero           = big.NewInt(0)
)

// HStar2LeafHash represents a leaf for the HStar2 sparse Merkle tree
// implementation.
type HStar2LeafHash struct {
	// TODO(al): remove big.Int
	Index    *big.Int
	LeafHash []byte
}

// HStar2 is a recursive implementation for calculating the root hash of a sparse
// Merkle tree.
type HStar2 struct {
	treeID int64
	hasher hashers.MapHasher
}

// NewHStar2 creates a new HStar2 tree calculator based on the passed in MapHasher.
func NewHStar2(treeID int64, hasher hashers.MapHasher) HStar2 {
	return HStar2{
		treeID: treeID,
		hasher: hasher,
	}
}

// HStar2Root calculates the root of a sparse Merkle tree of a given depth
// which contains the given set of non-null leaves.
func (s *HStar2) HStar2Root(depth int, values []HStar2LeafHash) ([]byte, error) {
	sort.Sort(ByIndex{values})
	return s.hStar2b(0, depth, values, smtZero, nil, nil)
}

// SparseGetNodeFunc should return any pre-existing node hash for the node address.
type SparseGetNodeFunc func(depth int, index *big.Int) ([]byte, error)

// SparseSetNodeFunc should store the passed node hash, associating it with the address.
type SparseSetNodeFunc func(depth int, index *big.Int, hash []byte) error

// HStar2Nodes calculates the root hash of a pre-existing sparse Merkle tree
// plus the extra values passed in.  Get and set are used to fetch and store
// internal node values. Values must not contain multiple leaves for the same
// index.
//
// prefix is the location of this subtree within the larger tree. Root is at nil.
// subtreeDepth is the number of levels in this subtree.
func (s *HStar2) HStar2Nodes(prefix []byte, subtreeDepth int, values []HStar2LeafHash,
	get SparseGetNodeFunc, set SparseSetNodeFunc) ([]byte, error) {
	if glog.V(3) {
		glog.Infof("HStar2Nodes(%x, %v, %v)", prefix, subtreeDepth, len(values))
		for _, v := range values {
			glog.Infof("  %x: %x", v.Index.Bytes(), v.LeafHash)
		}
	}
	depth := len(prefix) * 8
	totalDepth := depth + subtreeDepth
	if totalDepth > s.hasher.BitLen() {
		return nil, ErrSubtreeOverrun
	}
	sort.Sort(ByIndex{values})
	offset := storage.NewNodeIDFromPrefixSuffix(prefix, storage.Suffix{}, s.hasher.BitLen()).BigInt()
	return s.hStar2b(depth, totalDepth, values, offset, get, set)
}

// hStar2b computes a sparse Merkle tree root value recursively.
func (s *HStar2) hStar2b(depth, maxDepth int, values []HStar2LeafHash, offset *big.Int,
	get SparseGetNodeFunc, set SparseSetNodeFunc) ([]byte, error) {
	if depth == maxDepth {
		switch {
		case len(values) == 0:
			return s.get(offset, depth, get)
		case len(values) == 1:
			return values[0].LeafHash, nil
		default:
			return nil, fmt.Errorf("hStar2b base case: len(values): %d, want 1", len(values))
		}
	}
	if len(values) == 0 {
		return s.get(offset, depth, get)
	}

	bitsLeft := s.hasher.BitLen() - depth
	split := new(big.Int).Lsh(smtOne, uint(bitsLeft-1))
	split.Add(split, offset)
	i := sort.Search(len(values), func(i int) bool { return values[i].Index.Cmp(split) >= 0 })
	lhs, err := s.hStar2b(depth+1, maxDepth, values[:i], offset, get, set)
	if err != nil {
		return nil, err
	}
	rhs, err := s.hStar2b(depth+1, maxDepth, values[i:], split, get, set)
	if err != nil {
		return nil, err
	}
	h := s.hasher.HashChildren(lhs, rhs)
	s.set(offset, depth, h, set)
	return h, nil
}

// get attempts to use getter. If getter fails, returns the HashEmpty value.
func (s *HStar2) get(index *big.Int, depth int, getter SparseGetNodeFunc) ([]byte, error) {
	// if we've got a function for getting existing node values, try it:
	if getter != nil {
		h, err := getter(depth, index)
		if err != nil {
			return nil, err
		}
		// if we got a value then we'll use that
		if h != nil {
			return h, nil
		}
	}
	// TODO(gdbelvin): Hashers should accept depth as their main argument.
	height := s.hasher.BitLen() - depth
	nodeID := storage.NewNodeIDFromBigInt(index.BitLen(), index, s.hasher.BitLen())
	return s.hasher.HashEmpty(s.treeID, nodeID.Path, height), nil
}

// set attempts to use setter if it not nil.
func (s *HStar2) set(index *big.Int, depth int, hash []byte, setter SparseSetNodeFunc) error {
	if setter != nil {
		return setter(depth, index, hash)
	}
	return nil
}

// HStar2LeafHash sorting boilerplate below.

// Leaves is a slice of HStar2LeafHash
type Leaves []HStar2LeafHash

// Len returns the number of leaves.
func (s Leaves) Len() int { return len(s) }

// Swap swaps two leaf locations.
func (s Leaves) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// ByIndex implements sort.Interface by providing Less and using Len and Swap methods from the embedded Leaves value.
type ByIndex struct{ Leaves }

// Less returns true if i.Index < j.Index
func (s ByIndex) Less(i, j int) bool { return s.Leaves[i].Index.Cmp(s.Leaves[j].Index) < 0 }
<commit_msg>merkle: logging for base recursion failure<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package merkle

import (
	"errors"
	"fmt"
	"math/big"
	"sort"

	"github.com/golang/glog"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/storage"
)

var (
	// ErrSubtreeOverrun indicates that a subtree exceeds the maximum tree depth.
	ErrSubtreeOverrun = errors.New("subtree with prefix exceeds maximum tree size")
	smtOne            = big.NewInt(1)
	smtZero           = big.NewInt(0)
)

// HStar2LeafHash represents a leaf for the HStar2 sparse Merkle tree
// implementation.
type HStar2LeafHash struct {
	// TODO(al): remove big.Int
	Index    *big.Int
	LeafHash []byte
}

// HStar2 is a recursive implementation for calculating the root hash of a sparse
// Merkle tree.
type HStar2 struct {
	treeID int64
	hasher hashers.MapHasher
}

// NewHStar2 creates a new HStar2 tree calculator based on the passed in MapHasher.
func NewHStar2(treeID int64, hasher hashers.MapHasher) HStar2 {
	return HStar2{
		treeID: treeID,
		hasher: hasher,
	}
}

// HStar2Root calculates the root of a sparse Merkle tree of a given depth
// which contains the given set of non-null leaves.
func (s *HStar2) HStar2Root(depth int, values []HStar2LeafHash) ([]byte, error) {
	sort.Sort(ByIndex{values})
	return s.hStar2b(0, depth, values, smtZero, nil, nil)
}

// SparseGetNodeFunc should return any pre-existing node hash for the node address.
type SparseGetNodeFunc func(depth int, index *big.Int) ([]byte, error)

// SparseSetNodeFunc should store the passed node hash, associating it with the address.
type SparseSetNodeFunc func(depth int, index *big.Int, hash []byte) error

// HStar2Nodes calculates the root hash of a pre-existing sparse Merkle tree
// plus the extra values passed in.  Get and set are used to fetch and store
// internal node values. Values must not contain multiple leaves for the same
// index.
//
// prefix is the location of this subtree within the larger tree. Root is at nil.
// subtreeDepth is the number of levels in this subtree.
func (s *HStar2) HStar2Nodes(prefix []byte, subtreeDepth int, values []HStar2LeafHash,
	get SparseGetNodeFunc, set SparseSetNodeFunc) ([]byte, error) {
	if glog.V(3) {
		glog.Infof("HStar2Nodes(%x, %v, %v)", prefix, subtreeDepth, len(values))
		for _, v := range values {
			glog.Infof("  %x: %x", v.Index.Bytes(), v.LeafHash)
		}
	}
	depth := len(prefix) * 8
	totalDepth := depth + subtreeDepth
	if totalDepth > s.hasher.BitLen() {
		return nil, ErrSubtreeOverrun
	}
	sort.Sort(ByIndex{values})
	offset := storage.NewNodeIDFromPrefixSuffix(prefix, storage.Suffix{}, s.hasher.BitLen()).BigInt()
	return s.hStar2b(depth, totalDepth, values, offset, get, set)
}

// hStar2b computes a sparse Merkle tree root value recursively.
func (s *HStar2) hStar2b(depth, maxDepth int, values []HStar2LeafHash, offset *big.Int,
	get SparseGetNodeFunc, set SparseSetNodeFunc) ([]byte, error) {
	if depth == maxDepth {
		switch {
		case len(values) == 0:
			return s.get(offset, depth, get)
		case len(values) == 1:
			return values[0].LeafHash, nil
		default:
			glog.Errorf("base case with too many values: %+v", values)
			return nil, fmt.Errorf("hStar2b base case: len(values): %d, want 1", len(values))
		}
	}
	if len(values) == 0 {
		return s.get(offset, depth, get)
	}

	bitsLeft := s.hasher.BitLen() - depth
	split := new(big.Int).Lsh(smtOne, uint(bitsLeft-1))
	split.Add(split, offset)
	i := sort.Search(len(values), func(i int) bool { return values[i].Index.Cmp(split) >= 0 })
	lhs, err := s.hStar2b(depth+1, maxDepth, values[:i], offset, get, set)
	if err != nil {
		return nil, err
	}
	rhs, err := s.hStar2b(depth+1, maxDepth, values[i:], split, get, set)
	if err != nil {
		return nil, err
	}
	h := s.hasher.HashChildren(lhs, rhs)
	s.set(offset, depth, h, set)
	return h, nil
}

// get attempts to use getter. If getter fails, returns the HashEmpty value.
func (s *HStar2) get(index *big.Int, depth int, getter SparseGetNodeFunc) ([]byte, error) {
	// if we've got a function for getting existing node values, try it:
	if getter != nil {
		h, err := getter(depth, index)
		if err != nil {
			return nil, err
		}
		// if we got a value then we'll use that
		if h != nil {
			return h, nil
		}
	}
	// TODO(gdbelvin): Hashers should accept depth as their main argument.
	height := s.hasher.BitLen() - depth
	nodeID := storage.NewNodeIDFromBigInt(index.BitLen(), index, s.hasher.BitLen())
	return s.hasher.HashEmpty(s.treeID, nodeID.Path, height), nil
}

// set attempts to use setter if it not nil.
func (s *HStar2) set(index *big.Int, depth int, hash []byte, setter SparseSetNodeFunc) error {
	if setter != nil {
		return setter(depth, index, hash)
	}
	return nil
}

// HStar2LeafHash sorting boilerplate below.

// Leaves is a slice of HStar2LeafHash
type Leaves []HStar2LeafHash

// Len returns the number of leaves.
func (s Leaves) Len() int { return len(s) }

// Swap swaps two leaf locations.
func (s Leaves) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// ByIndex implements sort.Interface by providing Less and using Len and Swap methods from the embedded Leaves value.
type ByIndex struct{ Leaves }

// Less returns true if i.Index < j.Index
func (s ByIndex) Less(i, j int) bool { return s.Leaves[i].Index.Cmp(s.Leaves[j].Index) < 0 }
<|endoftext|>


<commit_before>//
// Copyright 2017 Pedro Salgado
//
// 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 resque

import (
	"encoding/json"
	"fmt"

	"gopkg.in/redis.v5"
)

// Queue a job queue.
type Queue struct {
	redis        *redis.Client
	jobClassName string
	Name         string
}

// newQueue links a job queue to a redis instance.
func newQueue(jcn string, c *redis.Client) *Queue {
	return &Queue{
		redis:        c,
		jobClassName: jcn,
		Name:         fmt.Sprintf("resque:queue:%s", jcn),
	}
}

// Receive gets a job from the queue.
func (q Queue) Receive() (*Job, error) {
	cmd := q.redis.LPop(q.Name)
	json_str, err := cmd.Bytes()
	if err != nil {
		return nil, err
	}

	job := &Job{}
	err = json.Unmarshal(json_str, job)
	if err != nil {
		return nil, err
	}

	return job, err
}

// Send places a job on the queue.
func (q Queue) Send(args []JobArgument) error {
	json_str, err := json.Marshal(Job{Class: q.jobClassName})
	if err != nil {
		return err
	}

	return q.redis.RPush(q.Name, json_str).Err()
}
<commit_msg>queue.go: fixed lint warnings.<commit_after>//
// Copyright 2017 Pedro Salgado
//
// 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 resque

import (
	"encoding/json"
	"fmt"

	"gopkg.in/redis.v5"
)

// Queue a job queue.
type Queue struct {
	redis        *redis.Client
	jobClassName string
	Name         string
}

// newQueue links a job queue to a redis instance.
func newQueue(jcn string, c *redis.Client) *Queue {
	return &Queue{
		redis:        c,
		jobClassName: jcn,
		Name:         fmt.Sprintf("resque:queue:%s", jcn),
	}
}

// Receive gets a job from the queue.
func (q Queue) Receive() (*Job, error) {
	cmd := q.redis.LPop(q.Name)
	jsonStr, err := cmd.Bytes()
	if err != nil {
		return nil, err
	}

	job := &Job{}
	err = json.Unmarshal(jsonStr, job)
	if err != nil {
		return nil, err
	}

	return job, err
}

// Send places a job on the queue.
func (q Queue) Send(args []JobArgument) error {
	jsonStr, err := json.Marshal(Job{Class: q.jobClassName})
	if err != nil {
		return err
	}

	return q.redis.RPush(q.Name, jsonStr).Err()
}
<|endoftext|>


<commit_before>package c3

type queue struct {
	head    *entry
	tail    *entry
	version int
	length  int
}

type entry struct {
	item interface{}
	next *entry
}

func (q *queue) Len() int {
	return q.length
}

func (q *queue) Clear() {
	if q.length == 0 {
		return
	}
	q.head = nil
	q.tail = nil
	q.length = 0
	q.version++
}

func (q *queue) Contains(item interface{}) bool {
	for e := q.head; e != nil; e = e.next {
		if e.item == item {
			return true
		}
	}
	return false
}

func (q *queue) Peek() (interface{}, bool) {
	if q.head != nil {
		return q.head.item, true
	}
	return nil, false
}

func (q *queue) Dequeue() (interface{}, bool) {
	if q.head != nil {
		e := q.head
		q.head = e.next
		if e.next == nil {
			q.tail = nil
		}
		q.length--
		q.version++
		return e.item, true
	}
	return nil, false
}

func (q *queue) Enqueue(item interface{}) bool {
	e := &entry{item, nil}
	if q.tail == nil {
		q.head = e
		q.tail = e
	} else {
		q.tail.next = e
		q.tail = e
	}
	q.version++
	q.length++
	return true
}

func (q *queue) Iterator() Iterator {
	return &queueIterator{q, nil, false, q.version}
}

func (q *queue) Consumer() Consumer {
	return &queueConsumer{q, nil}
}
<commit_msg>added begin of entry reuse to minimize garbage created by the queue<commit_after>package c3

type queue struct {
	head    *entry
	tail    *entry
	free *entry
	version int
	length  int
}

type entry struct {
	item interface{}
	next *entry
}

func (q *queue) Len() int {
	return q.length
}

func (q *queue) Clear() {
	if q.length == 0 {
		return
	}
	q.head = nil
	q.tail = nil
	q.free = nil
	q.length = 0
	q.version++
}

func (q *queue) Contains(item interface{}) bool {
	for e := q.head; e != nil; e = e.next {
		if e.item == item {
			return true
		}
	}
	return false
}

func (q *queue) Peek() (interface{}, bool) {
	if q.head != nil {
		return q.head.item, true
	}
	return nil, false
}

func (q *queue) Dequeue() (interface{}, bool) {
// TODO add removed entry to the free list.
	if q.head != nil {
		e := q.head
		q.head = e.next
		if e.next == nil {
			q.tail = nil
		}
		q.length--
		q.version++
		return e.item, true
	}
	return nil, false
}

func (q *queue) Enqueue(item interface{}) bool {
// TODO check if there's any free entries.
	e := &entry{item, nil}
	if q.tail == nil {
		q.head = e
		q.tail = e
	} else {
		q.tail.next = e
		q.tail = e
	}
	q.version++
	q.length++
	return true
}

func (q *queue) Iterator() Iterator {
	return &queueIterator{q, nil, false, q.version}
}

func (q *queue) Consumer() Consumer {
	return &queueConsumer{q, nil}
}
<|endoftext|>