<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// The watcher package provides an interface for observing changes
// to arbitrary MongoDB documents that are maintained via the
// mgo/txn transaction package.
package watcher

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.state.watcher")

// A Watcher can watch any number of collections and documents for changes.
type Watcher struct {
	tomb tomb.Tomb
	log  *mgo.Collection

	// watches holds the observers managed by Watch/Unwatch.
	watches map[watchKey][]watchInfo

	// current holds the current txn-revno values for all the observed
	// documents known to exist. Documents not observed or deleted are
	// omitted from this map and are considered to have revno -1.
	current map[watchKey]int64

	// needSync is set when a synchronization should take
	// place.
	needSync bool

	// syncEvents and requestEvents contain the events to be
	// dispatched to the watcher channels. They're queued during
	// processing and flushed at the end to simplify the algorithm.
	// The two queues are separated because events from sync are
	// handled in reverse order due to the way the algorithm works.
	syncEvents, requestEvents []event

	// request is used to deliver requests from the public API into
	// the the goroutine loop.
	request chan interface{}

	// lastId is the most recent transaction id observed by a sync.
	lastId interface{}
}

// A Change holds information about a document change.
type Change struct {
	// C and Id hold the collection name and document _id field value.
	C  string
	Id interface{}

	// Revno is the latest known value for the document's txn-revno
	// field, or -1 if the document was deleted.
	Revno int64
}

type watchKey struct {
	c  string
	id interface{} // nil when watching collection
}

func (k watchKey) String() string {
	coll := "collection " + k.c
	if k.id == nil {
		return coll
	}
	return fmt.Sprintf("document %v in %s", k.id, coll)
}

// match returns whether the receiving watch key,
// which may refer to a particular item or
// an entire collection, matches k1, which refers
// to a particular item.
func (k watchKey) match(k1 watchKey) bool {
	if k.c != k1.c {
		return false
	}
	if k.id == nil {
		// k refers to entire collection
		return true
	}
	return k.id == k1.id
}

type watchInfo struct {
	ch     chan<- Change
	revno  int64
	filter func(interface{}) bool
}

type event struct {
	ch    chan<- Change
	key   watchKey
	revno int64
}

// New returns a new Watcher observing the changelog collection,
// which must be a capped collection maintained by mgo/txn.
func New(changelog *mgo.Collection) *Watcher {
	w := &Watcher{
		log:     changelog,
		watches: make(map[watchKey][]watchInfo),
		current: make(map[watchKey]int64),
		request: make(chan interface{}),
	}
	go func() {
		w.tomb.Kill(w.loop())
		w.tomb.Done()
	}()
	return w
}

// Stop stops all the watcher activities.
func (w *Watcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

// Dead returns a channel that is closed when the watcher has stopped.
func (w *Watcher) Dead() <-chan struct{} {
	return w.tomb.Dead()
}

// Err returns the error with which the watcher stopped.
// It returns nil if the watcher stopped cleanly, tomb.ErrStillAlive
// if the watcher is still running properly, or the respective error
// if the watcher is terminating or has terminated with an error.
func (w *Watcher) Err() error {
	return w.tomb.Err()
}

type reqWatch struct {
	key  watchKey
	info watchInfo
}

type reqUnwatch struct {
	key watchKey
	ch  chan<- Change
}

type reqSync struct{}

func (w *Watcher) sendReq(req interface{}) {
	select {
	case w.request <- req:
	case <-w.tomb.Dying():
	}
}

// Watch starts watching the given collection and document id.
// An event will be sent onto ch whenever a matching document's txn-revno
// field is observed to change after a transaction is applied. The revno
// parameter holds the currently known revision number for the document.
// Non-existent documents are represented by a -1 revno.
func (w *Watcher) Watch(collection string, id interface{}, revno int64, ch chan<- Change) {
	if id == nil {
		panic("watcher: cannot watch a document with nil id")
	}
	w.sendReq(reqWatch{watchKey{collection, id}, watchInfo{ch, revno, nil}})
}

// WatchCollection starts watching the given collection.
// An event will be sent onto ch whenever the txn-revno field is observed
// to change after a transaction is applied for any document in the collection.
func (w *Watcher) WatchCollection(collection string, ch chan<- Change) {
	w.WatchCollectionWithFilter(collection, ch, nil)
}

// WatchCollectionWithFilter starts watching the given collection.
// An event will be sent onto ch whenever the txn-revno field is observed
// to change after a transaction is applied for any document in the collection, so long as the
// specified filter function returns true when called with the document id value.
func (w *Watcher) WatchCollectionWithFilter(collection string, ch chan<- Change, filter func(interface{}) bool) {
	w.sendReq(reqWatch{watchKey{collection, nil}, watchInfo{ch, 0, filter}})
}

// Unwatch stops watching the given collection and document id via ch.
func (w *Watcher) Unwatch(collection string, id interface{}, ch chan<- Change) {
	if id == nil {
		panic("watcher: cannot unwatch a document with nil id")
	}
	w.sendReq(reqUnwatch{watchKey{collection, id}, ch})
}

// UnwatchCollection stops watching the given collection via ch.
func (w *Watcher) UnwatchCollection(collection string, ch chan<- Change) {
	w.sendReq(reqUnwatch{watchKey{collection, nil}, ch})
}

// StartSync forces the watcher to load new events from the database.
func (w *Watcher) StartSync() {
	w.sendReq(reqSync{})
}

// Period is the delay between each sync.
// It must not be changed when any watchers are active.
var Period time.Duration = 5 * time.Second

// loop implements the main watcher loop.
func (w *Watcher) loop() error {
	next := time.After(Period)
	w.needSync = true
	if err := w.initLastId(); err != nil {
		return err
	}
	for {
		if w.needSync {
			if err := w.sync(); err != nil {
				return err
			}
			w.flush()
			next = time.After(Period)
		}
		select {
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case <-next:
			next = time.After(Period)
			w.needSync = true
		case req := <-w.request:
			w.handle(req)
			w.flush()
		}
	}
}

// flush sends all pending events to their respective channels.
func (w *Watcher) flush() {
	// refreshEvents are stored newest first.
	for i := len(w.syncEvents) - 1; i >= 0; i-- {
		e := &w.syncEvents[i]
		for e.ch != nil {
			select {
			case <-w.tomb.Dying():
				return
			case req := <-w.request:
				w.handle(req)
				continue
			case e.ch <- Change{e.key.c, e.key.id, e.revno}:
			}
			break
		}
	}
	// requestEvents are stored oldest first, and
	// may grow during the loop.
	for i := 0; i < len(w.requestEvents); i++ {
		e := &w.requestEvents[i]
		for e.ch != nil {
			select {
			case <-w.tomb.Dying():
				return
			case req := <-w.request:
				w.handle(req)
				continue
			case e.ch <- Change{e.key.c, e.key.id, e.revno}:
			}
			break
		}
	}
	w.syncEvents = w.syncEvents[:0]
	w.requestEvents = w.requestEvents[:0]
}

// handle deals with requests delivered by the public API
// onto the background watcher goroutine.
func (w *Watcher) handle(req interface{}) {
	logger.Tracef("got request: %#v", req)
	switch r := req.(type) {
	case reqSync:
		w.needSync = true
	case reqWatch:
		for _, info := range w.watches[r.key] {
			if info.ch == r.info.ch {
				panic(fmt.Errorf("tried to re-add channel %v for %s", info.ch, r.key))
			}
		}
		if revno, ok := w.current[r.key]; ok && (revno > r.info.revno || revno == -1 && r.info.revno >= 0) {
			r.info.revno = revno
			w.requestEvents = append(w.requestEvents, event{r.info.ch, r.key, revno})
		}
		w.watches[r.key] = append(w.watches[r.key], r.info)
	case reqUnwatch:
		watches := w.watches[r.key]
		removed := false
		for i, info := range watches {
			if info.ch == r.ch {
				watches[i] = watches[len(watches)-1]
				w.watches[r.key] = watches[:len(watches)-1]
				removed = true
				break
			}
		}
		if !removed {
			panic(fmt.Errorf("tried to remove missing channel %v for %s", r.ch, r.key))
		}
		for i := range w.requestEvents {
			e := &w.requestEvents[i]
			if r.key.match(e.key) && e.ch == r.ch {
				e.ch = nil
			}
		}
		for i := range w.syncEvents {
			e := &w.syncEvents[i]
			if r.key.match(e.key) && e.ch == r.ch {
				e.ch = nil
			}
		}
	default:
		panic(fmt.Errorf("unknown request: %T", req))
	}
}

type logInfo struct {
	Docs   []interface{} `bson:"d"`
	Revnos []int64       `bson:"r"`
}

// initLastId reads the most recent changelog document and initializes
// lastId with it. This causes all history that precedes the creation
// of the watcher to be ignored.
func (w *Watcher) initLastId() error {
	var entry struct {
		Id interface{} `bson:"_id"`
	}
	session := w.log.Database.Session.Copy()
	defer session.Close()
	err := w.log.With(session).Find(nil).Sort("-$natural").One(&entry)
	if err != nil && err != mgo.ErrNotFound {
		return err
	}
	w.lastId = entry.Id
	return nil
}

// sync updates the watcher knowledge from the database, and
// queues events to observing channels.
func (w *Watcher) sync() error {
	w.needSync = false
	// Iterate through log events in reverse insertion order (newest first).
	session := w.log.Database.Session.Copy()
	defer session.Close()
	iter := w.log.With(session).Find(nil).Batch(10).Sort("-$natural").Iter()
	seen := make(map[watchKey]bool)
	first := true
	lastId := w.lastId
	var entry bson.D
	for iter.Next(&entry) {
		if len(entry) == 0 {
			logger.Tracef("got empty changelog document")
		}
		id := entry[0]
		if id.Name != "_id" {
			panic("watcher: _id field isn't first entry")
		}
		if first {
			w.lastId = id.Value
			first = false
		}
		if id.Value == lastId {
			break
		}
		logger.Tracef("got changelog document: %#v", entry)
		for _, c := range entry[1:] {
			// See txn's Runner.ChangeLog for the structure of log entries.
			var d, r []interface{}
			dr, _ := c.Value.(bson.D)
			for _, item := range dr {
				switch item.Name {
				case "d":
					d, _ = item.Value.([]interface{})
				case "r":
					r, _ = item.Value.([]interface{})
				}
			}
			if len(d) == 0 || len(d) != len(r) {
				logger.Warningf("changelog has invalid collection document: %#v", c)
				continue
			}
			for i := len(d) - 1; i >= 0; i-- {
				key := watchKey{c.Name, d[i]}
				if seen[key] {
					continue
				}
				seen[key] = true
				revno, ok := r[i].(int64)
				if !ok {
					logger.Warningf("changelog has revno with type %T: %#v", r[i], r[i])
					continue
				}
				if revno < 0 {
					revno = -1
				}
				if w.current[key] == revno {
					continue
				}
				w.current[key] = revno
				// Queue notifications for per-collection watches.
				for _, info := range w.watches[watchKey{c.Name, nil}] {
					if info.filter != nil && !info.filter(d[i]) {
						continue
					}
					w.syncEvents = append(w.syncEvents, event{info.ch, key, revno})
				}
				// Queue notifications for per-document watches.
				infos := w.watches[key]
				for i, info := range infos {
					if revno > info.revno || revno < 0 && info.revno >= 0 {
						infos[i].revno = revno
						w.syncEvents = append(w.syncEvents, event{info.ch, key, revno})
					}
				}
			}
		}
	}
	if err := iter.Close(); err != nil {
		return errors.Errorf("watcher iteration error: %v", err)
	}
	return nil
}
<commit_msg>Do not clone db sessions in state watcher<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// The watcher package provides an interface for observing changes
// to arbitrary MongoDB documents that are maintained via the
// mgo/txn transaction package.
package watcher

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.state.watcher")

// A Watcher can watch any number of collections and documents for changes.
type Watcher struct {
	tomb tomb.Tomb
	log  *mgo.Collection

	// watches holds the observers managed by Watch/Unwatch.
	watches map[watchKey][]watchInfo

	// current holds the current txn-revno values for all the observed
	// documents known to exist. Documents not observed or deleted are
	// omitted from this map and are considered to have revno -1.
	current map[watchKey]int64

	// needSync is set when a synchronization should take
	// place.
	needSync bool

	// syncEvents and requestEvents contain the events to be
	// dispatched to the watcher channels. They're queued during
	// processing and flushed at the end to simplify the algorithm.
	// The two queues are separated because events from sync are
	// handled in reverse order due to the way the algorithm works.
	syncEvents, requestEvents []event

	// request is used to deliver requests from the public API into
	// the the goroutine loop.
	request chan interface{}

	// lastId is the most recent transaction id observed by a sync.
	lastId interface{}
}

// A Change holds information about a document change.
type Change struct {
	// C and Id hold the collection name and document _id field value.
	C  string
	Id interface{}

	// Revno is the latest known value for the document's txn-revno
	// field, or -1 if the document was deleted.
	Revno int64
}

type watchKey struct {
	c  string
	id interface{} // nil when watching collection
}

func (k watchKey) String() string {
	coll := "collection " + k.c
	if k.id == nil {
		return coll
	}
	return fmt.Sprintf("document %v in %s", k.id, coll)
}

// match returns whether the receiving watch key,
// which may refer to a particular item or
// an entire collection, matches k1, which refers
// to a particular item.
func (k watchKey) match(k1 watchKey) bool {
	if k.c != k1.c {
		return false
	}
	if k.id == nil {
		// k refers to entire collection
		return true
	}
	return k.id == k1.id
}

type watchInfo struct {
	ch     chan<- Change
	revno  int64
	filter func(interface{}) bool
}

type event struct {
	ch    chan<- Change
	key   watchKey
	revno int64
}

// New returns a new Watcher observing the changelog collection,
// which must be a capped collection maintained by mgo/txn.
func New(changelog *mgo.Collection) *Watcher {
	w := &Watcher{
		log:     changelog,
		watches: make(map[watchKey][]watchInfo),
		current: make(map[watchKey]int64),
		request: make(chan interface{}),
	}
	go func() {
		w.tomb.Kill(w.loop())
		w.tomb.Done()
	}()
	return w
}

// Stop stops all the watcher activities.
func (w *Watcher) Stop() error {
	w.tomb.Kill(nil)
	return w.tomb.Wait()
}

// Dead returns a channel that is closed when the watcher has stopped.
func (w *Watcher) Dead() <-chan struct{} {
	return w.tomb.Dead()
}

// Err returns the error with which the watcher stopped.
// It returns nil if the watcher stopped cleanly, tomb.ErrStillAlive
// if the watcher is still running properly, or the respective error
// if the watcher is terminating or has terminated with an error.
func (w *Watcher) Err() error {
	return w.tomb.Err()
}

type reqWatch struct {
	key  watchKey
	info watchInfo
}

type reqUnwatch struct {
	key watchKey
	ch  chan<- Change
}

type reqSync struct{}

func (w *Watcher) sendReq(req interface{}) {
	select {
	case w.request <- req:
	case <-w.tomb.Dying():
	}
}

// Watch starts watching the given collection and document id.
// An event will be sent onto ch whenever a matching document's txn-revno
// field is observed to change after a transaction is applied. The revno
// parameter holds the currently known revision number for the document.
// Non-existent documents are represented by a -1 revno.
func (w *Watcher) Watch(collection string, id interface{}, revno int64, ch chan<- Change) {
	if id == nil {
		panic("watcher: cannot watch a document with nil id")
	}
	w.sendReq(reqWatch{watchKey{collection, id}, watchInfo{ch, revno, nil}})
}

// WatchCollection starts watching the given collection.
// An event will be sent onto ch whenever the txn-revno field is observed
// to change after a transaction is applied for any document in the collection.
func (w *Watcher) WatchCollection(collection string, ch chan<- Change) {
	w.WatchCollectionWithFilter(collection, ch, nil)
}

// WatchCollectionWithFilter starts watching the given collection.
// An event will be sent onto ch whenever the txn-revno field is observed
// to change after a transaction is applied for any document in the collection, so long as the
// specified filter function returns true when called with the document id value.
func (w *Watcher) WatchCollectionWithFilter(collection string, ch chan<- Change, filter func(interface{}) bool) {
	w.sendReq(reqWatch{watchKey{collection, nil}, watchInfo{ch, 0, filter}})
}

// Unwatch stops watching the given collection and document id via ch.
func (w *Watcher) Unwatch(collection string, id interface{}, ch chan<- Change) {
	if id == nil {
		panic("watcher: cannot unwatch a document with nil id")
	}
	w.sendReq(reqUnwatch{watchKey{collection, id}, ch})
}

// UnwatchCollection stops watching the given collection via ch.
func (w *Watcher) UnwatchCollection(collection string, ch chan<- Change) {
	w.sendReq(reqUnwatch{watchKey{collection, nil}, ch})
}

// StartSync forces the watcher to load new events from the database.
func (w *Watcher) StartSync() {
	w.sendReq(reqSync{})
}

// Period is the delay between each sync.
// It must not be changed when any watchers are active.
var Period time.Duration = 5 * time.Second

// loop implements the main watcher loop.
func (w *Watcher) loop() error {
	next := time.After(Period)
	w.needSync = true
	if err := w.initLastId(); err != nil {
		return err
	}
	for {
		if w.needSync {
			if err := w.sync(); err != nil {
				return err
			}
			w.flush()
			next = time.After(Period)
		}
		select {
		case <-w.tomb.Dying():
			return tomb.ErrDying
		case <-next:
			next = time.After(Period)
			w.needSync = true
		case req := <-w.request:
			w.handle(req)
			w.flush()
		}
	}
}

// flush sends all pending events to their respective channels.
func (w *Watcher) flush() {
	// refreshEvents are stored newest first.
	for i := len(w.syncEvents) - 1; i >= 0; i-- {
		e := &w.syncEvents[i]
		for e.ch != nil {
			select {
			case <-w.tomb.Dying():
				return
			case req := <-w.request:
				w.handle(req)
				continue
			case e.ch <- Change{e.key.c, e.key.id, e.revno}:
			}
			break
		}
	}
	// requestEvents are stored oldest first, and
	// may grow during the loop.
	for i := 0; i < len(w.requestEvents); i++ {
		e := &w.requestEvents[i]
		for e.ch != nil {
			select {
			case <-w.tomb.Dying():
				return
			case req := <-w.request:
				w.handle(req)
				continue
			case e.ch <- Change{e.key.c, e.key.id, e.revno}:
			}
			break
		}
	}
	w.syncEvents = w.syncEvents[:0]
	w.requestEvents = w.requestEvents[:0]
}

// handle deals with requests delivered by the public API
// onto the background watcher goroutine.
func (w *Watcher) handle(req interface{}) {
	logger.Tracef("got request: %#v", req)
	switch r := req.(type) {
	case reqSync:
		w.needSync = true
	case reqWatch:
		for _, info := range w.watches[r.key] {
			if info.ch == r.info.ch {
				panic(fmt.Errorf("tried to re-add channel %v for %s", info.ch, r.key))
			}
		}
		if revno, ok := w.current[r.key]; ok && (revno > r.info.revno || revno == -1 && r.info.revno >= 0) {
			r.info.revno = revno
			w.requestEvents = append(w.requestEvents, event{r.info.ch, r.key, revno})
		}
		w.watches[r.key] = append(w.watches[r.key], r.info)
	case reqUnwatch:
		watches := w.watches[r.key]
		removed := false
		for i, info := range watches {
			if info.ch == r.ch {
				watches[i] = watches[len(watches)-1]
				w.watches[r.key] = watches[:len(watches)-1]
				removed = true
				break
			}
		}
		if !removed {
			panic(fmt.Errorf("tried to remove missing channel %v for %s", r.ch, r.key))
		}
		for i := range w.requestEvents {
			e := &w.requestEvents[i]
			if r.key.match(e.key) && e.ch == r.ch {
				e.ch = nil
			}
		}
		for i := range w.syncEvents {
			e := &w.syncEvents[i]
			if r.key.match(e.key) && e.ch == r.ch {
				e.ch = nil
			}
		}
	default:
		panic(fmt.Errorf("unknown request: %T", req))
	}
}

type logInfo struct {
	Docs   []interface{} `bson:"d"`
	Revnos []int64       `bson:"r"`
}

// initLastId reads the most recent changelog document and initializes
// lastId with it. This causes all history that precedes the creation
// of the watcher to be ignored.
func (w *Watcher) initLastId() error {
	var entry struct {
		Id interface{} `bson:"_id"`
	}
	err := w.log.Find(nil).Sort("-$natural").One(&entry)
	if err != nil && err != mgo.ErrNotFound {
		return err
	}
	w.lastId = entry.Id
	return nil
}

// sync updates the watcher knowledge from the database, and
// queues events to observing channels.
func (w *Watcher) sync() error {
	w.needSync = false
	// Iterate through log events in reverse insertion order (newest first).
	iter := w.log.Find(nil).Batch(10).Sort("-$natural").Iter()
	seen := make(map[watchKey]bool)
	first := true
	lastId := w.lastId
	var entry bson.D
	for iter.Next(&entry) {
		if len(entry) == 0 {
			logger.Tracef("got empty changelog document")
		}
		id := entry[0]
		if id.Name != "_id" {
			panic("watcher: _id field isn't first entry")
		}
		if first {
			w.lastId = id.Value
			first = false
		}
		if id.Value == lastId {
			break
		}
		logger.Tracef("got changelog document: %#v", entry)
		for _, c := range entry[1:] {
			// See txn's Runner.ChangeLog for the structure of log entries.
			var d, r []interface{}
			dr, _ := c.Value.(bson.D)
			for _, item := range dr {
				switch item.Name {
				case "d":
					d, _ = item.Value.([]interface{})
				case "r":
					r, _ = item.Value.([]interface{})
				}
			}
			if len(d) == 0 || len(d) != len(r) {
				logger.Warningf("changelog has invalid collection document: %#v", c)
				continue
			}
			for i := len(d) - 1; i >= 0; i-- {
				key := watchKey{c.Name, d[i]}
				if seen[key] {
					continue
				}
				seen[key] = true
				revno, ok := r[i].(int64)
				if !ok {
					logger.Warningf("changelog has revno with type %T: %#v", r[i], r[i])
					continue
				}
				if revno < 0 {
					revno = -1
				}
				if w.current[key] == revno {
					continue
				}
				w.current[key] = revno
				// Queue notifications for per-collection watches.
				for _, info := range w.watches[watchKey{c.Name, nil}] {
					if info.filter != nil && !info.filter(d[i]) {
						continue
					}
					w.syncEvents = append(w.syncEvents, event{info.ch, key, revno})
				}
				// Queue notifications for per-document watches.
				infos := w.watches[key]
				for i, info := range infos {
					if revno > info.revno || revno < 0 && info.revno >= 0 {
						infos[i].revno = revno
						w.syncEvents = append(w.syncEvents, event{info.ch, key, revno})
					}
				}
			}
		}
	}
	if err := iter.Close(); err != nil {
		return errors.Errorf("watcher iteration error: %v", err)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

package memory

import (
	"hash/fnv"
	"runtime"
	"sync"
	"time"

	"gopkg.in/yaml.v2"

	"github.com/chihaya/chihaya"
	"github.com/chihaya/chihaya/server/store"
)

func init() {
	store.RegisterPeerStoreDriver("memory", &peerStoreDriver{})
}

type peerStoreDriver struct{}

func (d *peerStoreDriver) New(storecfg *store.Config) (store.PeerStore, error) {
	cfg, err := newPeerStoreConfig(storecfg)
	if err != nil {
		return nil, err
	}

	return &peerStore{
		shards: make([]*peerShard, cfg.Shards),
	}, nil
}

type peerStoreConfig struct {
	Shards int `yaml:"shards"`
}

func newPeerStoreConfig(storecfg *store.Config) (*peerStoreConfig, error) {
	bytes, err := yaml.Marshal(storecfg.PeerStoreConfig)
	if err != nil {
		return nil, err
	}

	var cfg peerStoreConfig
	err = yaml.Unmarshal(bytes, &cfg)
	if err != nil {
		return nil, err
	}

	return &cfg, nil
}

type peer struct {
	chihaya.Peer
	LastAction time.Time
}

type peerShard struct {
	peers map[string]map[string]peer
	sync.RWMutex
}

type peerStore struct {
	shards []*peerShard
}

var _ store.PeerStore = &peerStore{}

func (s *peerStore) shardIndex(infoHash chihaya.InfoHash) uint32 {
	idx := fnv.New32()
	idx.Write([]byte(infoHash))
	return idx.Sum32() % uint32(len(s.shards))
}

func peerKey(p chihaya.Peer) string {
	return string(p.IP) + string(p.ID)
}

func seedersKey(infoHash chihaya.InfoHash) string {
	return string(infoHash) + "-s"
}

func leechersKey(infoHash chihaya.InfoHash) string {
	return string(infoHash) + "-l"
}

func (s *peerStore) PutSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	key := seedersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[key] == nil {
		shard.peers[key] = make(map[string]peer)
	}

	shard.peers[key][peerKey(p)] = peer{
		Peer:       p,
		LastAction: time.Now(),
	}

	return nil
}

func (s *peerStore) DeleteSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	key := seedersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[key] == nil {
		return nil
	}

	delete(shard.peers[key], peerKey(p))

	if len(shard.peers[key]) == 0 {
		delete(shard.peers, key)
	}

	return nil
}

func (s *peerStore) PutLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	key := leechersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[key] == nil {
		shard.peers[key] = make(map[string]peer)
	}

	shard.peers[key][peerKey(p)] = peer{
		Peer:       p,
		LastAction: time.Now(),
	}

	return nil
}

func (s *peerStore) DeleteLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	key := leechersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[key] == nil {
		return nil
	}

	delete(shard.peers[key], peerKey(p))

	if len(shard.peers[key]) == 0 {
		delete(shard.peers, key)
	}

	return nil
}

func (s *peerStore) GraduateLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	lkey := leechersKey(infoHash)
	skey := seedersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[lkey] != nil {
		delete(shard.peers[lkey], peerKey(p))
	}

	if shard.peers[skey] == nil {
		shard.peers[skey] = make(map[string]peer)
	}

	shard.peers[skey][peerKey(p)] = peer{
		Peer:       p,
		LastAction: time.Now(),
	}

	return nil
}

func (s *peerStore) CollectGarbage(cutoff time.Time) error {
	for _, shard := range s.shards {
		shard.RLock()
		var keys []string
		for key := range shard.peers {
			keys = append(keys, key)
		}
		shard.RUnlock()
		runtime.Gosched()

		for _, key := range keys {
			shard.Lock()

			for peerKey, p := range shard.peers[key] {
				if p.LastAction.Before(cutoff) {
					delete(shard.peers[key], peerKey)
				}
			}

			if len(shard.peers[key] == 0) {
				delete(shard.peers, key)
			}

			shard.Unlock()
			runtime.Gosched()
		}

		runtime.Gosched()
	}

	return nil
}

func (s *peerStore) AnnouncePeers(infoHash chihaya.InfoHash, seeder bool, numWant int) (peers, peers6 []chihaya.Peer, err error) {
	lkey := leechersKey(infoHash)
	skey := seedersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.RLock()
	defer shard.RUnlock()

	if seeder {
		// Append leechers as possible.
		leechers := shard.peers[lkey]
		for _, p := range leechers {
			if numWant == 0 {
				break
			}

			if p.IP.To4() == nil {
				peers6 = append(peers6, p.Peer)
			} else {
				peers = append(peers, p.Peer)
			}
			numWant--
		}
	} else {
		// Append as many seeders as possible.
		seeders := shard.peers[skey]
		for _, p := range seeders {
			if numWant == 0 {
				break
			}

			if p.IP.To4() == nil {
				peers6 = append(peers6, p.Peer)
			} else {
				peers = append(peers, p.Peer)
			}
			numWant--
		}

		// Append leechers until we reach numWant.
		leechers := shard.peers[lkey]
		if numWant > 0 {
			for _, p := range leechers {
				if numWant == 0 {
					break
				}

				if p.IP.To4() == nil {
					peers6 = append(peers6, p.Peer)
				} else {
					peers = append(peers, p.Peer)
				}
				numWant--
			}
		}
	}

	return
}
<commit_msg>memory: fix paren bug<commit_after>// Copyright 2016 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

package memory

import (
	"hash/fnv"
	"runtime"
	"sync"
	"time"

	"gopkg.in/yaml.v2"

	"github.com/chihaya/chihaya"
	"github.com/chihaya/chihaya/server/store"
)

func init() {
	store.RegisterPeerStoreDriver("memory", &peerStoreDriver{})
}

type peerStoreDriver struct{}

func (d *peerStoreDriver) New(storecfg *store.Config) (store.PeerStore, error) {
	cfg, err := newPeerStoreConfig(storecfg)
	if err != nil {
		return nil, err
	}

	return &peerStore{
		shards: make([]*peerShard, cfg.Shards),
	}, nil
}

type peerStoreConfig struct {
	Shards int `yaml:"shards"`
}

func newPeerStoreConfig(storecfg *store.Config) (*peerStoreConfig, error) {
	bytes, err := yaml.Marshal(storecfg.PeerStoreConfig)
	if err != nil {
		return nil, err
	}

	var cfg peerStoreConfig
	err = yaml.Unmarshal(bytes, &cfg)
	if err != nil {
		return nil, err
	}

	return &cfg, nil
}

type peer struct {
	chihaya.Peer
	LastAction time.Time
}

type peerShard struct {
	peers map[string]map[string]peer
	sync.RWMutex
}

type peerStore struct {
	shards []*peerShard
}

var _ store.PeerStore = &peerStore{}

func (s *peerStore) shardIndex(infoHash chihaya.InfoHash) uint32 {
	idx := fnv.New32()
	idx.Write([]byte(infoHash))
	return idx.Sum32() % uint32(len(s.shards))
}

func peerKey(p chihaya.Peer) string {
	return string(p.IP) + string(p.ID)
}

func seedersKey(infoHash chihaya.InfoHash) string {
	return string(infoHash) + "-s"
}

func leechersKey(infoHash chihaya.InfoHash) string {
	return string(infoHash) + "-l"
}

func (s *peerStore) PutSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	key := seedersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[key] == nil {
		shard.peers[key] = make(map[string]peer)
	}

	shard.peers[key][peerKey(p)] = peer{
		Peer:       p,
		LastAction: time.Now(),
	}

	return nil
}

func (s *peerStore) DeleteSeeder(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	key := seedersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[key] == nil {
		return nil
	}

	delete(shard.peers[key], peerKey(p))

	if len(shard.peers[key]) == 0 {
		delete(shard.peers, key)
	}

	return nil
}

func (s *peerStore) PutLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	key := leechersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[key] == nil {
		shard.peers[key] = make(map[string]peer)
	}

	shard.peers[key][peerKey(p)] = peer{
		Peer:       p,
		LastAction: time.Now(),
	}

	return nil
}

func (s *peerStore) DeleteLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	key := leechersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[key] == nil {
		return nil
	}

	delete(shard.peers[key], peerKey(p))

	if len(shard.peers[key]) == 0 {
		delete(shard.peers, key)
	}

	return nil
}

func (s *peerStore) GraduateLeecher(infoHash chihaya.InfoHash, p chihaya.Peer) error {
	lkey := leechersKey(infoHash)
	skey := seedersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.Lock()
	defer shard.Unlock()

	if shard.peers[lkey] != nil {
		delete(shard.peers[lkey], peerKey(p))
	}

	if shard.peers[skey] == nil {
		shard.peers[skey] = make(map[string]peer)
	}

	shard.peers[skey][peerKey(p)] = peer{
		Peer:       p,
		LastAction: time.Now(),
	}

	return nil
}

func (s *peerStore) CollectGarbage(cutoff time.Time) error {
	for _, shard := range s.shards {
		shard.RLock()
		var keys []string
		for key := range shard.peers {
			keys = append(keys, key)
		}
		shard.RUnlock()
		runtime.Gosched()

		for _, key := range keys {
			shard.Lock()

			for peerKey, p := range shard.peers[key] {
				if p.LastAction.Before(cutoff) {
					delete(shard.peers[key], peerKey)
				}
			}

			if len(shard.peers[key]) == 0 {
				delete(shard.peers, key)
			}

			shard.Unlock()
			runtime.Gosched()
		}

		runtime.Gosched()
	}

	return nil
}

func (s *peerStore) AnnouncePeers(infoHash chihaya.InfoHash, seeder bool, numWant int) (peers, peers6 []chihaya.Peer, err error) {
	lkey := leechersKey(infoHash)
	skey := seedersKey(infoHash)

	shard := s.shards[s.shardIndex(infoHash)]
	shard.RLock()
	defer shard.RUnlock()

	if seeder {
		// Append leechers as possible.
		leechers := shard.peers[lkey]
		for _, p := range leechers {
			if numWant == 0 {
				break
			}

			if p.IP.To4() == nil {
				peers6 = append(peers6, p.Peer)
			} else {
				peers = append(peers, p.Peer)
			}
			numWant--
		}
	} else {
		// Append as many seeders as possible.
		seeders := shard.peers[skey]
		for _, p := range seeders {
			if numWant == 0 {
				break
			}

			if p.IP.To4() == nil {
				peers6 = append(peers6, p.Peer)
			} else {
				peers = append(peers, p.Peer)
			}
			numWant--
		}

		// Append leechers until we reach numWant.
		leechers := shard.peers[lkey]
		if numWant > 0 {
			for _, p := range leechers {
				if numWant == 0 {
					break
				}

				if p.IP.To4() == nil {
					peers6 = append(peers6, p.Peer)
				} else {
					peers = append(peers, p.Peer)
				}
				numWant--
			}
		}
	}

	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2012-2016 Eli Janssen
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package statsd

import (
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"time"
)

var bufPool = newBufferPool()

// The StatSender interface wraps all the statsd metric methods
type StatSender interface {
	Inc(string, int64, float32) error
	Dec(string, int64, float32) error
	Gauge(string, int64, float32) error
	GaugeDelta(string, int64, float32) error
	Timing(string, int64, float32) error
	TimingDuration(string, time.Duration, float32) error
	Set(string, string, float32) error
	SetInt(string, int64, float32) error
	Raw(string, string, float32) error
}

// The Statter interface defines the behavior of a stat client
type Statter interface {
	StatSender
	NewSubStatter(string) SubStatter
	SetPrefix(string)
	Close() error
}

// The SubStatter interface defines the behavior of a stat child/subclient
type SubStatter interface {
	StatSender
	SetSamplerFunc(SamplerFunc)
	NewSubStatter(string) SubStatter
}

// The SamplerFunc type defines a function that can serve
// as a Client sampler function.
type SamplerFunc func(float32) bool

// DefaultSampler is the default rate sampler function
func DefaultSampler(rate float32) bool {
	if rate < 1 {
		if rand.Float32() < rate {
			return true
		}
		return false
	}
	return true
}

// A Client is a statsd client.
type Client struct {
	// prefix for statsd name
	prefix string
	// packet sender
	sender Sender
	// sampler method
	sampler SamplerFunc
}

// Close closes the connection and cleans up.
func (s *Client) Close() error {
	if s == nil {
		return nil
	}

	err := s.sender.Close()
	return err
}

// Inc increments a statsd count type.
// stat is a string name for the metric.
// value is the integer value
// rate is the sample rate (0.0 to 1.0)
func (s *Client) Inc(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|c", rate)
}

// Dec decrements a statsd count type.
// stat is a string name for the metric.
// value is the integer value.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Dec(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", -value, "|c", rate)
}

// Gauge submits/updates a statsd gauge type.
// stat is a string name for the metric.
// value is the integer value.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Gauge(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|g", rate)
}

// GaugeDelta submits a delta to a statsd gauge.
// stat is the string name for the metric.
// value is the (positive or negative) change.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) GaugeDelta(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	// if negative, the submit formatter will prefix with a - already
	// so only special case the positive value
	if value >= 0 {
		return s.submit(stat, "+", value, "|g", rate)
	}
	return s.submit(stat, "", value, "|g", rate)
}

// Timing submits a statsd timing type.
// stat is a string name for the metric.
// delta is the time duration value in milliseconds
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Timing(stat string, delta int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", delta, "|ms", rate)
}

// TimingDuration submits a statsd timing type.
// stat is a string name for the metric.
// delta is the timing value as time.Duration
// rate is the sample rate (0.0 to 1.0).
func (s *Client) TimingDuration(stat string, delta time.Duration, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	ms := float64(delta) / float64(time.Millisecond)
	return s.submit(stat, "", ms, "|ms", rate)
}

// Set submits a stats set type
// stat is a string name for the metric.
// value is the string value
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Set(stat string, value string, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|s", rate)
}

// SetInt submits a number as a stats set type.
// stat is a string name for the metric.
// value is the integer value
// rate is the sample rate (0.0 to 1.0).
func (s *Client) SetInt(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|s", rate)
}

// Raw submits a preformatted value.
// stat is the string name for the metric.
// value is a preformatted "raw" value string.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Raw(stat string, value string, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "", rate)
}

// SetSamplerFunc sets a sampler function to something other than the default
// sampler is a function that determines whether the metric is
// to be accepted, or discarded.
// An example use case is for submitted pre-sampled metrics.
func (s *Client) SetSamplerFunc(sampler SamplerFunc) {
	s.sampler = sampler
}

// submit an already sampled raw stat
func (s *Client) submit(stat, vprefix string, value interface{}, suffix string, rate float32) error {
	data := bufPool.Get()
	defer bufPool.Put(data)

	if s.prefix != "" {
		data.WriteString(s.prefix)
		data.WriteString(".")
	}

	data.WriteString(stat)
	data.WriteString(":")

	if vprefix != "" {
		data.WriteString(vprefix)
	}

	// sadly, no way to jam this back into the bytes.Buffer without
	// doing a few allocations... avoiding those is the whole point here...
	// so from here on out just use it as a raw []byte
	b := data.Bytes()

	switch v := value.(type) {
	case string:
		b = append(b, v...)
	case int64:
		b = strconv.AppendInt(b, v, 10)
	case float64:
		b = strconv.AppendFloat(b, v, 'f', -1, 64)
	default:
		return fmt.Errorf("No matching type format")
	}

	if suffix != "" {
		b = append(b, suffix...)
	}

	if rate < 1 {
		b = append(b, "|@"...)
		b = strconv.AppendFloat(b, float64(rate), 'f', 6, 32)
	}

	_, err := s.sender.Send(b)
	return err
}

// check for nil client, and perform sampling calculation
func (s *Client) includeStat(rate float32) bool {
	if s == nil {
		return false
	}

	// test for nil in case someone builds their own
	// client without calling new (result is nil sampler)
	if s.sampler != nil {
		return s.sampler(rate)
	}
	return DefaultSampler(rate)
}

// SetPrefix sets/updates the statsd client prefix.
// Note: Does not change the prefix of any SubStatters.
func (s *Client) SetPrefix(prefix string) {
	if s == nil {
		return
	}

	s.prefix = prefix
}

// NewSubStatter returns a SubStatter with appended prefix
func (s *Client) NewSubStatter(prefix string) SubStatter {
	var c *Client
	if s != nil {
		c = &Client{
			prefix:  joinPathComp(s.prefix, prefix),
			sender:  s.sender,
			sampler: s.sampler,
		}
	}
	return c
}

// NewClient returns a pointer to a new Client, and an error.
//
// addr is a string of the format "hostname:port", and must be parsable by
// net.ResolveUDPAddr.
//
// prefix is the statsd client prefix. Can be "" if no prefix is desired.
func NewClient(addr, prefix string) (Statter, error) {
	sender, err := NewSimpleSender(addr)
	if err != nil {
		return nil, err
	}

	return NewClientWithSender(sender, prefix)
}

// NewClientWithSender returns a pointer to a new Client and an error.
//
// sender is an instance of a statsd.Sender interface and may not be nil
//
// prefix is the stastd client prefix. Can be "" if no prefix is desired.
func NewClientWithSender(sender Sender, prefix string) (Statter, error) {
	if sender == nil {
		return nil, fmt.Errorf("Client sender may not be nil")
	}

	return &Client{prefix: prefix, sender: sender}, nil
}

// joinPathComp is a helper that ensures we combine path components with a dot
// when it's appropriate to do so; prefix is the existing prefix and suffix is
// the new component being added.
//
// It returns the joined prefix.
func joinPathComp(prefix, suffix string) string {
	suffix = strings.TrimLeft(suffix, ".")
	if prefix != "" && suffix != "" {
		return prefix + "." + suffix
	}
	return prefix + suffix
}

// Dial is a compatibility alias for NewClient
var Dial = NewClient

// New is a compatibility alias for NewClient
var New = NewClient
<commit_msg>avoid method call overhead here<commit_after>// Copyright (c) 2012-2016 Eli Janssen
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package statsd

import (
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"time"
)

var bufPool = newBufferPool()

// The StatSender interface wraps all the statsd metric methods
type StatSender interface {
	Inc(string, int64, float32) error
	Dec(string, int64, float32) error
	Gauge(string, int64, float32) error
	GaugeDelta(string, int64, float32) error
	Timing(string, int64, float32) error
	TimingDuration(string, time.Duration, float32) error
	Set(string, string, float32) error
	SetInt(string, int64, float32) error
	Raw(string, string, float32) error
}

// The Statter interface defines the behavior of a stat client
type Statter interface {
	StatSender
	NewSubStatter(string) SubStatter
	SetPrefix(string)
	Close() error
}

// The SubStatter interface defines the behavior of a stat child/subclient
type SubStatter interface {
	StatSender
	SetSamplerFunc(SamplerFunc)
	NewSubStatter(string) SubStatter
}

// The SamplerFunc type defines a function that can serve
// as a Client sampler function.
type SamplerFunc func(float32) bool

// DefaultSampler is the default rate sampler function
func DefaultSampler(rate float32) bool {
	if rate < 1 {
		if rand.Float32() < rate {
			return true
		}
		return false
	}
	return true
}

// A Client is a statsd client.
type Client struct {
	// prefix for statsd name
	prefix string
	// packet sender
	sender Sender
	// sampler method
	sampler SamplerFunc
}

// Close closes the connection and cleans up.
func (s *Client) Close() error {
	if s == nil {
		return nil
	}

	err := s.sender.Close()
	return err
}

// Inc increments a statsd count type.
// stat is a string name for the metric.
// value is the integer value
// rate is the sample rate (0.0 to 1.0)
func (s *Client) Inc(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|c", rate)
}

// Dec decrements a statsd count type.
// stat is a string name for the metric.
// value is the integer value.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Dec(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", -value, "|c", rate)
}

// Gauge submits/updates a statsd gauge type.
// stat is a string name for the metric.
// value is the integer value.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Gauge(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|g", rate)
}

// GaugeDelta submits a delta to a statsd gauge.
// stat is the string name for the metric.
// value is the (positive or negative) change.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) GaugeDelta(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	// if negative, the submit formatter will prefix with a - already
	// so only special case the positive value
	if value >= 0 {
		return s.submit(stat, "+", value, "|g", rate)
	}
	return s.submit(stat, "", value, "|g", rate)
}

// Timing submits a statsd timing type.
// stat is a string name for the metric.
// delta is the time duration value in milliseconds
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Timing(stat string, delta int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", delta, "|ms", rate)
}

// TimingDuration submits a statsd timing type.
// stat is a string name for the metric.
// delta is the timing value as time.Duration
// rate is the sample rate (0.0 to 1.0).
func (s *Client) TimingDuration(stat string, delta time.Duration, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	ms := float64(delta) / float64(time.Millisecond)
	return s.submit(stat, "", ms, "|ms", rate)
}

// Set submits a stats set type
// stat is a string name for the metric.
// value is the string value
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Set(stat string, value string, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|s", rate)
}

// SetInt submits a number as a stats set type.
// stat is a string name for the metric.
// value is the integer value
// rate is the sample rate (0.0 to 1.0).
func (s *Client) SetInt(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|s", rate)
}

// Raw submits a preformatted value.
// stat is the string name for the metric.
// value is a preformatted "raw" value string.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Raw(stat string, value string, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "", rate)
}

// SetSamplerFunc sets a sampler function to something other than the default
// sampler is a function that determines whether the metric is
// to be accepted, or discarded.
// An example use case is for submitted pre-sampled metrics.
func (s *Client) SetSamplerFunc(sampler SamplerFunc) {
	s.sampler = sampler
}

// submit an already sampled raw stat
func (s *Client) submit(stat, vprefix string, value interface{}, suffix string, rate float32) error {
	data := bufPool.Get()
	defer bufPool.Put(data)

	if s.prefix != "" {
		data.WriteString(s.prefix)
		data.WriteString(".")
	}

	data.WriteString(stat)
	data.WriteString(":")

	if vprefix != "" {
		data.WriteString(vprefix)
	}

	// sadly, no way to jam this back into the bytes.Buffer without
	// doing a few allocations... avoiding those is the whole point here...
	// so from here on out just use it as a raw []byte
	b := data.Bytes()

	switch v := value.(type) {
	case string:
		b = append(b, v...)
	case int64:
		b = strconv.AppendInt(b, v, 10)
	case float64:
		b = strconv.AppendFloat(b, v, 'f', -1, 64)
	default:
		return fmt.Errorf("No matching type format")
	}

	if suffix != "" {
		b = append(b, suffix...)
	}

	if rate < 1 {
		b = append(b, "|@"...)
		b = strconv.AppendFloat(b, float64(rate), 'f', 6, 32)
	}

	_, err := s.sender.Send(b)
	return err
}

// check for nil client, and perform sampling calculation
func (s *Client) includeStat(rate float32) bool {
	if s == nil {
		return false
	}

	// test for nil in case someone builds their own
	// client without calling new (result is nil sampler)
	if s.sampler != nil {
		return s.sampler(rate)
	}
	return DefaultSampler(rate)
}

// SetPrefix sets/updates the statsd client prefix.
// Note: Does not change the prefix of any SubStatters.
func (s *Client) SetPrefix(prefix string) {
	if s == nil {
		return
	}

	s.prefix = prefix
}

// NewSubStatter returns a SubStatter with appended prefix
func (s *Client) NewSubStatter(prefix string) SubStatter {
	var c *Client
	if s != nil {
		c = &Client{
			prefix:  joinPathComp(s.prefix, prefix),
			sender:  s.sender,
			sampler: s.sampler,
		}
	}
	return c
}

// NewClient returns a pointer to a new Client, and an error.
//
// addr is a string of the format "hostname:port", and must be parsable by
// net.ResolveUDPAddr.
//
// prefix is the statsd client prefix. Can be "" if no prefix is desired.
func NewClient(addr, prefix string) (Statter, error) {
	sender, err := NewSimpleSender(addr)
	if err != nil {
		return nil, err
	}

	return &Client{prefix: prefix, sender: sender}, nil
}

// NewClientWithSender returns a pointer to a new Client and an error.
//
// sender is an instance of a statsd.Sender interface and may not be nil
//
// prefix is the stastd client prefix. Can be "" if no prefix is desired.
func NewClientWithSender(sender Sender, prefix string) (Statter, error) {
	if sender == nil {
		return nil, fmt.Errorf("Client sender may not be nil")
	}

	return &Client{prefix: prefix, sender: sender}, nil
}

// joinPathComp is a helper that ensures we combine path components with a dot
// when it's appropriate to do so; prefix is the existing prefix and suffix is
// the new component being added.
//
// It returns the joined prefix.
func joinPathComp(prefix, suffix string) string {
	suffix = strings.TrimLeft(suffix, ".")
	if prefix != "" && suffix != "" {
		return prefix + "." + suffix
	}
	return prefix + suffix
}

// Dial is a compatibility alias for NewClient
var Dial = NewClient

// New is a compatibility alias for NewClient
var New = NewClient
<|endoftext|>


<commit_before>package statsd

import (
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"time"
)

var bufPool = newBufferPool()

type StatSender interface {
	Inc(string, int64, float32) error
	Dec(string, int64, float32) error
	Gauge(string, int64, float32) error
	GaugeDelta(string, int64, float32) error
	Timing(string, int64, float32) error
	TimingDuration(string, time.Duration, float32) error
	Set(string, string, float32) error
	SetInt(string, int64, float32) error
	Raw(string, string, float32) error
}

type Statter interface {
	StatSender
	NewSubStatter(string) SubStatter
	SetPrefix(string)
	Close() error
}

type SubStatter interface {
	StatSender
	NewSubStatter(string) SubStatter
}

type Client struct {
	// prefix for statsd name
	prefix string
	// packet sender
	sender Sender
}

// Close closes the connection and cleans up.
func (s *Client) Close() error {
	if s == nil {
		return nil
	}

	err := s.sender.Close()
	return err
}

// Increments a statsd count type.
// stat is a string name for the metric.
// value is the integer value
// rate is the sample rate (0.0 to 1.0)
func (s *Client) Inc(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|c", rate)
}

// Decrements a statsd count type.
// stat is a string name for the metric.
// value is the integer value.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Dec(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", -value, "|c", rate)
}

// Submits/Updates a statsd gauge type.
// stat is a string name for the metric.
// value is the integer value.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Gauge(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|g", rate)
}

// Submits a delta to a statsd gauge.
// stat is the string name for the metric.
// value is the (positive or negative) change.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) GaugeDelta(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	// if negative, the submit formatter will prefix with a - already
	// so in that case, use an empty prefix here.
	symbol := ""
	if value >= 0 {
		symbol = "+"
	}
	return s.submit(stat, symbol, value, "|g", rate)
}

// Submits a statsd timing type.
// stat is a string name for the metric.
// delta is the time duration value in milliseconds
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Timing(stat string, delta int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", delta, "|ms", rate)
}

// Submits a statsd timing type.
// stat is a string name for the metric.
// delta is the timing value as time.Duration
// rate is the sample rate (0.0 to 1.0).
func (s *Client) TimingDuration(stat string, delta time.Duration, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	ms := float64(delta) / float64(time.Millisecond)
	return s.submit(stat, "", ms, "|ms", rate)
}

// Submits a stats set type
// stat is a string name for the metric.
// value is the string value
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Set(stat string, value string, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|s", rate)
}

// Submits a number as a stats set type.
// stat is a string name for the metric.
// value is the integer value
// rate is the sample rate (0.0 to 1.0).
func (s *Client) SetInt(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|s", rate)
}

// Raw submits a preformatted value.
// stat is the string name for the metric.
// value is a preformatted "raw" value string.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Raw(stat string, value string, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "", rate)
}

// submit an already sampled raw stat
func (s *Client) submit(stat, vprefix string, value interface{}, suffix string, rate float32) error {
	data := bufPool.Get()
	defer bufPool.Put(data)

	if s.prefix != "" {
		data.WriteString(s.prefix)
		data.WriteString(".")
	}

	data.WriteString(stat)
	data.WriteString(":")

	if vprefix != "" {
		data.WriteString(vprefix)
	}

	// sadly, no way to jam this back into the bytes.Buffer without
	// doing a few allocations... avoiding those is the whole point here...
	// so from here on out just use it as a raw []byte
	b := data.Bytes()

	switch v := value.(type) {
	case string:
		b = append(b, v...)
	case int64:
		b = strconv.AppendInt(b, v, 10)
	case float64:
		b = strconv.AppendFloat(b, v, 'f', -1, 64)
	default:
		return fmt.Errorf("No matching type format")
	}

	if suffix != "" {
		b = append(b, suffix...)
	}

	if rate < 1 {
		b = append(b, "|@"...)
		b = strconv.AppendFloat(b, float64(rate), 'f', 6, 32)
	}

	_, err := s.sender.Send(b)
	return err
}

// check for nil client, and perform sampling calculation
func (s *Client) includeStat(rate float32) bool {
	if s == nil {
		return false
	}

	if rate < 1 {
		if rand.Float32() < rate {
			return true
		}
		return false
	}
	return true
}

// Sets/Updates the statsd client prefix.
// Note: Does not change the prefix of any SubStatters.
func (s *Client) SetPrefix(prefix string) {
	if s == nil {
		return
	}

	s.prefix = prefix
}

// Returns a SubStatter with appended prefix
func (s *Client) NewSubStatter(prefix string) SubStatter {
	var c *Client
	if s != nil {
		subfix := dotprefix(prefix)

		c = &Client{
			prefix: s.prefix + subfix,
			sender: s.sender,
		}
	}
	return c
}

// Returns a pointer to a new Client, and an error.
//
// addr is a string of the format "hostname:port", and must be parsable by
// net.ResolveUDPAddr.
//
// prefix is the statsd client prefix. Can be "" if no prefix is desired.
func NewClient(addr, prefix string) (Statter, error) {
	sender, err := NewSimpleSender(addr)
	if err != nil {
		return nil, err
	}

	client := &Client{
		prefix: prefix,
		sender: sender,
	}

	return client, nil
}

// helper method to ensure a dot is added only when necessary
func dotprefix(prefix string) string {
	if prefix != "" && !strings.HasPrefix(prefix, ".") {
		return "." + prefix
	}

	return prefix
}

// Compatibility aliases
var Dial = New
var New = NewClient
<commit_msg>dont create string<commit_after>package statsd

import (
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"time"
)

var bufPool = newBufferPool()

type StatSender interface {
	Inc(string, int64, float32) error
	Dec(string, int64, float32) error
	Gauge(string, int64, float32) error
	GaugeDelta(string, int64, float32) error
	Timing(string, int64, float32) error
	TimingDuration(string, time.Duration, float32) error
	Set(string, string, float32) error
	SetInt(string, int64, float32) error
	Raw(string, string, float32) error
}

type Statter interface {
	StatSender
	NewSubStatter(string) SubStatter
	SetPrefix(string)
	Close() error
}

type SubStatter interface {
	StatSender
	NewSubStatter(string) SubStatter
}

type Client struct {
	// prefix for statsd name
	prefix string
	// packet sender
	sender Sender
}

// Close closes the connection and cleans up.
func (s *Client) Close() error {
	if s == nil {
		return nil
	}

	err := s.sender.Close()
	return err
}

// Increments a statsd count type.
// stat is a string name for the metric.
// value is the integer value
// rate is the sample rate (0.0 to 1.0)
func (s *Client) Inc(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|c", rate)
}

// Decrements a statsd count type.
// stat is a string name for the metric.
// value is the integer value.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Dec(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", -value, "|c", rate)
}

// Submits/Updates a statsd gauge type.
// stat is a string name for the metric.
// value is the integer value.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Gauge(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|g", rate)
}

// Submits a delta to a statsd gauge.
// stat is the string name for the metric.
// value is the (positive or negative) change.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) GaugeDelta(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	// if negative, the submit formatter will prefix with a - already
	// so only special case the positive value
	if value >= 0 {
		return s.submit(stat, "+", value, "|g", rate)
	}
	return s.submit(stat, "", value, "|g", rate)
}

// Submits a statsd timing type.
// stat is a string name for the metric.
// delta is the time duration value in milliseconds
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Timing(stat string, delta int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", delta, "|ms", rate)
}

// Submits a statsd timing type.
// stat is a string name for the metric.
// delta is the timing value as time.Duration
// rate is the sample rate (0.0 to 1.0).
func (s *Client) TimingDuration(stat string, delta time.Duration, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	ms := float64(delta) / float64(time.Millisecond)
	return s.submit(stat, "", ms, "|ms", rate)
}

// Submits a stats set type
// stat is a string name for the metric.
// value is the string value
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Set(stat string, value string, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|s", rate)
}

// Submits a number as a stats set type.
// stat is a string name for the metric.
// value is the integer value
// rate is the sample rate (0.0 to 1.0).
func (s *Client) SetInt(stat string, value int64, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "|s", rate)
}

// Raw submits a preformatted value.
// stat is the string name for the metric.
// value is a preformatted "raw" value string.
// rate is the sample rate (0.0 to 1.0).
func (s *Client) Raw(stat string, value string, rate float32) error {
	if !s.includeStat(rate) {
		return nil
	}

	return s.submit(stat, "", value, "", rate)
}

// submit an already sampled raw stat
func (s *Client) submit(stat, vprefix string, value interface{}, suffix string, rate float32) error {
	data := bufPool.Get()
	defer bufPool.Put(data)

	if s.prefix != "" {
		data.WriteString(s.prefix)
		data.WriteString(".")
	}

	data.WriteString(stat)
	data.WriteString(":")

	if vprefix != "" {
		data.WriteString(vprefix)
	}

	// sadly, no way to jam this back into the bytes.Buffer without
	// doing a few allocations... avoiding those is the whole point here...
	// so from here on out just use it as a raw []byte
	b := data.Bytes()

	switch v := value.(type) {
	case string:
		b = append(b, v...)
	case int64:
		b = strconv.AppendInt(b, v, 10)
	case float64:
		b = strconv.AppendFloat(b, v, 'f', -1, 64)
	default:
		return fmt.Errorf("No matching type format")
	}

	if suffix != "" {
		b = append(b, suffix...)
	}

	if rate < 1 {
		b = append(b, "|@"...)
		b = strconv.AppendFloat(b, float64(rate), 'f', 6, 32)
	}

	_, err := s.sender.Send(b)
	return err
}

// check for nil client, and perform sampling calculation
func (s *Client) includeStat(rate float32) bool {
	if s == nil {
		return false
	}

	if rate < 1 {
		if rand.Float32() < rate {
			return true
		}
		return false
	}
	return true
}

// Sets/Updates the statsd client prefix.
// Note: Does not change the prefix of any SubStatters.
func (s *Client) SetPrefix(prefix string) {
	if s == nil {
		return
	}

	s.prefix = prefix
}

// Returns a SubStatter with appended prefix
func (s *Client) NewSubStatter(prefix string) SubStatter {
	var c *Client
	if s != nil {
		subfix := dotprefix(prefix)

		c = &Client{
			prefix: s.prefix + subfix,
			sender: s.sender,
		}
	}
	return c
}

// Returns a pointer to a new Client, and an error.
//
// addr is a string of the format "hostname:port", and must be parsable by
// net.ResolveUDPAddr.
//
// prefix is the statsd client prefix. Can be "" if no prefix is desired.
func NewClient(addr, prefix string) (Statter, error) {
	sender, err := NewSimpleSender(addr)
	if err != nil {
		return nil, err
	}

	client := &Client{
		prefix: prefix,
		sender: sender,
	}

	return client, nil
}

// helper method to ensure a dot is added only when necessary
func dotprefix(prefix string) string {
	if prefix != "" && !strings.HasPrefix(prefix, ".") {
		return "." + prefix
	}

	return prefix
}

// Compatibility aliases
var Dial = New
var New = NewClient
<|endoftext|>


<commit_before>package main

import (
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"strings"
	"text/template"
)

var pattern = "@struct2schema"

// SchemaInfo - saves table infos
type SchemaInfo struct {
	TableName string
	Fields    []SchemaField
	LastIdx   int
}

// SchemaField - saves schema field
type SchemaField struct {
	Name      string
	ValueType string
}

func processFile(inputPath string, templateStr string) {
	log.Printf("Processing file %s", inputPath)
	fset := token.NewFileSet() // positions are relative to fset
	f, err := parser.ParseFile(fset, inputPath, nil, parser.ParseComments)
	if err != nil {
		panic(err)
	}

	// ast.Print(fset, f)

	var schemaInfo SchemaInfo

	for _, decl := range f.Decls {
		schemaInfo.Fields = []SchemaField{}
		var ok bool

		ok = getTableInfo(decl, &schemaInfo)
		if !ok {
			continue
		}

		t := template.Must(template.New("sqlCommand").Parse(templateStr))
		err := t.Execute(os.Stdout, schemaInfo)
		if err != nil {
			log.Println("executing template:", err)
		}
	}
}

func getTableInfo(decl ast.Decl, schemaInfo *SchemaInfo) (found bool) {
	genDecl, ok := decl.(*ast.GenDecl)
	if !ok {
		return
	}
	if genDecl.Doc == nil {
		return
	}

	tableStructFound := false
	for _, comment := range genDecl.Doc.List {
		if strings.Contains(comment.Text, pattern) {
			tableStructFound = true
			break
		}
	}

	if !tableStructFound {
		return
	}

	for _, spec := range genDecl.Specs {
		switch spec.(type) {
		case *ast.TypeSpec:
			schemaInfo.TableName, found = getTableName(spec)

			if found == true {
				typeSpec := spec.(*ast.TypeSpec)
				fieldLen := 0

				switch typeSpec.Type.(type) {
				case *ast.StructType:
					structSpec := typeSpec.Type.(*ast.StructType)
					for _, elem := range structSpec.Fields.List {
						newField := SchemaField{
							Name:      elem.Names[0].Name,
							ValueType: elem.Type.(*ast.Ident).Name,
						}
						schemaInfo.Fields = append(schemaInfo.Fields, newField)
						fieldLen++
					}
				}

				schemaInfo.LastIdx = fieldLen - 1
			}
		}
	}

	if schemaInfo.TableName == "" {
		return
	}

	found = true
	return
}

func getTableName(spec ast.Spec) (tableName string, ok bool) {
	typeSpec := spec.(*ast.TypeSpec)
	if typeSpec.Name != nil {
		ok = true
		tableName = typeSpec.Name.Name
	}
	return
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("struct2schema: ")

	const sqlTemplateStr = `
CREATE TABLE IF NOT EXISTS {{.TableName}} ( {{$lastIdx := .LastIdx}} {{ range $idx, $field := .Fields }}
  {{.Name}} {{.ValueType}}{{ if ne $lastIdx $idx }}, {{end}}
{{ end }} )
`

	for _, path := range os.Args[1:] {
		processFile(path, sqlTemplateStr)
	}
}
<commit_msg>Add type as flag, and convert golang variable type to DB type<commit_after>package main

import (
	"flag"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"strings"
	"text/template"
)

var (
	pattern = "@struct2schema"
	dbType  = flag.String("dbType", "sqlite3", "Database used for the generated SQL command, available choise: mysql/sqlite3")
)

// SchemaInfo - saves table infos
type SchemaInfo struct {
	TableName string
	Fields    []SchemaField
	LastIdx   int
}

// SchemaField - saves schema field
type SchemaField struct {
	Name      string
	ValueType string
}

func processFile(inputPath string, templateStr string) {
	log.Printf("Processing file %s", inputPath)
	fset := token.NewFileSet() // positions are relative to fset
	f, err := parser.ParseFile(fset, inputPath, nil, parser.ParseComments)
	if err != nil {
		panic(err)
	}

	// ast.Print(fset, f)

	var schemaInfo SchemaInfo

	for _, decl := range f.Decls {
		schemaInfo.Fields = []SchemaField{}
		var ok bool

		ok = getTableInfo(decl, &schemaInfo)
		if !ok {
			continue
		}

		// Generate SQL command via pre-defined template string
		t := template.Must(template.New("sqlCommand").Parse(templateStr))
		err := t.Execute(os.Stdout, schemaInfo)
		if err != nil {
			log.Println("executing template:", err)
		}
	}
}

func getTableInfo(decl ast.Decl, schemaInfo *SchemaInfo) (found bool) {
	genDecl, ok := decl.(*ast.GenDecl)

	// Skip nil or error nodes
	if !ok {
		return
	}
	if genDecl.Doc == nil {
		return
	}

	// table structure should have commant before its code block, or it will not be handled by this generate
	tableStructFound := false
	for _, comment := range genDecl.Doc.List {
		if strings.Contains(comment.Text, pattern) {
			tableStructFound = true
			break
		}
	}

	if !tableStructFound {
		return
	}

	for _, spec := range genDecl.Specs {
		switch spec.(type) {
		case *ast.TypeSpec:
			schemaInfo.TableName, found = getTableName(spec)

			if found == true {
				typeSpec := spec.(*ast.TypeSpec)
				fieldLen := 0

				switch typeSpec.Type.(type) {
				case *ast.StructType:
					structSpec := typeSpec.Type.(*ast.StructType)
					for _, elem := range structSpec.Fields.List {
						newField := SchemaField{
							Name:      elem.Names[0].Name,
							ValueType: typeConvert(elem.Type.(*ast.Ident).Name),
						}

						schemaInfo.Fields = append(schemaInfo.Fields, newField)
						fieldLen++
					}
				}

				schemaInfo.LastIdx = fieldLen - 1
			}
		}
	}

	if schemaInfo.TableName == "" {
		return
	}

	found = true
	return
}

// Convert golang type to specified DB field type
// TODO: Convert to getting from file or generate tool
func typeConvert(golangFieldType string) (dbFieldType string) {
	switch golangFieldType {
	case "uint", "int":
		switch *dbType {
		case "sqlite3":
			dbFieldType = "INTEGER"
		case "mysql":
			dbFieldType = "INT"
		}
	case "uint8", "int8", "byte":
		switch *dbType {
		case "sqlite3":
			dbFieldType = "INTEGER"
		case "mysql":
			dbFieldType = "TINYINT"
		}
	case "uint16", "int16":
		switch *dbType {
		case "sqlite3":
			dbFieldType = "INTEGER"
		case "mysql":
			dbFieldType = "SMALLINT"
		}
	case "uint32", "int32", "rune":
		switch *dbType {
		case "sqlite3":
			dbFieldType = "INTEGER"
		case "mysql":
			dbFieldType = "INT"
		}
	case "uint64", "int64":
		switch *dbType {
		case "sqlite3":
			dbFieldType = "INTEGER"
		case "mysql":
			dbFieldType = "BIGINT"
		}

	case "float32", "float64":
		switch *dbType {
		case "sqlite3":
			dbFieldType = "REAL"
		case "mysql":
			dbFieldType = "FLOAT"
		}

	case "string":
		switch *dbType {
		case "sqlite3":
			dbFieldType = "TEXT"
		case "mysql":
			dbFieldType = "MEDIUMTEXT"
		}
	}
	return
}

func getTableName(spec ast.Spec) (tableName string, ok bool) {
	typeSpec := spec.(*ast.TypeSpec)
	if typeSpec.Name != nil {
		ok = true
		tableName = typeSpec.Name.Name
	}
	return
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("struct2schema: ")

	const sqlTemplateStr = `
CREATE TABLE IF NOT EXISTS {{.TableName}} ( {{$lastIdx := .LastIdx}} {{ range $idx, $field := .Fields }}
  {{.Name}} {{.ValueType}}{{ if ne $lastIdx $idx }}, {{end}}
{{ end }} )
`

	for _, path := range os.Args[1:] {
		processFile(path, sqlTemplateStr)
	}
}
<|endoftext|>


<commit_before>package geom

type Point struct {
	geom0
}

var _ T = &Point{}

func NewPoint(layout Layout) *Point {
	return NewPointFlat(layout, make([]float64, layout.Stride()))
}

func NewPointFlat(layout Layout, flatCoords []float64) *Point {
	p := new(Point)
	p.layout = layout
	p.stride = layout.Stride()
	p.flatCoords = flatCoords
	return p
}

func (p *Point) Area() float64 {
	return 0
}

func (p *Point) Clone() *Point {
	flatCoords := make([]float64, len(p.flatCoords))
	copy(flatCoords, p.flatCoords)
	return NewPointFlat(p.layout, flatCoords)
}

func (p *Point) Length() float64 {
	return 0
}

func (p *Point) MustSetCoords(coords Coord) *Point {
	Must(p.SetCoords(coords))
	return p
}

func (p *Point) SetCoords(coords Coord) (*Point, error) {
	if err := p.setCoords(coords); err != nil {
		return nil, err
	}
	return p, nil
}
<commit_msg>Add initial Point documentation<commit_after>package geom

// A Point represents a single point.
type Point struct {
	geom0
}

var _ T = &Point{}

// NewPoint allocates a new Point with layout l and all values zero.
func NewPoint(l Layout) *Point {
	return NewPointFlat(l, make([]float64, l.Stride()))
}

// NewPointFlat allocates a new Point with layout l and flat coordinates flatCoords.
func NewPointFlat(l Layout, flatCoords []float64) *Point {
	p := new(Point)
	p.layout = l
	p.stride = l.Stride()
	p.flatCoords = flatCoords
	return p
}

// Area returns p's area, i.e. zero.
func (p *Point) Area() float64 {
	return 0
}

// Clone returns a copy of p that does not alias p.
func (p *Point) Clone() *Point {
	flatCoords := make([]float64, len(p.flatCoords))
	copy(flatCoords, p.flatCoords)
	return NewPointFlat(p.layout, flatCoords)
}

// Length returns the length of p, i.e. zero.
func (p *Point) Length() float64 {
	return 0
}

// MustSetCoords is like SetCoords but panics on any error.
func (p *Point) MustSetCoords(coords Coord) *Point {
	Must(p.SetCoords(coords))
	return p
}

// SetCoords sets the coordinates of p.
func (p *Point) SetCoords(coords Coord) (*Point, error) {
	if err := p.setCoords(coords); err != nil {
		return nil, err
	}
	return p, nil
}
<|endoftext|>


<commit_before>package pouch

import "database/sql"

// A Pouch is anything which can act as a backing Storage and which
// we can query for entities.
type Pouch interface {
	Queryable
	Storage
}

// Storage is the interface implemented by anything which can
// be interacted with to store or retrieve entities which
// can be found, created, updated and deleted.
type Storage interface {
	FindAll([]Findable) error
	Find(Findable) error

	Create(Createable) error
	CreateAll([]Createable) error

	Update(Updateable) error
	UpdateAll([]Updateable) error

	Delete(Deleteable) error
	DeleteAll([]Deleteable) error
}

// A Query is a direct gateway to interact with the storage and
// retrieval of entities which is backed by a Storage system.
// A Query can also be modified to either add or remove criterions
// that are used to filter searches for entities in the backing
// Storage system.
type Query interface {
	Queryable
	Storage
	FindEntities(Findable, *[]Findable) error
}

// A Creatable entity is one that knows where it is meant to be
// explicitly stored and knows what parts of it need to be stored.
type Createable interface {
	Insertable
	Tableable
}

// A Gettable entity is one that knows what needs to be retrieved
// from an underlying Storage system in order to rebuild itself.
// It also knows where it should be retrieved from inside a Storage
// system.
type Gettable interface {
	GetFieldsFor([]string) []interface{}
	GetAllFields() ([]string, []interface{})
	Tableable
}

// A Tableable entity is one which knows where in a Storage system
// it is supposed to be stored and retrieved from.
type Tableable interface {
	Table() string
}

// A Findable entity is one that is Gettable from a Storage system
// and that knows how to identify itself in the underlying Storage.
type Findable interface {
	Identifiable
	Gettable
	FindableCopy() Findable
}

// An Insertable entity is one which knows what data from itself
// needs to be stored in an underlying storage system.
type Insertable interface {
	FieldsFor([]string) []interface{}
	InsertableFields() ([]string, []interface{})
	SetIdentifier(interface{}) error
}

// An Identifiable entity is one which knows how to find itself
// in a Storage system.
type Identifiable interface {
	IdentifiableFields() ([]string, []interface{})
}

// An Updateable entity is one which knows how to update itself
// in a Storage system.
type Updateable interface {
	Insertable
	Identifiable
	Tableable
}

// A Deleteable entity is one which knows how to delete only itself
// from a Storage system.
type Deleteable interface {
	Identifiable
	Tableable
}

// Anything that is Queryable knows how to filter queries for itself.
type Queryable interface {
	GroupBy(spec string) Query
	OrderBy(spec string) Query
	Where(frag string, val ...interface{}) Query
	Limit(lim int) Query
	Offset(off int) Query
}

// Executor is a convenience wrapper that allows both *sql.DB and
// *sql.DB to be as Pouches.
type Executor interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) *sql.Row
}
<commit_msg>Add documentation to Storage functions<commit_after>package pouch

import "database/sql"

// A Pouch is anything which can act as a backing Storage and which
// we can query for entities.
type Pouch interface {
	Queryable
	Storage
}

// Storage is the interface implemented by anything which can
// be interacted with to store or retrieve entities which
// can be found, created, updated and deleted. For specific
// information, with respect to error cases and implementation
// details, please see the implementation for your backing
// storage medium.
type Storage interface {
	// Find retrieves an entity from the backing storage medium.
	Find(Findable) error
	// FindAll takes a slice of findable entities (that contain
	// enough information to uniquely identify them in the backing
	// storage medium) and retrieves the rest of each Findable
	// entities information. It should be noted, that the Findable
	// entities to be retrieved to not have to be of the same
	// underlying Go type.
	FindAll([]Findable) error

	// Create stores a Creatable entity in the backing storage medium.
	Create(Createable) error
	// CreateAll stores all given Creatable entities in the backing
	// storage medium. It should be noted that, similarly to FindAll,
	// the underlying types do not need to be the same. HOWEVER,
	// this does mean that for certain backing storage media (i.e. SQL)
	// this function will have to run multiple queries. If you want a bulk
	// insert, please see (YET_TO_BE_IMPLEMENTED).
	CreateAll([]Createable) error

	// Update identifies the Updateable entity in the backing storage
	// media, and updates it the given information that Updateable
	// entity currently contains. Whether or not the provided
	// functionality is a full update or a partial update, is up to
	// underlying Pouch implementation.
	Update(Updateable) error
	// UpdateAll updates all of the provided entities in the backing
	// storage media. Like [Find/Create]All, the underlying types do not
	// have to be the same and as such this function is not meant to
	// guarantee that this is a single interaction for certain media
	// (i.e. SQL based systems). However, this can be (read should be)
	// specified to run inside a transaction.
	UpdateAll([]Updateable) error

	// Delete deletes the given entity from a backing storage media.
	Delete(Deleteable) error
	// DeleteAll removes all given entities from the backing storage media.
	DeleteAll([]Deleteable) error
}

// A Query is a direct gateway to interact with the storage and
// retrieval of entities which is backed by a storage system.
// A Query can also be modified to either add or remove criterions
// that are used to filter searches for entities in the backing
// storage system.
type Query interface {
	Queryable
	Storage

	// FindEntities retrieves all entities that satisfy the Query's
	// current criterions. This can thus also be used to retrieve
	// all entities of a given type, to implement pagination, etc.
	FindEntities(Findable, *[]Findable) error
}

// A Creatable entity is one that knows where it is meant to be
// explicitly stored and knows what parts of it need to be stored.
type Createable interface {
	Insertable
	Tableable
}

// A Gettable entity is one that knows what needs to be retrieved
// from an underlying Storage system in order to rebuild itself.
// It also knows where it should be retrieved from inside a Storage
// system.
type Gettable interface {
	GetFieldsFor([]string) []interface{}
	GetAllFields() ([]string, []interface{})
	Tableable
}

// A Tableable entity is one which knows where in a Storage system
// it is supposed to be stored and retrieved from.
type Tableable interface {
	Table() string
}

// A Findable entity is one that is Gettable from a Storage system
// and that knows how to identify itself in the underlying Storage.
type Findable interface {
	Identifiable
	Gettable
	FindableCopy() Findable
}

// An Insertable entity is one which knows what data from itself
// needs to be stored in an underlying storage system.
type Insertable interface {
	FieldsFor([]string) []interface{}
	InsertableFields() ([]string, []interface{})
	SetIdentifier(interface{}) error
}

// An Identifiable entity is one which knows how to find itself
// in a Storage system.
type Identifiable interface {
	IdentifiableFields() ([]string, []interface{})
}

// An Updateable entity is one which knows how to update itself
// in a Storage system.
type Updateable interface {
	Insertable
	Identifiable
	Tableable
}

// A Deleteable entity is one which knows how to delete only itself
// from a Storage system.
type Deleteable interface {
	Identifiable
	Tableable
}

// Anything that is Queryable knows how to filter queries for itself.
type Queryable interface {
	GroupBy(spec string) Query
	OrderBy(spec string) Query
	Where(frag string, val ...interface{}) Query
	Limit(lim int) Query
	Offset(off int) Query
}

// Executor is a convenience wrapper that allows both *sql.DB and
// *sql.DB to be as Pouches.
type Executor interface {
	Exec(query string, args ...interface{}) (sql.Result, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) *sql.Row
}
<|endoftext|>


<commit_before>package ginpprof

import (
	"net/http/pprof"
	"strings"

	"github.com/gin-gonic/gin"
)

// Wrap adds several routes from package `net/http/pprof` to *gin.Engine object
func Wrap(router *gin.Engine) {
	WrapGroup(&router.RouterGroup)
}

// WrapGroup adds several routes from package `net/http/pprof` to *gin.RouterGroup object
func WrapGroup(router *gin.RouterGroup) {
	routers := []struct {
		Method  string
		Path    string
		Handler gin.HandlerFunc
	}{
		{"GET", "/debug/pprof/", IndexHandler()},
		{"GET", "/debug/pprof/heap", HeapHandler()},
		{"GET", "/debug/pprof/goroutine", GoroutineHandler()},
		{"GET", "/debug/pprof/block", BlockHandler()},
		{"GET", "/debug/pprof/threadcreate", ThreadCreateHandler()},
		{"GET", "/debug/pprof/cmdline", CmdlineHandler()},
		{"GET", "/debug/pprof/profile", ProfileHandler()},
		{"GET", "/debug/pprof/symbol", SymbolHandler()},
		{"POST", "/debug/pprof/symbol", SymbolHandler()},
		{"GET", "/debug/pprof/trace", TraceHandler()},
		{"GET", "/debug/pprof/mutex", MutexHandler()},
	}

	basePath := strings.TrimSuffix(router.BasePath(), "/")
	var prefix string

	switch {
	case basePath == "":
		prefix = ""
	case strings.HasSuffix(basePath, "/debug"):
		prefix = "/debug"
	case strings.HasSuffix(basePath, "/debug/pprof"):
		prefix = "/debug/pprof"
	}

	for _, r := range routers {
		router.Handle(r.Method, strings.TrimPrefix(r.Path, prefix), r.Handler)
	}
}

// IndexHandler will pass the call from /debug/pprof to pprof
func IndexHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Index(ctx.Writer, ctx.Request)
	}
}

// HeapHandler will pass the call from /debug/pprof/heap to pprof
func HeapHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("heap").ServeHTTP(ctx.Writer, ctx.Request)
	}
}

// GoroutineHandler will pass the call from /debug/pprof/goroutine to pprof
func GoroutineHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("goroutine").ServeHTTP(ctx.Writer, ctx.Request)
	}
}

// BlockHandler will pass the call from /debug/pprof/block to pprof
func BlockHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("block").ServeHTTP(ctx.Writer, ctx.Request)
	}
}

// ThreadCreateHandler will pass the call from /debug/pprof/threadcreate to pprof
func ThreadCreateHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("threadcreate").ServeHTTP(ctx.Writer, ctx.Request)
	}
}

// CmdlineHandler will pass the call from /debug/pprof/cmdline to pprof
func CmdlineHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Cmdline(ctx.Writer, ctx.Request)
	}
}

// ProfileHandler will pass the call from /debug/pprof/profile to pprof
func ProfileHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Profile(ctx.Writer, ctx.Request)
	}
}

// SymbolHandler will pass the call from /debug/pprof/symbol to pprof
func SymbolHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Symbol(ctx.Writer, ctx.Request)
	}
}

// TraceHandler will pass the call from /debug/pprof/trace to pprof
func TraceHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Trace(ctx.Writer, ctx.Request)
	}
}

// MutexHandler will pass the call from /debug/pprof/mutex to pprof
func MutexHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("mutex").ServeHTTP(ctx.Writer, ctx.Request)
	}
}
<commit_msg>re-add Wrapper for backword compatibility<commit_after>package ginpprof

import (
	"net/http/pprof"
	"strings"

	"github.com/gin-gonic/gin"
)

// Wrap adds several routes from package `net/http/pprof` to *gin.Engine object
func Wrap(router *gin.Engine) {
	WrapGroup(&router.RouterGroup)
}

// Wrapper make sure we are backward compatible
var Wrapper = Wrap

// WrapGroup adds several routes from package `net/http/pprof` to *gin.RouterGroup object
func WrapGroup(router *gin.RouterGroup) {
	routers := []struct {
		Method  string
		Path    string
		Handler gin.HandlerFunc
	}{
		{"GET", "/debug/pprof/", IndexHandler()},
		{"GET", "/debug/pprof/heap", HeapHandler()},
		{"GET", "/debug/pprof/goroutine", GoroutineHandler()},
		{"GET", "/debug/pprof/block", BlockHandler()},
		{"GET", "/debug/pprof/threadcreate", ThreadCreateHandler()},
		{"GET", "/debug/pprof/cmdline", CmdlineHandler()},
		{"GET", "/debug/pprof/profile", ProfileHandler()},
		{"GET", "/debug/pprof/symbol", SymbolHandler()},
		{"POST", "/debug/pprof/symbol", SymbolHandler()},
		{"GET", "/debug/pprof/trace", TraceHandler()},
		{"GET", "/debug/pprof/mutex", MutexHandler()},
	}

	basePath := strings.TrimSuffix(router.BasePath(), "/")
	var prefix string

	switch {
	case basePath == "":
		prefix = ""
	case strings.HasSuffix(basePath, "/debug"):
		prefix = "/debug"
	case strings.HasSuffix(basePath, "/debug/pprof"):
		prefix = "/debug/pprof"
	}

	for _, r := range routers {
		router.Handle(r.Method, strings.TrimPrefix(r.Path, prefix), r.Handler)
	}
}

// IndexHandler will pass the call from /debug/pprof to pprof
func IndexHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Index(ctx.Writer, ctx.Request)
	}
}

// HeapHandler will pass the call from /debug/pprof/heap to pprof
func HeapHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("heap").ServeHTTP(ctx.Writer, ctx.Request)
	}
}

// GoroutineHandler will pass the call from /debug/pprof/goroutine to pprof
func GoroutineHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("goroutine").ServeHTTP(ctx.Writer, ctx.Request)
	}
}

// BlockHandler will pass the call from /debug/pprof/block to pprof
func BlockHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("block").ServeHTTP(ctx.Writer, ctx.Request)
	}
}

// ThreadCreateHandler will pass the call from /debug/pprof/threadcreate to pprof
func ThreadCreateHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("threadcreate").ServeHTTP(ctx.Writer, ctx.Request)
	}
}

// CmdlineHandler will pass the call from /debug/pprof/cmdline to pprof
func CmdlineHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Cmdline(ctx.Writer, ctx.Request)
	}
}

// ProfileHandler will pass the call from /debug/pprof/profile to pprof
func ProfileHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Profile(ctx.Writer, ctx.Request)
	}
}

// SymbolHandler will pass the call from /debug/pprof/symbol to pprof
func SymbolHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Symbol(ctx.Writer, ctx.Request)
	}
}

// TraceHandler will pass the call from /debug/pprof/trace to pprof
func TraceHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Trace(ctx.Writer, ctx.Request)
	}
}

// MutexHandler will pass the call from /debug/pprof/mutex to pprof
func MutexHandler() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		pprof.Handler("mutex").ServeHTTP(ctx.Writer, ctx.Request)
	}
}
<|endoftext|>