<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	linuxproc "github.com/c9s/goprocinfo/linux"
)

var version = "0.0.1"

// set up cli vars
var argIP = flag.String("listen_ip", "", "IP to listen on, defaults to all IPs")
var argPort = flag.Int("port", 8801, "port to listen")
var versionFlag = flag.Bool("version", false, "print cAdvisor-companion version and exit")

// customProcs is our extended linuxproc.Process type
// with custom cpuUsage attribute
type customProcs struct {
	linuxproc.Process
	CPUUsage float64
}

// byCPU helps us sort array of customProcs by CPUUsage
type byCPU []customProcs

// byRSS helps us sort array of customProcs by Status.VmRSS
type byRSS []customProcs

func (p byRSS) Len() int {
	return len(p)
}
func (p byRSS) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}
func (p byRSS) Less(i, j int) bool {
	return p[i].Status.VmRSS < p[j].Status.VmRSS
}

func (p byCPU) Len() int {
	return len(p)
}
func (p byCPU) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}
func (p byCPU) Less(i, j int) bool {
	return p[i].CPUUsage < p[j].CPUUsage
}

// dockerSnapshot containes process list and cpu usage for some docker container
type dockerSnapshot struct {
	processes []linuxproc.Process
	cpuStat   uint64
}

// historyEntry containes all dockerSnapshots for some moment in time
type historyEntry map[string]dockerSnapshot

// history holds RRD-like array of last 60 history entries
var history [60]historyEntry

// getProcesses returns list of processes for given dockerID
func getProcesses(dockerID string) ([]linuxproc.Process, error) {
	tasksPath := fmt.Sprintf("/sys/fs/cgroup/cpu/docker/%s/tasks", dockerID)
	dataBytes, err := ioutil.ReadFile(tasksPath)
	if err != nil {
		return nil, err
	}

	tasksString := string(dataBytes)
	tasks := strings.Split(tasksString, "\n")
	var procs []linuxproc.Process
	for _, t := range tasks {
		pid, err := strconv.ParseUint(t, 10, 64)
		if err == nil {
			p, err := linuxproc.ReadProcess(pid, "/proc/")
			if err != nil {
				continue
			}
			if p.Cmdline != "" && p.Status.VmRSS > 0 {
				procs = append(procs, *p)
			}
		}
	}
	return procs, nil
}

// getCPUTotalUsage returns amount of CPU time, used by given docker container
func getCPUTotalUsage(dockerID string) (uint64, error) {
	cpuAcctPath := fmt.Sprintf("/sys/fs/cgroup/cpuacct/docker/%s/cpuacct.stat", dockerID)
	dataBytes, err := ioutil.ReadFile(cpuAcctPath)
	if err != nil {
		return 0, nil
	}
	cpuAcctString := string(dataBytes)
	cpuAcct := strings.Split(cpuAcctString, "\n")
	var total uint64
	total = 0
	for _, s := range cpuAcct {
		splitted := strings.Split(s, " ")
		if len(splitted) > 1 {
			usage, _ := strconv.ParseUint(splitted[1], 10, 64)
			total += usage
		}
	}

	return total, nil
}

// findProc searches for given pid in list of processes and returns
// its process if found
func findProc(pid uint64, procs []linuxproc.Process) *linuxproc.Process {
	for _, p := range procs {
		if p.Status.Pid == pid {
			return &p
		}
	}
	return nil
}

// getLastData returns latest data from history with added calculated CPU usage
func getLastData(dockerID string) ([]customProcs, error) {
	last := len(history) - 1
	first := last
	for i, e := range history {
		if len(e[dockerID].processes) > 0 {
			first = i
			break
		}
	}
	entry1 := history[first][dockerID]
	entry2 := history[last][dockerID]
	var procs []customProcs
	for _, p2 := range entry2.processes {
		p1 := findProc(p2.Status.Pid, entry1.processes)
		if p1 != nil {
			user := int64(p2.Stat.Utime-p1.Stat.Utime) + (p2.Stat.Cutime - p1.Stat.Cutime)
			system := int64(p2.Stat.Stime-p1.Stat.Stime) + (p2.Stat.Cstime - p1.Stat.Cstime)
			percent := (float64(user+system) / float64(entry2.cpuStat-entry1.cpuStat)) * 100
			procs = append(procs, customProcs{p2, percent})
		}

	}
	return procs, nil
}

// getTopCPU returns `limit` entries with top CPU usage
func getTopCPU(dockerID string, limit int) ([]customProcs, error) {
	procs, err := getLastData(dockerID)
	if err != nil {
		return nil, err
	}
	sort.Sort(sort.Reverse(byCPU(procs)))
	if limit > len(procs) {
		limit = len(procs)
	}
	var result []customProcs
	for _, p := range procs[:limit] {
		result = append(result, p)
		//fmt.Printf("%f%% CPU %s\n", p.CPUUsage, p.Cmdline)
	}
	return result, nil
}

// getTopMem returns `limit` entries with top VmRSS usage
func getTopMem(dockerID string, limit int) ([]customProcs, error) {
	procs, err := getLastData(dockerID)
	if err != nil {
		return nil, err
	}
	sort.Sort(sort.Reverse(byRSS(procs)))
	if limit > len(procs) {
		limit = len(procs)
	}
	var result []customProcs
	for _, p := range procs[:limit] {
		result = append(result, p)
		//fmt.Printf("%dKb %s\n", p.Status.VmRSS, p.Cmdline)
	}
	return result, nil
}

// getDockerIDs collects all docker ids from cgroups pseudo-filesystem
func getDockerIDs() ([]string, error) {
	d, err := os.Open("/sys/fs/cgroup/cpu/docker")
	if err != nil {
		return nil, err
	}
	defer d.Close()

	results := make([]string, 0, 50)
	for {
		fis, err := d.Readdir(10)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		for _, fi := range fis {
			// We only care about directories, since all pids are dirs
			if !fi.IsDir() {
				continue
			}
			name := fi.Name()
			if len(name) > '8' {
				results = append(results, name)
			}
		}
	}

	return results, nil
}

// httpHandler handles http requests
func httpHandler(res http.ResponseWriter, req *http.Request) {
	var validPath = regexp.MustCompile("^/([a-zA-Z0-9]+)/(mem|cpu|all)$")
	m := validPath.FindStringSubmatch(req.URL.Path)
	if m == nil {
		http.NotFound(res, req)
		return
	}
	limitStr := req.URL.Query().Get("limit")
	limit, err := strconv.ParseInt(limitStr, 10, 0)
	if err != nil || limit < 1 {
		limit = 5
	}
	dockerID := m[1]
	var result []customProcs
	switch m[2] {
	case "cpu":
		result, err = getTopCPU(dockerID, int(limit))
	case "mem":
		result, err = getTopMem(dockerID, int(limit))
	case "all":
		result, err = getLastData(dockerID)
	}
	if err != nil {
		fmt.Println(err)
	}
	jsonResult, _ := json.Marshal(result)
	res.Header().Set(
		"Content-Type",
		"text/json",
	)
	io.WriteString(res, string(jsonResult))
}

// collectData scrapes procs data for all docker containers
// every second, and sends through channel to getData()
func collectData() {
	dockerIDs, err := getDockerIDs()
	if err != nil {
		fmt.Println(err)
	}
	for {
		entry := make(historyEntry)
		for _, id := range dockerIDs {
			res, _ := getProcesses(id)
			cpu, _ := getCPUTotalUsage(id)
			entry[id] = dockerSnapshot{res, cpu}
		}
		for i, v := range history[1:] {
			history[i] = v
		}
		history[59] = entry
		time.Sleep(time.Second)
	}
}

func main() {
	flag.Parse()

	if *versionFlag {
		fmt.Printf("cAdvisor-companion version %s\n", version)
		os.Exit(0)
	}

	go collectData()
	//go processData()
	addr := fmt.Sprintf("%s:%d", *argIP, *argPort)
	fmt.Printf("Starting cAdvisor-companion version: %q on port %d\n", version, *argPort)
	http.HandleFunc("/", httpHandler)
	err := http.ListenAndServe(addr, nil)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}
<commit_msg>ADD cgroup parser<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	linuxproc "github.com/c9s/goprocinfo/linux"
)

var version = "0.0.1"

// set up cli vars
var argIP = flag.String("listen_ip", "", "IP to listen on, defaults to all IPs")
var argPort = flag.Int("port", 8801, "port to listen")
var versionFlag = flag.Bool("version", false, "print cAdvisor-companion version and exit")

// customProcs is our extended linuxproc.Process type
// with custom cpuUsage attribute
type customProcs struct {
	linuxproc.Process
	CPUUsage float64
}

// byCPU helps us sort array of customProcs by CPUUsage
type byCPU []customProcs

// byRSS helps us sort array of customProcs by Status.VmRSS
type byRSS []customProcs

func (p byRSS) Len() int {
	return len(p)
}
func (p byRSS) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}
func (p byRSS) Less(i, j int) bool {
	return p[i].Status.VmRSS < p[j].Status.VmRSS
}

func (p byCPU) Len() int {
	return len(p)
}
func (p byCPU) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}
func (p byCPU) Less(i, j int) bool {
	return p[i].CPUUsage < p[j].CPUUsage
}

// dockerSnapshot containes process list and cpu usage for some docker container
type dockerSnapshot struct {
	processes []linuxproc.Process
	cpuStat   uint64
}

// historyEntry containes all dockerSnapshots for some moment in time
type historyEntry map[string]dockerSnapshot

// history holds RRD-like array of last 60 history entries
var history [60]historyEntry

// readCgroup reads and parses /proc/{pid}/cgroup file
func readCgroup(pid uint64, procPath string) (string, error) {
	cgroupPath := fmt.Sprintf("%s/%d/cgroup", procPath, pid)
	dataBytes, err := ioutil.ReadFile(cgroupPath)
	if err != nil {
		return "", err
	}
	cgroupString := string(dataBytes)
	lines := strings.Split(cgroupString, "\n")
	var validLine = regexp.MustCompile("^[0-9]+:([a-z,]+):([a-z0-9/]+)$")
	for _, l := range lines {
		m := validLine.FindStringSubmatch(l)
		if m == nil {
			continue
		}
		// we care only about cpu cgroup
		if strings.Contains(m[1], "cpu") {
			return m[2], nil
		}
	}
	return "", nil
}

// getProcesses returns list of processes for given dockerID
func getProcesses(dockerID string) ([]linuxproc.Process, error) {
	tasksPath := fmt.Sprintf("/sys/fs/cgroup/cpu/docker/%s/tasks", dockerID)
	dataBytes, err := ioutil.ReadFile(tasksPath)
	if err != nil {
		return nil, err
	}

	tasksString := string(dataBytes)
	tasks := strings.Split(tasksString, "\n")
	var procs []linuxproc.Process
	for _, t := range tasks {
		pid, err := strconv.ParseUint(t, 10, 64)
		if err == nil {
			p, err := linuxproc.ReadProcess(pid, "/proc/")
			cgroup, err := readCgroup(pid, "/proc/")
			fmt.Println(cgroup)
			if err != nil {
				continue
			}
			if p.Cmdline != "" && p.Status.VmRSS > 0 {
				procs = append(procs, *p)
			}
		}
	}
	return procs, nil
}

// getCPUTotalUsage returns amount of CPU time, used by given docker container
func getCPUTotalUsage(dockerID string) (uint64, error) {
	cpuAcctPath := fmt.Sprintf("/sys/fs/cgroup/cpuacct/docker/%s/cpuacct.stat", dockerID)
	dataBytes, err := ioutil.ReadFile(cpuAcctPath)
	if err != nil {
		return 0, nil
	}
	cpuAcctString := string(dataBytes)
	cpuAcct := strings.Split(cpuAcctString, "\n")
	var total uint64
	total = 0
	for _, s := range cpuAcct {
		splitted := strings.Split(s, " ")
		if len(splitted) > 1 {
			usage, _ := strconv.ParseUint(splitted[1], 10, 64)
			total += usage
		}
	}

	return total, nil
}

// findProc searches for given pid in list of processes and returns
// its process if found
func findProc(pid uint64, procs []linuxproc.Process) *linuxproc.Process {
	for _, p := range procs {
		if p.Status.Pid == pid {
			return &p
		}
	}
	return nil
}

// getLastData returns latest data from history with added calculated CPU usage
func getLastData(dockerID string) ([]customProcs, error) {
	last := len(history) - 1
	first := last
	for i, e := range history {
		if len(e[dockerID].processes) > 0 {
			first = i
			break
		}
	}
	entry1 := history[first][dockerID]
	entry2 := history[last][dockerID]
	var procs []customProcs
	for _, p2 := range entry2.processes {
		p1 := findProc(p2.Status.Pid, entry1.processes)
		if p1 != nil {
			user := int64(p2.Stat.Utime-p1.Stat.Utime) + (p2.Stat.Cutime - p1.Stat.Cutime)
			system := int64(p2.Stat.Stime-p1.Stat.Stime) + (p2.Stat.Cstime - p1.Stat.Cstime)
			percent := (float64(user+system) / float64(entry2.cpuStat-entry1.cpuStat)) * 100
			procs = append(procs, customProcs{p2, percent})
		}

	}
	return procs, nil
}

// getTopCPU returns `limit` entries with top CPU usage
func getTopCPU(dockerID string, limit int) ([]customProcs, error) {
	procs, err := getLastData(dockerID)
	if err != nil {
		return nil, err
	}
	sort.Sort(sort.Reverse(byCPU(procs)))
	if limit > len(procs) {
		limit = len(procs)
	}
	var result []customProcs
	for _, p := range procs[:limit] {
		result = append(result, p)
		//fmt.Printf("%f%% CPU %s\n", p.CPUUsage, p.Cmdline)
	}
	return result, nil
}

// getTopMem returns `limit` entries with top VmRSS usage
func getTopMem(dockerID string, limit int) ([]customProcs, error) {
	procs, err := getLastData(dockerID)
	if err != nil {
		return nil, err
	}
	sort.Sort(sort.Reverse(byRSS(procs)))
	if limit > len(procs) {
		limit = len(procs)
	}
	var result []customProcs
	for _, p := range procs[:limit] {
		result = append(result, p)
		//fmt.Printf("%dKb %s\n", p.Status.VmRSS, p.Cmdline)
	}
	return result, nil
}

// getDockerIDs collects all docker ids from cgroups pseudo-filesystem
func getDockerIDs() ([]string, error) {
	d, err := os.Open("/sys/fs/cgroup/cpu/docker")
	if err != nil {
		return nil, err
	}
	defer d.Close()

	results := make([]string, 0, 50)
	for {
		fis, err := d.Readdir(10)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}

		for _, fi := range fis {
			// We only care about directories, since all pids are dirs
			if !fi.IsDir() {
				continue
			}
			name := fi.Name()
			if len(name) > '8' {
				results = append(results, name)
			}
		}
	}

	return results, nil
}

// httpHandler handles http requests
func httpHandler(res http.ResponseWriter, req *http.Request) {
	var validPath = regexp.MustCompile("^/([a-zA-Z0-9]+)/(mem|cpu|all)$")
	m := validPath.FindStringSubmatch(req.URL.Path)
	if m == nil {
		http.NotFound(res, req)
		return
	}
	limitStr := req.URL.Query().Get("limit")
	limit, err := strconv.ParseInt(limitStr, 10, 0)
	if err != nil || limit < 1 {
		limit = 5
	}
	dockerID := m[1]
	var result []customProcs
	switch m[2] {
	case "cpu":
		result, err = getTopCPU(dockerID, int(limit))
	case "mem":
		result, err = getTopMem(dockerID, int(limit))
	case "all":
		result, err = getLastData(dockerID)
	}
	if err != nil {
		fmt.Println(err)
	}
	jsonResult, _ := json.Marshal(result)
	res.Header().Set(
		"Content-Type",
		"text/json",
	)
	io.WriteString(res, string(jsonResult))
}

// collectData scrapes procs data for all docker containers
// every second, and sends through channel to getData()
func collectData() {
	dockerIDs, err := getDockerIDs()
	if err != nil {
		fmt.Println(err)
	}
	for {
		entry := make(historyEntry)
		for _, id := range dockerIDs {
			res, _ := getProcesses(id)
			cpu, _ := getCPUTotalUsage(id)
			entry[id] = dockerSnapshot{res, cpu}
		}
		for i, v := range history[1:] {
			history[i] = v
		}
		history[59] = entry
		time.Sleep(time.Second)
	}
}

func main() {
	flag.Parse()

	if *versionFlag {
		fmt.Printf("cAdvisor-companion version %s\n", version)
		os.Exit(0)
	}

	go collectData()
	//go processData()
	addr := fmt.Sprintf("%s:%d", *argIP, *argPort)
	fmt.Printf("Starting cAdvisor-companion version: %q on port %d\n", version, *argPort)
	http.HandleFunc("/", httpHandler)
	err := http.ListenAndServe(addr, nil)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package leveldb

import (
	"path/filepath"

	"github.com/n0stack/n0stack/n0core/pkg/datastore/store"

	"github.com/pkg/errors"
	"github.com/syndtr/goleveldb/leveldb"
	lerrors "github.com/syndtr/goleveldb/leveldb/errors"
	"github.com/syndtr/goleveldb/leveldb/util"
)

type LeveldbStore struct {
	db *leveldb.DB

	prefix string
}

func NewLeveldbStore(directory string) (*LeveldbStore, error) {
	db, err := leveldb.OpenFile("test.db", nil)
	if err != nil {
		return nil, errors.Wrap(err, "failed to connect database")
	}

	return &LeveldbStore{
		db: db,
	}, nil
}

func (ds *LeveldbStore) Close() error {
	return ds.db.Close()
}

func (ds *LeveldbStore) AddPrefix(prefix string) *LeveldbStore {
	return &LeveldbStore{
		db:     ds.db,
		prefix: filepath.Join(ds.prefix, prefix),
	}
}

func (ds *LeveldbStore) List() ([][]byte, error) {
	res := make([][]byte, 0)

	iter := ds.db.NewIterator(util.BytesPrefix([]byte(ds.prefix)), nil)
	for iter.Next() {
		res = append(res, iter.Value())
	}
	iter.Release()

	if err := iter.Error(); err != nil {
		return nil, err
	}

	return res, nil
}

func (ds *LeveldbStore) Get(key string) ([]byte, error) {
	v, err := ds.db.Get(ds.getKey(key), nil)
	if err != nil {
		if err == lerrors.ErrNotFound {
			return nil, store.NewNotFound(key)
		}

		return nil, err
	}

	return v, nil
}

func (ds *LeveldbStore) Apply(key string, value []byte) error {
	return ds.db.Put(ds.getKey(key), value, nil)
}

func (ds *LeveldbStore) Delete(key string) error {
	return ds.db.Delete(ds.getKey(key), nil)
}

func (ds LeveldbStore) getKey(key string) []byte {
	return []byte(filepath.Join(ds.prefix, key))
}
<commit_msg>fix leveldb store which was not store on selected directory<commit_after>package leveldb

import (
	"path/filepath"

	"github.com/n0stack/n0stack/n0core/pkg/datastore/store"

	"github.com/pkg/errors"
	"github.com/syndtr/goleveldb/leveldb"
	lerrors "github.com/syndtr/goleveldb/leveldb/errors"
	"github.com/syndtr/goleveldb/leveldb/util"
)

type LeveldbStore struct {
	db *leveldb.DB

	prefix string
}

func NewLeveldbStore(directory string) (*LeveldbStore, error) {
	db, err := leveldb.OpenFile(directory, nil)
	if err != nil {
		return nil, errors.Wrap(err, "failed to connect database")
	}

	return &LeveldbStore{
		db: db,
	}, nil
}

func (ds *LeveldbStore) Close() error {
	return ds.db.Close()
}

func (ds *LeveldbStore) AddPrefix(prefix string) *LeveldbStore {
	return &LeveldbStore{
		db:     ds.db,
		prefix: filepath.Join(ds.prefix, prefix),
	}
}

func (ds *LeveldbStore) List() ([][]byte, error) {
	res := make([][]byte, 0)

	iter := ds.db.NewIterator(util.BytesPrefix([]byte(ds.prefix)), nil)
	for iter.Next() {
		res = append(res, iter.Value())
	}
	iter.Release()

	if err := iter.Error(); err != nil {
		return nil, err
	}

	return res, nil
}

func (ds *LeveldbStore) Get(key string) ([]byte, error) {
	v, err := ds.db.Get(ds.getKey(key), nil)
	if err != nil {
		if err == lerrors.ErrNotFound {
			return nil, store.NewNotFound(key)
		}

		return nil, err
	}

	return v, nil
}

func (ds *LeveldbStore) Apply(key string, value []byte) error {
	return ds.db.Put(ds.getKey(key), value, nil)
}

func (ds *LeveldbStore) Delete(key string) error {
	return ds.db.Delete(ds.getKey(key), nil)
}

func (ds LeveldbStore) getKey(key string) []byte {
	return []byte(filepath.Join(ds.prefix, key))
}
<|endoftext|>


<commit_before>package cache2go

import (
	"container/list"
	"sync"
	"time"
)

// Cache is an LRU cache. It is not safe for concurrent access.
type Cache struct {
	mu sync.RWMutex
	// MaxEntries is the maximum number of cache entries before
	// an item is evicted. Zero means no limit.
	maxEntries int

	// OnEvicted optionally specificies a callback function to be
	// executed when an entry is purged from the cache.
	OnEvicted func(key string, value interface{})

	lruIndex   *list.List
	ttlIndex   []*list.Element
	cache      map[interface{}]*list.Element
	expiration time.Duration
}

type entry struct {
	key       string
	value     interface{}
	timestamp time.Time
}

// New creates a new Cache.
// If maxEntries is zero, the cache has no limit and it's assumed
// that eviction is done by the caller.
func New(maxEntries int, expire time.Duration) *Cache {
	c := &Cache{
		maxEntries: maxEntries,
		expiration: expire,
		lruIndex:   list.New(),
		ttlIndex:   make([]*list.Element, 0),
		cache:      make(map[interface{}]*list.Element),
	}
	if c.expiration > 0 {
		go c.cleanExpired()
	}
	return c
}

// cleans expired entries performing minimal checks
func (c *Cache) cleanExpired() {
	for {
		c.mu.RLock()
		if len(c.ttlIndex) == 0 {
			c.mu.RUnlock()
			time.Sleep(c.expiration)
			continue
		}
		e := c.ttlIndex[0]
		c.mu.RUnlock()

		en := e.Value.(*entry)
		if time.Now().After(en.timestamp.Add(c.expiration)) {
			c.mu.Lock()
			c.removeElement(e)
			c.mu.Unlock()
		} else {
			time.Sleep(time.Now().Sub(en.timestamp.Add(c.expiration)))
		}
	}
}

// Add adds a value to the cache
func (c *Cache) Set(key string, value interface{}) {
	c.mu.Lock()
	if c.cache == nil {
		c.cache = make(map[interface{}]*list.Element)
		c.lruIndex = list.New()
		c.ttlIndex = make([]*list.Element, 0)
	}

	if e, ok := c.cache[key]; ok {
		c.lruIndex.MoveToFront(e)

		en := e.Value.(*entry)
		en.value = value
		en.timestamp = time.Now()

		c.mu.Unlock()
		return
	}
	e := c.lruIndex.PushFront(&entry{key: key, value: value, timestamp: time.Now()})
	c.ttlIndex = append(c.ttlIndex, e)
	c.cache[key] = e
	c.mu.Unlock()

	if c.maxEntries != 0 && c.lruIndex.Len() > c.maxEntries {
		c.RemoveOldest()
	}
}

// Get looks up a key's value from the cache.
func (c *Cache) Get(key string) (value interface{}, ok bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if c.cache == nil {
		return
	}
	if e, hit := c.cache[key]; hit {
		c.lruIndex.MoveToFront(e)
		e.Value.(*entry).timestamp = time.Now()
		return e.Value.(*entry).value, true
	}
	return
}

// Remove removes the provided key from the cache.
func (c *Cache) Remove(key string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.cache == nil {
		return
	}
	if e, hit := c.cache[key]; hit {
		c.removeElement(e)
	}
}

// RemoveOldest removes the oldest item from the cache.
func (c *Cache) RemoveOldest() {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.cache == nil {
		return
	}
	e := c.lruIndex.Back()
	if e != nil {
		c.removeElement(e)
	}
}

func (c *Cache) removeElement(e *list.Element) {
	c.lruIndex.Remove(e)
	if c.expiration > 0 {
		for i, se := range c.ttlIndex {
			if se == e {
				//delete
				copy(c.ttlIndex[i:], c.ttlIndex[i+1:])
				c.ttlIndex[len(c.ttlIndex)-1] = nil
				c.ttlIndex = c.ttlIndex[:len(c.ttlIndex)-1]
				break
			}
		}
	}
	kv := e.Value.(*entry)
	delete(c.cache, kv.key)
	if c.OnEvicted != nil {
		c.OnEvicted(kv.key, kv.value)
	}
}

// Len returns the number of items in the cache.
func (c *Cache) Len() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if c.cache == nil {
		return 0
	}
	return c.lruIndex.Len()
}

// empties the whole cache
func (c *Cache) Flush() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.lruIndex = list.New()
	c.ttlIndex = make([]*list.Element, 0)
	c.cache = make(map[interface{}]*list.Element)
}
<commit_msg>no update of timestamp on get<commit_after>package cache2go

import (
	"container/list"
	"sync"
	"time"
)

// Cache is an LRU cache. It is not safe for concurrent access.
type Cache struct {
	mu sync.RWMutex
	// MaxEntries is the maximum number of cache entries before
	// an item is evicted. Zero means no limit.
	maxEntries int

	// OnEvicted optionally specificies a callback function to be
	// executed when an entry is purged from the cache.
	OnEvicted func(key string, value interface{})

	lruIndex   *list.List
	ttlIndex   []*list.Element
	cache      map[interface{}]*list.Element
	expiration time.Duration
}

type entry struct {
	key       string
	value     interface{}
	timestamp time.Time
}

// New creates a new Cache.
// If maxEntries is zero, the cache has no limit and it's assumed
// that eviction is done by the caller.
func New(maxEntries int, expire time.Duration) *Cache {
	c := &Cache{
		maxEntries: maxEntries,
		expiration: expire,
		lruIndex:   list.New(),
		ttlIndex:   make([]*list.Element, 0),
		cache:      make(map[interface{}]*list.Element),
	}
	if c.expiration > 0 {
		go c.cleanExpired()
	}
	return c
}

// cleans expired entries performing minimal checks
func (c *Cache) cleanExpired() {
	for {
		c.mu.RLock()
		if len(c.ttlIndex) == 0 {
			c.mu.RUnlock()
			time.Sleep(c.expiration)
			continue
		}
		e := c.ttlIndex[0]
		c.mu.RUnlock()

		en := e.Value.(*entry)
		if time.Now().After(en.timestamp.Add(c.expiration)) {
			c.mu.Lock()
			c.removeElement(e)
			c.mu.Unlock()
		} else {
			time.Sleep(time.Now().Sub(en.timestamp.Add(c.expiration)))
		}
	}
}

// Add adds a value to the cache
func (c *Cache) Set(key string, value interface{}) {
	c.mu.Lock()
	if c.cache == nil {
		c.cache = make(map[interface{}]*list.Element)
		c.lruIndex = list.New()
		c.ttlIndex = make([]*list.Element, 0)
	}

	if e, ok := c.cache[key]; ok {
		c.lruIndex.MoveToFront(e)

		en := e.Value.(*entry)
		en.value = value
		en.timestamp = time.Now()

		c.mu.Unlock()
		return
	}
	e := c.lruIndex.PushFront(&entry{key: key, value: value, timestamp: time.Now()})
	c.ttlIndex = append(c.ttlIndex, e)
	c.cache[key] = e
	c.mu.Unlock()

	if c.maxEntries != 0 && c.lruIndex.Len() > c.maxEntries {
		c.RemoveOldest()
	}
}

// Get looks up a key's value from the cache.
func (c *Cache) Get(key string) (value interface{}, ok bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if c.cache == nil {
		return
	}
	if e, hit := c.cache[key]; hit {
		c.lruIndex.MoveToFront(e)
		return e.Value.(*entry).value, true
	}
	return
}

// Remove removes the provided key from the cache.
func (c *Cache) Remove(key string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.cache == nil {
		return
	}
	if e, hit := c.cache[key]; hit {
		c.removeElement(e)
	}
}

// RemoveOldest removes the oldest item from the cache.
func (c *Cache) RemoveOldest() {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.cache == nil {
		return
	}
	e := c.lruIndex.Back()
	if e != nil {
		c.removeElement(e)
	}
}

func (c *Cache) removeElement(e *list.Element) {
	c.lruIndex.Remove(e)
	if c.expiration > 0 {
		for i, se := range c.ttlIndex {
			if se == e {
				//delete
				copy(c.ttlIndex[i:], c.ttlIndex[i+1:])
				c.ttlIndex[len(c.ttlIndex)-1] = nil
				c.ttlIndex = c.ttlIndex[:len(c.ttlIndex)-1]
				break
			}
		}
	}
	kv := e.Value.(*entry)
	delete(c.cache, kv.key)
	if c.OnEvicted != nil {
		c.OnEvicted(kv.key, kv.value)
	}
}

// Len returns the number of items in the cache.
func (c *Cache) Len() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if c.cache == nil {
		return 0
	}
	return c.lruIndex.Len()
}

// empties the whole cache
func (c *Cache) Flush() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.lruIndex = list.New()
	c.ttlIndex = make([]*list.Element, 0)
	c.cache = make(map[interface{}]*list.Element)
}
<|endoftext|>


<commit_before>package httpretry

import (
	"bytes"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"
)

func TestRetry(t *testing.T) {
	t.Parallel()
	requests := []func(w http.ResponseWriter, r *http.Request){
		func(w http.ResponseWriter, r *http.Request) {
			time.Sleep(time.Second)
			writeTestData(w, 404, "never reached")
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "" {
				t.Errorf("Unexpected Range header on request 2: %s", v)
			}

			head := w.Header()
			head.Set("Accept-Ranges", "bytes")
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "5")
			w.WriteHeader(200)
			w.Write([]byte("ab"))
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "bytes=2-4" {
				t.Errorf("Unexpected Range header on request 3: %s", v)
			}

			head := w.Header()
			head.Set("Content-Range", "bytes 2-4/4")
			head.Set("Accept-Ranges", "bytes")
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "3")
			w.WriteHeader(206)
			w.Write([]byte("cd"))
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "bytes=4-4" {
				t.Errorf("Unexpected Range header on request 4: %s", v)
			}

			time.Sleep(time.Second)
			writeTestData(w, 404, "never reached")
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "bytes=4-4" {
				t.Errorf("Unexpected Range header on request 5: %s", v)
			}

			head := w.Header()
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "4")
			w.WriteHeader(500)
			w.Write([]byte("boom"))
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "bytes=4-4" {
				t.Errorf("Unexpected Range header on request 6: %s", v)
			}

			head := w.Header()
			head.Set("Content-Range", "bytes 4-4/4")
			head.Set("Accept-Ranges", "bytes")
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "1")
			w.WriteHeader(206)
			w.Write([]byte("e"))
		},
	}
	i := 0

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if i < len(requests) {
			requests[i](w, r)
			i += 1
		} else {
			head := w.Header()
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "7")
			w.WriteHeader(404)
			w.Write([]byte("missing"))
		}
	}))
	defer ts.Close()

	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Fatal(err)
	}

	code, head, reader := Getter(req, nil)

	if code != 200 {
		t.Errorf("Unexpected status %d", code)
	}

	if ctype := head.Get("Content-Type"); ctype != "text/plain" {
		t.Errorf("Unexpected Content Type: %s", ctype)
	}

	buf := &bytes.Buffer{}
	written, err := io.Copy(buf, reader)
	if err != nil {
		t.Errorf("Copy error: %s", err)
	}

	if written != 5 {
		t.Errorf("Wrote %d", written)
	}

	if b := buf.String(); b != "abcde" {
		t.Errorf("Got %s", b)
	}

	reader.Close()
}

func TestSingleSuccess(t *testing.T) {
	t.Parallel()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeTestData(w, 200, "ok")
	}))
	defer ts.Close()

	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Fatal(err)
	}

	code, head, reader := Getter(req, nil)

	if code != 200 {
		t.Errorf("Unexpected status %d", code)
	}

	if ctype := head.Get("Content-Type"); ctype != "text/plain" {
		t.Errorf("Unexpected Content Type: %s", ctype)
	}

	buf := &bytes.Buffer{}
	written, err := io.Copy(buf, reader)
	if err != nil {
		t.Errorf("Copy error: %s", err)
	}

	if written != 2 {
		t.Errorf("Wrote %d", written)
	}

	if b := buf.String(); b != "ok" {
		t.Errorf("Got %s", b)
	}

	reader.Close()
}

func TestSkipRetryWithoutAcceptRange(t *testing.T) {
	t.Parallel()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		head := w.Header()
		head.Set("Content-Type", "text/plain")
		head.Set("Content-Length", "2")
		w.WriteHeader(200)
		w.Write([]byte("o"))
	}))
	defer ts.Close()

	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Fatal(err)
	}

	code, head, reader := Getter(req, nil)

	if code != 200 {
		t.Errorf("Unexpected status %d", code)
	}

	if ctype := head.Get("Content-Type"); ctype != "text/plain" {
		t.Errorf("Unexpected Content Type: %s", ctype)
	}

	buf := &bytes.Buffer{}
	written, err := io.Copy(buf, reader)
	if err != nil {
		t.Errorf("Copy error: %s", err)
	}

	if written != 1 {
		t.Errorf("Wrote %d", written)
	}

	if b := buf.String(); b != "o" {
		t.Errorf("Got %s", b)
	}

	reader.Close()
}

func TestSkipRetryWith400(t *testing.T) {
	t.Parallel()
	status := 200
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeTestData(w, status, "client error")
	}))
	defer ts.Close()

	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Fatal(err)
	}

	for status = 400; status < 500; status++ {
		code, head, reader := Getter(req, nil)
		reader.Close()
		if code != status {
			t.Errorf("Expected status %d, got %d", status, code)
		}

		if ctype := head.Get("Content-Type"); ctype != "text/plain" {
			t.Fatalf("Unexpected Content Type: %s", ctype)
		}
	}
}

func writeTestData(w http.ResponseWriter, status int, body string) {
	by := []byte(body)
	head := w.Header()
	head.Set("Accept-Ranges", "bytes")
	head.Set("Content-Type", "text/plain")
	head.Set("Content-Length", strconv.Itoa(len(by)))
	w.WriteHeader(status)
	w.Write(by)
}

func init() {
	tport := http.DefaultTransport.(*http.Transport)
	tport.ResponseHeaderTimeout = 500 * time.Millisecond
}
<commit_msg>assert that we can read the body of a 400 response<commit_after>package httpretry

import (
	"bytes"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"
)

func TestRetry(t *testing.T) {
	t.Parallel()
	requests := []func(w http.ResponseWriter, r *http.Request){
		func(w http.ResponseWriter, r *http.Request) {
			time.Sleep(time.Second)
			writeTestData(w, 404, "never reached")
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "" {
				t.Errorf("Unexpected Range header on request 2: %s", v)
			}

			head := w.Header()
			head.Set("Accept-Ranges", "bytes")
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "5")
			w.WriteHeader(200)
			w.Write([]byte("ab"))
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "bytes=2-4" {
				t.Errorf("Unexpected Range header on request 3: %s", v)
			}

			head := w.Header()
			head.Set("Content-Range", "bytes 2-4/4")
			head.Set("Accept-Ranges", "bytes")
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "3")
			w.WriteHeader(206)
			w.Write([]byte("cd"))
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "bytes=4-4" {
				t.Errorf("Unexpected Range header on request 4: %s", v)
			}

			time.Sleep(time.Second)
			writeTestData(w, 404, "never reached")
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "bytes=4-4" {
				t.Errorf("Unexpected Range header on request 5: %s", v)
			}

			head := w.Header()
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "4")
			w.WriteHeader(500)
			w.Write([]byte("boom"))
		},
		func(w http.ResponseWriter, r *http.Request) {
			if v := r.Header.Get("Range"); v != "bytes=4-4" {
				t.Errorf("Unexpected Range header on request 6: %s", v)
			}

			head := w.Header()
			head.Set("Content-Range", "bytes 4-4/4")
			head.Set("Accept-Ranges", "bytes")
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "1")
			w.WriteHeader(206)
			w.Write([]byte("e"))
		},
	}
	i := 0

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if i < len(requests) {
			requests[i](w, r)
			i += 1
		} else {
			head := w.Header()
			head.Set("Content-Type", "text/plain")
			head.Set("Content-Length", "7")
			w.WriteHeader(404)
			w.Write([]byte("missing"))
		}
	}))
	defer ts.Close()

	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Fatal(err)
	}

	code, head, reader := Getter(req, nil)

	if code != 200 {
		t.Errorf("Unexpected status %d", code)
	}

	if ctype := head.Get("Content-Type"); ctype != "text/plain" {
		t.Errorf("Unexpected Content Type: %s", ctype)
	}

	buf := &bytes.Buffer{}
	written, err := io.Copy(buf, reader)
	if err != nil {
		t.Errorf("Copy error: %s", err)
	}

	if written != 5 {
		t.Errorf("Wrote %d", written)
	}

	if b := buf.String(); b != "abcde" {
		t.Errorf("Got %s", b)
	}

	reader.Close()
}

func TestSingleSuccess(t *testing.T) {
	t.Parallel()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeTestData(w, 200, "ok")
	}))
	defer ts.Close()

	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Fatal(err)
	}

	code, head, reader := Getter(req, nil)

	if code != 200 {
		t.Errorf("Unexpected status %d", code)
	}

	if ctype := head.Get("Content-Type"); ctype != "text/plain" {
		t.Errorf("Unexpected Content Type: %s", ctype)
	}

	buf := &bytes.Buffer{}
	written, err := io.Copy(buf, reader)
	if err != nil {
		t.Errorf("Copy error: %s", err)
	}

	if written != 2 {
		t.Errorf("Wrote %d", written)
	}

	if b := buf.String(); b != "ok" {
		t.Errorf("Got %s", b)
	}

	reader.Close()
}

func TestSkipRetryWithoutAcceptRange(t *testing.T) {
	t.Parallel()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		head := w.Header()
		head.Set("Content-Type", "text/plain")
		head.Set("Content-Length", "2")
		w.WriteHeader(200)
		w.Write([]byte("o"))
	}))
	defer ts.Close()

	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Fatal(err)
	}

	code, head, reader := Getter(req, nil)

	if code != 200 {
		t.Errorf("Unexpected status %d", code)
	}

	if ctype := head.Get("Content-Type"); ctype != "text/plain" {
		t.Errorf("Unexpected Content Type: %s", ctype)
	}

	buf := &bytes.Buffer{}
	written, err := io.Copy(buf, reader)
	if err != nil {
		t.Errorf("Copy error: %s", err)
	}

	if written != 1 {
		t.Errorf("Wrote %d", written)
	}

	if b := buf.String(); b != "o" {
		t.Errorf("Got %s", b)
	}

	reader.Close()
}

func TestSkipRetryWith400(t *testing.T) {
	t.Parallel()
	status := 200
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		writeTestData(w, status, "client error")
	}))
	defer ts.Close()

	req, err := http.NewRequest("GET", ts.URL, nil)
	if err != nil {
		t.Fatal(err)
	}

	for status = 400; status < 500; status++ {
		code, head, reader := Getter(req, nil)

		if code != status {
			t.Errorf("Expected status %d, got %d", status, code)
		}

		if ctype := head.Get("Content-Type"); ctype != "text/plain" {
			t.Fatalf("Unexpected Content Type: %s", ctype)
		}

		buf := &bytes.Buffer{}
		written, err := io.Copy(buf, reader)
		if err != nil {
			t.Errorf("Copy error: %s", err)
		}

		if written != 12 {
			t.Errorf("Wrote %d", written)
		}

		if b := buf.String(); b != "client error" {
			t.Errorf("Got %s", b)
		}

		reader.Close()
	}
}

func writeTestData(w http.ResponseWriter, status int, body string) {
	by := []byte(body)
	head := w.Header()
	head.Set("Accept-Ranges", "bytes")
	head.Set("Content-Type", "text/plain")
	head.Set("Content-Length", strconv.Itoa(len(by)))
	w.WriteHeader(status)
	w.Write(by)
}

func init() {
	tport := http.DefaultTransport.(*http.Transport)
	tport.ResponseHeaderTimeout = 500 * time.Millisecond
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"path/filepath"

	"time"

	. "github.com/claudetech/loggo/default"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/sqlite"
	"golang.org/x/oauth2"
)

// Cache is the cache
type Cache struct {
	db        *gorm.DB
	dbAction  chan cacheAction
	tokenPath string
}

const (
	// StoreAction stores an object in cache
	StoreAction = iota
	// DeleteAction deletes an object in cache
	DeleteAction = iota
)

type cacheAction struct {
	action int
	object *APIObject
}

// APIObject is a Google Drive file object
type APIObject struct {
	ObjectID     string `gorm:"primary_key"`
	Name         string `gorm:"index"`
	IsDir        bool
	Size         uint64
	LastModified time.Time
	DownloadURL  string
	Parents      string `gorm:"index"`
	CreatedAt    time.Time
}

// PageToken is the last change id
type PageToken struct {
	gorm.Model
	Token string
}

// NewCache creates a new cache instance
func NewCache(cacheBasePath string, sqlDebug bool) (*Cache, error) {
	Log.Debugf("Opening cache connection")
	db, err := gorm.Open("sqlite3", filepath.Join(cacheBasePath, "cache"))
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not open cache database")
	}

	Log.Debugf("Migrating cache schema")
	db.AutoMigrate(&APIObject{})
	db.AutoMigrate(&PageToken{})
	db.LogMode(sqlDebug)

	cache := Cache{
		db:        db,
		dbAction:  make(chan cacheAction),
		tokenPath: filepath.Join(cacheBasePath, "token.json"),
	}

	go cache.startStoringQueue()

	return &cache, nil
}

func (c *Cache) startStoringQueue() {
	for {
		action := <-c.dbAction

		if nil != action.object {
			if action.action == DeleteAction || action.action == StoreAction {
				Log.Debugf("Deleting object %v", action.object.ObjectID)
				c.db.Delete(action.object)
			}
			if action.action == StoreAction {
				Log.Debugf("Storing object %v in cache", action.object.ObjectID)
				c.db.Create(action.object)
			}
		}
	}
}

// Close closes all handles
func (c *Cache) Close() error {
	Log.Debugf("Closing cache connection")

	close(c.dbAction)
	if err := c.db.Close(); nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Could not close cache connection")
	}

	return nil
}

// LoadToken loads a token from cache
func (c *Cache) LoadToken() (*oauth2.Token, error) {
	Log.Debugf("Loading token from cache")

	tokenFile, err := ioutil.ReadFile(c.tokenPath)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not read token file in %v", c.tokenPath)
	}

	var token oauth2.Token
	json.Unmarshal(tokenFile, &token)

	Log.Tracef("Got token from cache %v", token)

	return &token, nil
}

// StoreToken stores a token in the cache or updates the existing token element
func (c *Cache) StoreToken(token *oauth2.Token) error {
	Log.Debugf("Storing token to cache")

	tokenJSON, err := json.Marshal(token)
	if nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Could not generate token.json content")
	}

	if err := ioutil.WriteFile(c.tokenPath, tokenJSON, 0644); nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Could not generate token.json file")
	}

	return nil
}

// GetObject gets an object by id
func (c *Cache) GetObject(id string) (*APIObject, error) {
	Log.Tracef("Getting object %v", id)

	var object APIObject
	c.db.Where(&APIObject{ObjectID: id}).First(&object)

	Log.Tracef("Got object from cache %v", object)

	if "" != object.ObjectID {
		return &object, nil
	}

	return nil, fmt.Errorf("Could not find object %v in cache", id)
}

// GetObjectsByParent get all objects under parent id
func (c *Cache) GetObjectsByParent(parent string) ([]*APIObject, error) {
	Log.Tracef("Getting children for %v", parent)

	var objects []*APIObject
	c.db.Where("parents LIKE ?", fmt.Sprintf("%%|%v|%%", parent)).Find(&objects)

	Log.Tracef("Got objects from cache %v", objects)

	if 0 != len(objects) {
		return objects, nil
	}

	return nil, fmt.Errorf("Could not find children for parent %v in cache", parent)
}

// GetObjectByParentAndName finds a child element by name and its parent id
func (c *Cache) GetObjectByParentAndName(parent, name string) (*APIObject, error) {
	Log.Tracef("Getting object %v in parent %v", name, parent)

	var object APIObject
	c.db.Where("parents LIKE ? AND name = ?", fmt.Sprintf("%%|%v|%%", parent), name).First(&object)

	Log.Tracef("Got object from cache %v", object)

	if "" != object.ObjectID {
		return &object, nil
	}

	return nil, fmt.Errorf("Could not find object with name %v in parent %v", name, parent)
}

// DeleteObject deletes an object by id
func (c *Cache) DeleteObject(id string) error {
	c.dbAction <- cacheAction{
		action: DeleteAction,
		object: &APIObject{ObjectID: id},
	}
	return nil
}

// UpdateObject updates an object
func (c *Cache) UpdateObject(object *APIObject) error {
	c.dbAction <- cacheAction{
		action: StoreAction,
		object: object,
	}
	return nil
}

// StoreStartPageToken stores the page token for changes
func (c *Cache) StoreStartPageToken(token string) error {
	Log.Debugf("Storing page token %v in cache", token)

	c.db.Delete(&PageToken{})
	c.db.Create(&PageToken{
		Token: token,
	})

	return nil
}

// GetStartPageToken gets the start page token
func (c *Cache) GetStartPageToken() (string, error) {
	Log.Debugf("Getting start page token from cache")

	var pageToken PageToken
	c.db.First(&pageToken)

	Log.Tracef("Got start page token %v", pageToken.Token)

	if "" == pageToken.Token {
		return "", fmt.Errorf("Token not found in cache")
	}

	return pageToken.Token, nil
}
<commit_msg>really delete db objects<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"path/filepath"

	"time"

	. "github.com/claudetech/loggo/default"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/sqlite"
	"golang.org/x/oauth2"
)

// Cache is the cache
type Cache struct {
	db        *gorm.DB
	dbAction  chan cacheAction
	tokenPath string
}

const (
	// StoreAction stores an object in cache
	StoreAction = iota
	// DeleteAction deletes an object in cache
	DeleteAction = iota
)

type cacheAction struct {
	action int
	object *APIObject
}

// APIObject is a Google Drive file object
type APIObject struct {
	ObjectID     string `gorm:"primary_key"`
	Name         string `gorm:"index"`
	IsDir        bool
	Size         uint64
	LastModified time.Time
	DownloadURL  string
	Parents      string `gorm:"index"`
}

// PageToken is the last change id
type PageToken struct {
	gorm.Model
	Token string
}

// NewCache creates a new cache instance
func NewCache(cacheBasePath string, sqlDebug bool) (*Cache, error) {
	Log.Debugf("Opening cache connection")
	db, err := gorm.Open("sqlite3", filepath.Join(cacheBasePath, "cache"))
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not open cache database")
	}

	Log.Debugf("Migrating cache schema")
	db.AutoMigrate(&APIObject{})
	db.AutoMigrate(&PageToken{})
	db.LogMode(sqlDebug)

	cache := Cache{
		db:        db,
		dbAction:  make(chan cacheAction),
		tokenPath: filepath.Join(cacheBasePath, "token.json"),
	}

	go cache.startStoringQueue()

	return &cache, nil
}

func (c *Cache) startStoringQueue() {
	for {
		action := <-c.dbAction

		if nil != action.object {
			if action.action == DeleteAction || action.action == StoreAction {
				Log.Debugf("Deleting object %v", action.object.ObjectID)
				c.db.Unscoped().Delete(action.object)
			}
			if action.action == StoreAction {
				Log.Debugf("Storing object %v in cache", action.object.ObjectID)
				c.db.Unscoped().Create(action.object)
			}
		}
	}
}

// Close closes all handles
func (c *Cache) Close() error {
	Log.Debugf("Closing cache connection")

	close(c.dbAction)
	if err := c.db.Close(); nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Could not close cache connection")
	}

	return nil
}

// LoadToken loads a token from cache
func (c *Cache) LoadToken() (*oauth2.Token, error) {
	Log.Debugf("Loading token from cache")

	tokenFile, err := ioutil.ReadFile(c.tokenPath)
	if nil != err {
		Log.Debugf("%v", err)
		return nil, fmt.Errorf("Could not read token file in %v", c.tokenPath)
	}

	var token oauth2.Token
	json.Unmarshal(tokenFile, &token)

	Log.Tracef("Got token from cache %v", token)

	return &token, nil
}

// StoreToken stores a token in the cache or updates the existing token element
func (c *Cache) StoreToken(token *oauth2.Token) error {
	Log.Debugf("Storing token to cache")

	tokenJSON, err := json.Marshal(token)
	if nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Could not generate token.json content")
	}

	if err := ioutil.WriteFile(c.tokenPath, tokenJSON, 0644); nil != err {
		Log.Debugf("%v", err)
		return fmt.Errorf("Could not generate token.json file")
	}

	return nil
}

// GetObject gets an object by id
func (c *Cache) GetObject(id string) (*APIObject, error) {
	Log.Tracef("Getting object %v", id)

	var object APIObject
	c.db.Where(&APIObject{ObjectID: id}).First(&object)

	Log.Tracef("Got object from cache %v", object)

	if "" != object.ObjectID {
		return &object, nil
	}

	return nil, fmt.Errorf("Could not find object %v in cache", id)
}

// GetObjectsByParent get all objects under parent id
func (c *Cache) GetObjectsByParent(parent string) ([]*APIObject, error) {
	Log.Tracef("Getting children for %v", parent)

	var objects []*APIObject
	c.db.Where("parents LIKE ?", fmt.Sprintf("%%|%v|%%", parent)).Find(&objects)

	Log.Tracef("Got objects from cache %v", objects)

	if 0 != len(objects) {
		return objects, nil
	}

	return nil, fmt.Errorf("Could not find children for parent %v in cache", parent)
}

// GetObjectByParentAndName finds a child element by name and its parent id
func (c *Cache) GetObjectByParentAndName(parent, name string) (*APIObject, error) {
	Log.Tracef("Getting object %v in parent %v", name, parent)

	var object APIObject
	c.db.Where("parents LIKE ? AND name = ?", fmt.Sprintf("%%|%v|%%", parent), name).First(&object)

	Log.Tracef("Got object from cache %v", object)

	if "" != object.ObjectID {
		return &object, nil
	}

	return nil, fmt.Errorf("Could not find object with name %v in parent %v", name, parent)
}

// DeleteObject deletes an object by id
func (c *Cache) DeleteObject(id string) error {
	c.dbAction <- cacheAction{
		action: DeleteAction,
		object: &APIObject{ObjectID: id},
	}
	return nil
}

// UpdateObject updates an object
func (c *Cache) UpdateObject(object *APIObject) error {
	c.dbAction <- cacheAction{
		action: StoreAction,
		object: object,
	}
	return nil
}

// StoreStartPageToken stores the page token for changes
func (c *Cache) StoreStartPageToken(token string) error {
	Log.Debugf("Storing page token %v in cache", token)

	c.db.Unscoped().Delete(&PageToken{})
	c.db.Unscoped().Create(&PageToken{
		Token: token,
	})

	return nil
}

// GetStartPageToken gets the start page token
func (c *Cache) GetStartPageToken() (string, error) {
	Log.Debugf("Getting start page token from cache")

	var pageToken PageToken
	c.db.First(&pageToken)

	Log.Tracef("Got start page token %v", pageToken.Token)

	if "" == pageToken.Token {
		return "", fmt.Errorf("Token not found in cache")
	}

	return pageToken.Token, nil
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type storageHandler interface {
	listDatabases() ([]string, error)
	listCollections(dbname string) ([]string, error)
	listMetrics(dbname, collection string) ([]string, error)
	insertSample(dbname, collection string, sample map[string]interface{}) error
	findValues(dbname, collection, metric string) (map[string]float64, error)
	aggregate(dbname, collection, metric string) (map[string]interface{}, error)
	getHeatMap(dbname, collection, metric string) (*heatMap, error)
	getHistogram(dbname, collection, metric string) (map[string]float64, error)
}

type mongoHandler struct {
	Session *mgo.Session
}

func newMongoHandler(addrs []string, timeout time.Duration) (*mongoHandler, error) {
	dialInfo := &mgo.DialInfo{
		Addrs:   addrs,
		Timeout: timeout,
	}

	logger.Info("Connecting to database...")
	if session, err := mgo.DialWithInfo(dialInfo); err != nil {
		logger.Criticalf("Failed to connect to database: %s", err)
		return nil, err
	} else {
		logger.Info("Connection established.")
		session.SetMode(mgo.Monotonic, true)
		return &mongoHandler{session}, nil
	}
}

var dbPrefix = "perf"

func (mongo *mongoHandler) listDatabases() ([]string, error) {
	if err := mongo.Session.Ping(); err != nil {
		mongo.Session.Refresh()
	}
	allDbs, err := mongo.Session.DatabaseNames()
	if err != nil {
		return nil, err
	}

	dbs := []string{}
	for _, db := range allDbs {
		if strings.HasPrefix(db, dbPrefix) {
			dbs = append(dbs, strings.Replace(db, dbPrefix, "", 1))
		}
	}
	return dbs, nil
}

func (mongo *mongoHandler) listCollections(dbname string) ([]string, error) {
	session := mongo.Session.New()
	defer session.Close()
	_db := session.DB(dbPrefix + dbname)

	allCollections, err := _db.CollectionNames()
	if err != nil {
		return nil, err
	}

	collections := []string{}
	for _, collection := range allCollections {
		if collection != "system.indexes" {
			collections = append(collections, collection)
		}
	}
	return collections, err
}

func (mongo *mongoHandler) listMetrics(dbname, collection string) ([]string, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var metrics []string
	if err := _collection.Find(bson.M{}).Sort("m").Distinct("m", &metrics); err != nil {
		return nil, err
	}
	return metrics, nil
}

func (mongo *mongoHandler) findValues(dbname, collection, metric string) (map[string]float64, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var docs []map[string]interface{}
	if err := _collection.Find(bson.M{"m": metric}).Sort("ts").All(&docs); err != nil {
		return nil, err
	}
	values := map[string]float64{}
	for _, doc := range docs {
		values[doc["ts"].(string)] = doc["v"].(float64)
	}
	return values, nil
}

func (mongo *mongoHandler) insertSample(dbname, collection string, sample map[string]interface{}) error {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	if err := _collection.Insert(sample); err != nil {
		return err
	}
	logger.Infof("Successfully added new sample to %s.%s", dbname, collection)

	for _, key := range []string{"m", "ts", "v"} {
		if err := _collection.EnsureIndexKey(key); err != nil {
			return err
		}
	}
	return nil
}

func calcPercentile(data []float64, p float64) float64 {
	sort.Float64s(data)

	k := float64(len(data)-1) * p
	f := math.Floor(k)
	c := math.Ceil(k)
	if f == c {
		return data[int(k)]
	} else {
		return data[int(f)]*(c-k) + data[int(c)]*(k-f)
	}
}

const queryLimit = 10000

func (mongo *mongoHandler) aggregate(dbname, collection, metric string) (map[string]interface{}, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	pipe := _collection.Pipe(
		[]bson.M{
			{
				"$match": bson.M{
					"m": metric,
				},
			},
			{
				"$group": bson.M{
					"_id": bson.M{
						"metric": "$m",
					},
					"avg":   bson.M{"$avg": "$v"},
					"min":   bson.M{"$min": "$v"},
					"max":   bson.M{"$max": "$v"},
					"count": bson.M{"$sum": 1},
				},
			},
		},
	)
	summaries := []map[string]interface{}{}
	if err := pipe.All(&summaries); err != nil {
		return nil, err
	}
	if len(summaries) == 0 {
		return nil, nil
	}
	summary := summaries[0]
	delete(summary, "_id")

	count := summary["count"].(int)
	if count < queryLimit {
		// Don't perform in-memory aggregation if limit exceeded
		var docs []map[string]interface{}
		if err := _collection.Find(bson.M{"m": metric}).Select(bson.M{"v": 1}).All(&docs); err != nil {
			return nil, err
		}
		values := []float64{}
		for _, doc := range docs {
			values = append(values, doc["v"].(float64))
		}
		for _, percentile := range []float64{0.5, 0.8, 0.9, 0.95, 0.99} {
			p := fmt.Sprintf("p%v", percentile*100)
			summary[p] = calcPercentile(values, percentile)
		}
	} else {
		// Calculate percentiles using index-based sorting at database level
		var result []map[string]interface{}
		for _, percentile := range []float64{0.5, 0.8, 0.9, 0.95, 0.99} {
			skip := int(float64(count)*percentile) - 1
			if err := _collection.Find(bson.M{"m": metric}).Sort("v").Skip(skip).One(&result); err != nil {
				return nil, err
			}
			p := fmt.Sprintf("p%v", percentile*100)
			summary[p] = result[0]["v"].(float64)
		}
	}

	return summary, nil
}

type heatMap struct {
	MinTS    int64   `json:"minTimestamp"`
	MaxTS    int64   `json:"maxTimestamp"`
	MaxValue float64 `json:"maxValue"`
	Map      [][]int `json:"map"`
}

const (
	height = 60
	width  = 120
)

func newHeatMap() *heatMap {
	hm := heatMap{}
	hm.Map = [][]int{}
	for y := 0; y < height; y++ {
		hm.Map = append(hm.Map, []int{})
		for x := 0; x < width; x++ {
			hm.Map[y] = append(hm.Map[y], 0)
		}
	}
	return &hm
}

func (mongo *mongoHandler) getHeatMap(dbname, collection, metric string) (*heatMap, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var doc map[string]interface{}
	hm := newHeatMap()

	// Min timestamp
	if err := _collection.Find(bson.M{"m": metric}).Sort("ts").One(&doc); err != nil {
		return nil, err
	}
	if tsInt, err := strconv.ParseInt(doc["ts"].(string), 10, 64); err != nil {
		return nil, err
	} else {
		hm.MinTS = tsInt
	}
	// Max timestamp
	if err := _collection.Find(bson.M{"m": metric}).Sort("-ts").One(&doc); err != nil {
		return nil, err
	}
	if tsInt, err := strconv.ParseInt(doc["ts"].(string), 10, 64); err != nil {
		return nil, err
	} else {
		hm.MaxTS = tsInt
	}
	// Max value
	if err := _collection.Find(bson.M{"m": metric}).Sort("-v").One(&doc); err != nil {
		return nil, err
	}
	hm.MaxValue = doc["v"].(float64)

	if hm.MaxTS == hm.MinTS || hm.MaxValue == 0 {
		return hm, nil
	}

	iter := _collection.Find(bson.M{"m": metric}).Sort("ts").Iter()
	for iter.Next(&doc) {
		if tsInt, err := strconv.ParseInt(doc["ts"].(string), 10, 64); err != nil {
			return nil, err
		} else {
			x := math.Floor(width * float64(tsInt-hm.MinTS) / float64(hm.MaxTS-hm.MinTS))
			y := math.Floor(height * doc["v"].(float64) / hm.MaxValue)
			if x == width {
				x--
			}
			if y == height {
				y--
			}
			hm.Map[int(y)][int(x)]++
		}
	}
	if err := iter.Close(); err != nil {
		return nil, err
	}

	return hm, nil
}

const numBins = 6

func (mongo *mongoHandler) getHistogram(dbname, collection, metric string) (map[string]float64, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var skip int
	if count, err := _collection.Find(bson.M{"m": metric}).Count(); err != nil {
		return nil, err
	} else {
		skip = int(float64(count)*0.99) - 1
	}
	if skip == -1 {
		return nil, errors.New("not enough data points")
	}

	var doc map[string]interface{}
	if err := _collection.Find(bson.M{"m": metric}).Sort("v").Skip(skip).One(&doc); err != nil {
		return nil, err
	}
	p99 := doc["v"].(float64)
	if err := _collection.Find(bson.M{"m": metric}).Sort("v").One(&doc); err != nil {
		return nil, err
	}
	minValue := doc["v"].(float64)
	if p99 == minValue {
		return nil, errors.New("dataset lacks variation")
	}

	delta := (p99 - minValue) / numBins
	histogram := map[string]float64{}
	for i := 0; i < numBins; i++ {
		lr := minValue + float64(i)*delta
		rr := lr + delta
		rname := fmt.Sprintf("%f - %f", lr, rr)
		if count, err := _collection.Find(bson.M{"m": metric, "v": bson.M{"$gte": lr, "$lt": rr}}).Count(); err != nil {
			return nil, err
		} else {
			histogram[rname] = 100.0 * float64(count) / float64(skip)
		}
	}

	return histogram, nil
}
<commit_msg>Misc. lint fixes<commit_after>package main

import (
	"errors"
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type storageHandler interface {
	listDatabases() ([]string, error)
	listCollections(dbname string) ([]string, error)
	listMetrics(dbname, collection string) ([]string, error)
	insertSample(dbname, collection string, sample map[string]interface{}) error
	findValues(dbname, collection, metric string) (map[string]float64, error)
	aggregate(dbname, collection, metric string) (map[string]interface{}, error)
	getHeatMap(dbname, collection, metric string) (*heatMap, error)
	getHistogram(dbname, collection, metric string) (map[string]float64, error)
}

type mongoHandler struct {
	Session *mgo.Session
}

func newMongoHandler(addrs []string, timeout time.Duration) (*mongoHandler, error) {
	dialInfo := &mgo.DialInfo{
		Addrs:   addrs,
		Timeout: timeout,
	}

	logger.Info("Connecting to database...")
	if session, err := mgo.DialWithInfo(dialInfo); err != nil {
		logger.Criticalf("Failed to connect to database: %s", err)
		return nil, err
	} else {
		logger.Info("Connection established.")
		session.SetMode(mgo.Monotonic, true)
		return &mongoHandler{session}, nil
	}
}

var dbPrefix = "perf"

func (mongo *mongoHandler) listDatabases() ([]string, error) {
	if err := mongo.Session.Ping(); err != nil {
		mongo.Session.Refresh()
	}
	allDbs, err := mongo.Session.DatabaseNames()
	if err != nil {
		return nil, err
	}

	dbs := []string{}
	for _, db := range allDbs {
		if strings.HasPrefix(db, dbPrefix) {
			dbs = append(dbs, strings.Replace(db, dbPrefix, "", 1))
		}
	}
	return dbs, nil
}

func (mongo *mongoHandler) listCollections(dbname string) ([]string, error) {
	session := mongo.Session.New()
	defer session.Close()
	_db := session.DB(dbPrefix + dbname)

	allCollections, err := _db.CollectionNames()
	if err != nil {
		return nil, err
	}

	collections := []string{}
	for _, collection := range allCollections {
		if collection != "system.indexes" {
			collections = append(collections, collection)
		}
	}
	return collections, err
}

func (mongo *mongoHandler) listMetrics(dbname, collection string) ([]string, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var metrics []string
	if err := _collection.Find(bson.M{}).Sort("m").Distinct("m", &metrics); err != nil {
		return nil, err
	}
	return metrics, nil
}

func (mongo *mongoHandler) findValues(dbname, collection, metric string) (map[string]float64, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var docs []map[string]interface{}
	if err := _collection.Find(bson.M{"m": metric}).Sort("ts").All(&docs); err != nil {
		return nil, err
	}
	values := map[string]float64{}
	for _, doc := range docs {
		values[doc["ts"].(string)] = doc["v"].(float64)
	}
	return values, nil
}

func (mongo *mongoHandler) insertSample(dbname, collection string, sample map[string]interface{}) error {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	if err := _collection.Insert(sample); err != nil {
		return err
	}
	logger.Infof("Successfully added new sample to %s.%s", dbname, collection)

	for _, key := range []string{"m", "ts", "v"} {
		if err := _collection.EnsureIndexKey(key); err != nil {
			return err
		}
	}
	return nil
}

func calcPercentile(data []float64, p float64) float64 {
	sort.Float64s(data)

	k := float64(len(data)-1) * p
	f := math.Floor(k)
	c := math.Ceil(k)
	if f == c {
		return data[int(k)]
	}
	return data[int(f)]*(c-k) + data[int(c)]*(k-f)
}

const queryLimit = 10000

func (mongo *mongoHandler) aggregate(dbname, collection, metric string) (map[string]interface{}, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	pipe := _collection.Pipe(
		[]bson.M{
			{
				"$match": bson.M{
					"m": metric,
				},
			},
			{
				"$group": bson.M{
					"_id": bson.M{
						"metric": "$m",
					},
					"avg":   bson.M{"$avg": "$v"},
					"min":   bson.M{"$min": "$v"},
					"max":   bson.M{"$max": "$v"},
					"count": bson.M{"$sum": 1},
				},
			},
		},
	)
	summaries := []map[string]interface{}{}
	if err := pipe.All(&summaries); err != nil {
		return nil, err
	}
	if len(summaries) == 0 {
		return nil, nil
	}
	summary := summaries[0]
	delete(summary, "_id")

	count := summary["count"].(int)
	if count < queryLimit {
		// Don't perform in-memory aggregation if limit exceeded
		var docs []map[string]interface{}
		if err := _collection.Find(bson.M{"m": metric}).Select(bson.M{"v": 1}).All(&docs); err != nil {
			return nil, err
		}
		values := []float64{}
		for _, doc := range docs {
			values = append(values, doc["v"].(float64))
		}
		for _, percentile := range []float64{0.5, 0.8, 0.9, 0.95, 0.99} {
			p := fmt.Sprintf("p%v", percentile*100)
			summary[p] = calcPercentile(values, percentile)
		}
	} else {
		// Calculate percentiles using index-based sorting at database level
		var result []map[string]interface{}
		for _, percentile := range []float64{0.5, 0.8, 0.9, 0.95, 0.99} {
			skip := int(float64(count)*percentile) - 1
			if err := _collection.Find(bson.M{"m": metric}).Sort("v").Skip(skip).One(&result); err != nil {
				return nil, err
			}
			p := fmt.Sprintf("p%v", percentile*100)
			summary[p] = result[0]["v"].(float64)
		}
	}

	return summary, nil
}

type heatMap struct {
	MinTS    int64   `json:"minTimestamp"`
	MaxTS    int64   `json:"maxTimestamp"`
	MaxValue float64 `json:"maxValue"`
	Map      [][]int `json:"map"`
}

const (
	height = 60
	width  = 120
)

func newHeatMap() *heatMap {
	hm := heatMap{}
	hm.Map = [][]int{}
	for y := 0; y < height; y++ {
		hm.Map = append(hm.Map, []int{})
		for x := 0; x < width; x++ {
			hm.Map[y] = append(hm.Map[y], 0)
		}
	}
	return &hm
}

func (mongo *mongoHandler) getHeatMap(dbname, collection, metric string) (*heatMap, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	var doc map[string]interface{}
	hm := newHeatMap()

	// Min timestamp
	if err := _collection.Find(bson.M{"m": metric}).Sort("ts").One(&doc); err != nil {
		return nil, err
	}
	if tsInt, err := strconv.ParseInt(doc["ts"].(string), 10, 64); err != nil {
		return nil, err
	} else {
		hm.MinTS = tsInt
	}
	// Max timestamp
	if err := _collection.Find(bson.M{"m": metric}).Sort("-ts").One(&doc); err != nil {
		return nil, err
	}
	if tsInt, err := strconv.ParseInt(doc["ts"].(string), 10, 64); err != nil {
		return nil, err
	} else {
		hm.MaxTS = tsInt
	}
	// Max value
	if err := _collection.Find(bson.M{"m": metric}).Sort("-v").One(&doc); err != nil {
		return nil, err
	}
	hm.MaxValue = doc["v"].(float64)

	if hm.MaxTS == hm.MinTS || hm.MaxValue == 0 {
		return hm, nil
	}

	iter := _collection.Find(bson.M{"m": metric}).Sort("ts").Iter()
	for iter.Next(&doc) {
		tsInt, err := strconv.ParseInt(doc["ts"].(string), 10, 64)
		if err != nil {
			return nil, err
		}
		x := math.Floor(width * float64(tsInt-hm.MinTS) / float64(hm.MaxTS-hm.MinTS))
		y := math.Floor(height * doc["v"].(float64) / hm.MaxValue)
		if x == width {
			x--
		}
		if y == height {
			y--
		}
		hm.Map[int(y)][int(x)]++
	}
	if err := iter.Close(); err != nil {
		return nil, err
	}

	return hm, nil
}

const numBins = 6

func (mongo *mongoHandler) getHistogram(dbname, collection, metric string) (map[string]float64, error) {
	session := mongo.Session.New()
	defer session.Close()
	_collection := session.DB(dbPrefix + dbname).C(collection)

	count, err := _collection.Find(bson.M{"m": metric}).Count()
	if err != nil {
		return nil, err
	}
	if count <= 1 {
		return nil, errors.New("not enough data points")
	}
	skip := int(float64(count)*0.99) - 1

	var doc map[string]interface{}
	if err := _collection.Find(bson.M{"m": metric}).Sort("v").Skip(skip).One(&doc); err != nil {
		return nil, err
	}
	p99 := doc["v"].(float64)
	if err := _collection.Find(bson.M{"m": metric}).Sort("v").One(&doc); err != nil {
		return nil, err
	}
	minValue := doc["v"].(float64)
	if p99 == minValue {
		return nil, errors.New("dataset lacks variation")
	}

	delta := (p99 - minValue) / numBins
	histogram := map[string]float64{}
	for i := 0; i < numBins; i++ {
		lr := minValue + float64(i)*delta
		rr := lr + delta
		rname := fmt.Sprintf("%f - %f", lr, rr)
		count, err := _collection.Find(bson.M{"m": metric, "v": bson.M{"$gte": lr, "$lt": rr}}).Count()
		if err != nil {
			return nil, err
		}
		histogram[rname] = 100.0 * float64(count) / float64(skip)
	}

	return histogram, nil
}
<|endoftext|>


<commit_before>package hoverfly

import (
	"bytes"
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/boltdb/bolt"
)

// Cache - cache interface used to store and retrieve request/response payloads
type Cache interface {
	Set(key, value []byte) error
	Get(key []byte) ([]byte, error)
	GetAllRequests() ([]Payload, error)
	RecordsCount() (int, error)
	DeleteData() error
	GetAllKeys() (map[string]bool, error)
	CloseDB()
}

// NewBoltDBCache - returns new BoltCache instance
func NewBoltDBCache(db *bolt.DB, bucket []byte) *BoltCache {
	return &BoltCache{
		DS:             db,
		RequestsBucket: []byte(bucket),
	}
}

// RequestsBucketName - default name for BoltDB bucket
const RequestsBucketName = "rqbucket"

// Cache - provides access to BoltDB and holds current bucket name
type BoltCache struct {
	DS             *bolt.DB
	RequestsBucket []byte
}

// GetDB - returns open BoltDB database with read/write permissions or goes down in flames if
// something bad happends
func GetDB(name string) *bolt.DB {
	log.WithFields(log.Fields{
		"databaseName": name,
	}).Info("Initiating database")
	db, err := bolt.Open(name, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}

	return db
}

// CloseDB - closes database
func (c *BoltCache) CloseDB() {
	c.DS.Close()
}

// Set - saves given key and value pair to cache
func (c *BoltCache) Set(key, value []byte) error {
	err := c.DS.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists(c.RequestsBucket)
		if err != nil {
			return err
		}
		err = bucket.Put(key, value)
		if err != nil {
			return err
		}
		return nil
	})

	return err
}

// Get - searches for given key in the cache and returns value if found
func (c *BoltCache) Get(key []byte) (value []byte, err error) {

	err = c.DS.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(c.RequestsBucket)
		if bucket == nil {
			return fmt.Errorf("Bucket %q not found!", c.RequestsBucket)
		}
		// "Byte slices returned from Bolt are only valid during a transaction."
		var buffer bytes.Buffer
		val := bucket.Get(key)

		// If it doesn't exist then it will return nil
		if val == nil {
			return fmt.Errorf("key %q not found \n", key)
		}

		buffer.Write(val)
		value = buffer.Bytes()
		return nil
	})

	return
}

// GetAllRequests - returns all captured requests/responses
func (c *BoltCache) GetAllRequests() (payloads []Payload, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)
		if b == nil {
			// bucket doesn't exist
			return nil
		}
		c := b.Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			pl, err := decodePayload(v)
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
					"json":  v,
				}).Warning("Failed to deserialize bytes to payload.")
			} else {
				payloads = append(payloads, *pl)
			}
		}
		return nil
	})
	return
}

// RecordsCount - returns records count
func (c *BoltCache) RecordsCount() (count int, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)
		if b == nil {
			// bucket doesn't exist
			return nil
		}

		count = b.Stats().KeyN

		return nil
	})
	return
}

// DeleteData - deletes bucket with all saved data
func (c *BoltCache) DeleteData() error {
	err := c.DeleteBucket(c.RequestsBucket)
	return err
}

// DeleteBucket - deletes bucket with all saved data
func (c *BoltCache) DeleteBucket(name []byte) (err error) {
	err = c.DS.Update(func(tx *bolt.Tx) error {
		err = tx.DeleteBucket(name)
		if err != nil {
			log.WithFields(log.Fields{
				"error": err.Error(),
				"name":  string(name),
			}).Warning("Failed to delete bucket")

		}
		return err
	})
	return
}

// GetAllKeys - gets all current keys
func (c *BoltCache) GetAllKeys() (keys map[string]bool, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)

		keys = make(map[string]bool)

		if b == nil {
			// bucket doesn't exist
			return nil
		}
		c := b.Cursor()

		for k, _ := c.First(); k != nil; k, _ = c.Next() {
			keys[string(k)] = true
		}
		return nil
	})
	return
}
<commit_msg>docstring updated<commit_after>package hoverfly

import (
	"bytes"
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/boltdb/bolt"
)

// Cache - cache interface used to store and retrieve request/response payloads or anything else
type Cache interface {
	Set(key, value []byte) error
	Get(key []byte) ([]byte, error)
	GetAllRequests() ([]Payload, error)
	RecordsCount() (int, error)
	DeleteData() error
	GetAllKeys() (map[string]bool, error)
	CloseDB()
}

// NewBoltDBCache - returns new BoltCache instance
func NewBoltDBCache(db *bolt.DB, bucket []byte) *BoltCache {
	return &BoltCache{
		DS:             db,
		RequestsBucket: []byte(bucket),
	}
}

// RequestsBucketName - default name for BoltDB bucket
const RequestsBucketName = "rqbucket"

// BoltCache - container to implement Cache instance with BoltDB backend for storage
type BoltCache struct {
	DS             *bolt.DB
	RequestsBucket []byte
}

// GetDB - returns open BoltDB database with read/write permissions or goes down in flames if
// something bad happends
func GetDB(name string) *bolt.DB {
	log.WithFields(log.Fields{
		"databaseName": name,
	}).Info("Initiating database")
	db, err := bolt.Open(name, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}

	return db
}

// CloseDB - closes database
func (c *BoltCache) CloseDB() {
	c.DS.Close()
}

// Set - saves given key and value pair to cache
func (c *BoltCache) Set(key, value []byte) error {
	err := c.DS.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists(c.RequestsBucket)
		if err != nil {
			return err
		}
		err = bucket.Put(key, value)
		if err != nil {
			return err
		}
		return nil
	})

	return err
}

// Get - searches for given key in the cache and returns value if found
func (c *BoltCache) Get(key []byte) (value []byte, err error) {

	err = c.DS.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(c.RequestsBucket)
		if bucket == nil {
			return fmt.Errorf("Bucket %q not found!", c.RequestsBucket)
		}
		// "Byte slices returned from Bolt are only valid during a transaction."
		var buffer bytes.Buffer
		val := bucket.Get(key)

		// If it doesn't exist then it will return nil
		if val == nil {
			return fmt.Errorf("key %q not found \n", key)
		}

		buffer.Write(val)
		value = buffer.Bytes()
		return nil
	})

	return
}

// GetAllRequests - returns all captured requests/responses
func (c *BoltCache) GetAllRequests() (payloads []Payload, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)
		if b == nil {
			// bucket doesn't exist
			return nil
		}
		c := b.Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			pl, err := decodePayload(v)
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
					"json":  v,
				}).Warning("Failed to deserialize bytes to payload.")
			} else {
				payloads = append(payloads, *pl)
			}
		}
		return nil
	})
	return
}

// RecordsCount - returns records count
func (c *BoltCache) RecordsCount() (count int, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)
		if b == nil {
			// bucket doesn't exist
			return nil
		}

		count = b.Stats().KeyN

		return nil
	})
	return
}

// DeleteData - deletes bucket with all saved data
func (c *BoltCache) DeleteData() error {
	err := c.DeleteBucket(c.RequestsBucket)
	return err
}

// DeleteBucket - deletes bucket with all saved data
func (c *BoltCache) DeleteBucket(name []byte) (err error) {
	err = c.DS.Update(func(tx *bolt.Tx) error {
		err = tx.DeleteBucket(name)
		if err != nil {
			log.WithFields(log.Fields{
				"error": err.Error(),
				"name":  string(name),
			}).Warning("Failed to delete bucket")

		}
		return err
	})
	return
}

// GetAllKeys - gets all current keys
func (c *BoltCache) GetAllKeys() (keys map[string]bool, err error) {
	err = c.DS.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(c.RequestsBucket)

		keys = make(map[string]bool)

		if b == nil {
			// bucket doesn't exist
			return nil
		}
		c := b.Cursor()

		for k, _ := c.First(); k != nil; k, _ = c.Next() {
			keys[string(k)] = true
		}
		return nil
	})
	return
}
<|endoftext|>


<commit_before>package luar

import (
	"container/list"
	"reflect"
	"sync"

	"github.com/yuin/gopher-lua"
)

const (
	cacheKey = "github.com/layeh/gopher-luar"
	tagName  = "luar"
)

var mu sync.Mutex

type mtCache struct {
	regular, types map[reflect.Type]*lua.LTable
}

func newMTCache() *mtCache {
	return &mtCache{
		regular: make(map[reflect.Type]*lua.LTable),
		types:   make(map[reflect.Type]*lua.LTable),
	}
}

func getMTCache(L *lua.LState) *mtCache {
	registry, ok := L.Get(lua.RegistryIndex).(*lua.LTable)
	if !ok {
		panic("gopher-luar: corrupt lua registry")
	}
	lCache, ok := registry.RawGetString(cacheKey).(*lua.LUserData)
	if !ok {
		lCache = L.NewUserData()
		lCache.Value = newMTCache()
		registry.RawSetString(cacheKey, lCache)
	}
	cache, ok := lCache.Value.(*mtCache)
	if !ok {
		panic("gopher-luar: corrupt luar metatable cache")
	}
	return cache
}

func addMethods(L *lua.LState, vtype reflect.Type, tbl *lua.LTable) {
	for i := 0; i < vtype.NumMethod(); i++ {
		method := vtype.Method(i)
		if method.PkgPath != "" {
			continue
		}
		fn := New(L, method.Func.Interface())
		tbl.RawSetString(method.Name, fn)
		tbl.RawSetString(getUnexportedName(method.Name), fn)
	}
}

func addFields(L *lua.LState, vtype reflect.Type, tbl *lua.LTable) {
	type element struct {
		Type  reflect.Type
		Index []int
	}

	queue := list.New()
	queue.PushFront(element{
		Type: vtype,
	})

	for queue.Len() > 0 {
		e := queue.Back()
		elem := e.Value.(element)
		vtype := elem.Type
		if vtype.Kind() == reflect.Ptr {
			vtype = vtype.Elem()
		}
	fields:
		for i := 0; i < vtype.NumField(); i++ {
			field := vtype.Field(i)
			if field.PkgPath != "" {
				continue
			}
			var names []string
			tag := field.Tag.Get(tagName)
			if tag == "-" {
				continue
			}
			if tag != "" {
				names = []string{
					tag,
				}
			} else {
				names = []string{
					field.Name,
					getUnexportedName(field.Name),
				}
			}
			for _, key := range names {
				if tbl.RawGetString(key) != lua.LNil {
					continue fields
				}
			}
			index := make([]int, len(elem.Index)+1)
			copy(index, elem.Index)
			index[len(elem.Index)] = i

			ud := L.NewUserData()
			ud.Value = index
			for _, key := range names {
				tbl.RawSetString(key, ud)
			}
			if field.Anonymous {
				index := make([]int, len(elem.Index)+1)
				copy(index, elem.Index)
				index[len(elem.Index)] = i
				queue.PushFront(element{
					Type:  field.Type,
					Index: index,
				})
			}
		}

		queue.Remove(e)
	}
}

func getMetatable(L *lua.LState, value reflect.Value) *lua.LTable {
	mu.Lock()
	defer mu.Unlock()

	cache := getMTCache(L)

	vtype := value.Type()
	if vtype.Kind() == reflect.Ptr {
		vtype = vtype.Elem()
	}
	if v := cache.regular[vtype]; v != nil {
		return v
	}

	mt := L.NewTable()
	mt.RawSetString("__tostring", L.NewFunction(allTostring))
	mt.RawSetString("__metatable", L.NewTable())

	ptrMethods := L.NewTable()
	methods := L.NewTable()

	switch vtype.Kind() {
	case reflect.Array:
		mt.RawSetString("__index", L.NewFunction(arrayIndex))
		mt.RawSetString("__newindex", L.NewFunction(arrayNewIndex))
		mt.RawSetString("__len", L.NewFunction(arrayLen))
		mt.RawSetString("__eq", L.NewFunction(arrayEq))
	case reflect.Chan:
		methods.RawSetString("send", L.NewFunction(chanSend))
		methods.RawSetString("receive", L.NewFunction(chanReceive))
		methods.RawSetString("close", L.NewFunction(chanClose))

		mt.RawSetString("__index", L.NewFunction(chanIndex))
		mt.RawSetString("__len", L.NewFunction(chanLen))
		mt.RawSetString("__eq", L.NewFunction(chanEq))
	case reflect.Map:
		mt.RawSetString("__index", L.NewFunction(mapIndex))
		mt.RawSetString("__newindex", L.NewFunction(mapNewIndex))
		mt.RawSetString("__len", L.NewFunction(mapLen))
		mt.RawSetString("__call", L.NewFunction(mapCall))
		mt.RawSetString("__eq", L.NewFunction(mapEq))
	case reflect.Struct:
		fields := L.NewTable()
		addFields(L, vtype, fields)
		mt.RawSetString("fields", fields)

		mt.RawSetString("__index", L.NewFunction(structIndex))
		mt.RawSetString("__newindex", L.NewFunction(structNewIndex))
		mt.RawSetString("__eq", L.NewFunction(ptrEq)) // TODO: allow non-pointer structs to be compared
	case reflect.Slice:
		methods.RawSetString("capacity", L.NewFunction(sliceCapacity))
		methods.RawSetString("append", L.NewFunction(sliceAppend))

		mt.RawSetString("__index", L.NewFunction(sliceIndex))
		mt.RawSetString("__newindex", L.NewFunction(sliceNewIndex))
		mt.RawSetString("__len", L.NewFunction(sliceLen))
		mt.RawSetString("__eq", L.NewFunction(sliceEq))
	default:
		mt.RawSetString("__index", L.NewFunction(ptrIndex))
		mt.RawSetString("__eq", L.NewFunction(ptrEq))
	}

	mt.RawSetString("__pow", L.NewFunction(ptrPow))
	mt.RawSetString("__unm", L.NewFunction(ptrUnm))

	addMethods(L, reflect.PtrTo(vtype), ptrMethods)
	mt.RawSetString("ptr_methods", ptrMethods)

	addMethods(L, vtype, methods)
	mt.RawSetString("methods", methods)

	cache.regular[vtype] = mt
	return mt
}

func getTypeMetatable(L *lua.LState, t reflect.Type) *lua.LTable {
	mu.Lock()
	defer mu.Unlock()

	cache := getMTCache(L)

	if v := cache.types[t]; v != nil {
		return v
	}

	mt := L.NewTable()
	mt.RawSetString("__call", L.NewFunction(typeCall))
	mt.RawSetString("__eq", L.NewFunction(typeEq))

	cache.types[t] = mt
	return mt
}
<commit_msg>fix for Go 1.6 unexported fields<commit_after>package luar

import (
	"container/list"
	"reflect"
	"sync"

	"github.com/yuin/gopher-lua"
)

const (
	cacheKey = "github.com/layeh/gopher-luar"
	tagName  = "luar"
)

var mu sync.Mutex

type mtCache struct {
	regular, types map[reflect.Type]*lua.LTable
}

func newMTCache() *mtCache {
	return &mtCache{
		regular: make(map[reflect.Type]*lua.LTable),
		types:   make(map[reflect.Type]*lua.LTable),
	}
}

func getMTCache(L *lua.LState) *mtCache {
	registry, ok := L.Get(lua.RegistryIndex).(*lua.LTable)
	if !ok {
		panic("gopher-luar: corrupt lua registry")
	}
	lCache, ok := registry.RawGetString(cacheKey).(*lua.LUserData)
	if !ok {
		lCache = L.NewUserData()
		lCache.Value = newMTCache()
		registry.RawSetString(cacheKey, lCache)
	}
	cache, ok := lCache.Value.(*mtCache)
	if !ok {
		panic("gopher-luar: corrupt luar metatable cache")
	}
	return cache
}

func addMethods(L *lua.LState, vtype reflect.Type, tbl *lua.LTable) {
	for i := 0; i < vtype.NumMethod(); i++ {
		method := vtype.Method(i)
		if method.PkgPath != "" {
			continue
		}
		fn := New(L, method.Func.Interface())
		tbl.RawSetString(method.Name, fn)
		tbl.RawSetString(getUnexportedName(method.Name), fn)
	}
}

func addFields(L *lua.LState, vtype reflect.Type, tbl *lua.LTable) {
	type element struct {
		Type  reflect.Type
		Index []int
	}

	queue := list.New()
	queue.PushFront(element{
		Type: vtype,
	})

	for queue.Len() > 0 {
		e := queue.Back()
		elem := e.Value.(element)
		vtype := elem.Type
		if vtype.Kind() == reflect.Ptr {
			vtype = vtype.Elem()
		}
	fields:
		for i := 0; i < vtype.NumField(); i++ {
			field := vtype.Field(i)
			if field.PkgPath != "" && !field.Anonymous {
				continue
			}
			var names []string
			tag := field.Tag.Get(tagName)
			if tag == "-" {
				continue
			}
			if tag != "" {
				names = []string{
					tag,
				}
			} else {
				names = []string{
					field.Name,
					getUnexportedName(field.Name),
				}
			}
			for _, key := range names {
				if tbl.RawGetString(key) != lua.LNil {
					continue fields
				}
			}
			index := make([]int, len(elem.Index)+1)
			copy(index, elem.Index)
			index[len(elem.Index)] = i

			ud := L.NewUserData()
			ud.Value = index
			for _, key := range names {
				tbl.RawSetString(key, ud)
			}
			if field.Anonymous {
				index := make([]int, len(elem.Index)+1)
				copy(index, elem.Index)
				index[len(elem.Index)] = i
				queue.PushFront(element{
					Type:  field.Type,
					Index: index,
				})
			}
		}

		queue.Remove(e)
	}
}

func getMetatable(L *lua.LState, value reflect.Value) *lua.LTable {
	mu.Lock()
	defer mu.Unlock()

	cache := getMTCache(L)

	vtype := value.Type()
	if vtype.Kind() == reflect.Ptr {
		vtype = vtype.Elem()
	}
	if v := cache.regular[vtype]; v != nil {
		return v
	}

	mt := L.NewTable()
	mt.RawSetString("__tostring", L.NewFunction(allTostring))
	mt.RawSetString("__metatable", L.NewTable())

	ptrMethods := L.NewTable()
	methods := L.NewTable()

	switch vtype.Kind() {
	case reflect.Array:
		mt.RawSetString("__index", L.NewFunction(arrayIndex))
		mt.RawSetString("__newindex", L.NewFunction(arrayNewIndex))
		mt.RawSetString("__len", L.NewFunction(arrayLen))
		mt.RawSetString("__eq", L.NewFunction(arrayEq))
	case reflect.Chan:
		methods.RawSetString("send", L.NewFunction(chanSend))
		methods.RawSetString("receive", L.NewFunction(chanReceive))
		methods.RawSetString("close", L.NewFunction(chanClose))

		mt.RawSetString("__index", L.NewFunction(chanIndex))
		mt.RawSetString("__len", L.NewFunction(chanLen))
		mt.RawSetString("__eq", L.NewFunction(chanEq))
	case reflect.Map:
		mt.RawSetString("__index", L.NewFunction(mapIndex))
		mt.RawSetString("__newindex", L.NewFunction(mapNewIndex))
		mt.RawSetString("__len", L.NewFunction(mapLen))
		mt.RawSetString("__call", L.NewFunction(mapCall))
		mt.RawSetString("__eq", L.NewFunction(mapEq))
	case reflect.Struct:
		fields := L.NewTable()
		addFields(L, vtype, fields)
		mt.RawSetString("fields", fields)

		mt.RawSetString("__index", L.NewFunction(structIndex))
		mt.RawSetString("__newindex", L.NewFunction(structNewIndex))
		mt.RawSetString("__eq", L.NewFunction(ptrEq)) // TODO: allow non-pointer structs to be compared
	case reflect.Slice:
		methods.RawSetString("capacity", L.NewFunction(sliceCapacity))
		methods.RawSetString("append", L.NewFunction(sliceAppend))

		mt.RawSetString("__index", L.NewFunction(sliceIndex))
		mt.RawSetString("__newindex", L.NewFunction(sliceNewIndex))
		mt.RawSetString("__len", L.NewFunction(sliceLen))
		mt.RawSetString("__eq", L.NewFunction(sliceEq))
	default:
		mt.RawSetString("__index", L.NewFunction(ptrIndex))
		mt.RawSetString("__eq", L.NewFunction(ptrEq))
	}

	mt.RawSetString("__pow", L.NewFunction(ptrPow))
	mt.RawSetString("__unm", L.NewFunction(ptrUnm))

	addMethods(L, reflect.PtrTo(vtype), ptrMethods)
	mt.RawSetString("ptr_methods", ptrMethods)

	addMethods(L, vtype, methods)
	mt.RawSetString("methods", methods)

	cache.regular[vtype] = mt
	return mt
}

func getTypeMetatable(L *lua.LState, t reflect.Type) *lua.LTable {
	mu.Lock()
	defer mu.Unlock()

	cache := getMTCache(L)

	if v := cache.types[t]; v != nil {
		return v
	}

	mt := L.NewTable()
	mt.RawSetString("__call", L.NewFunction(typeCall))
	mt.RawSetString("__eq", L.NewFunction(typeEq))

	cache.types[t] = mt
	return mt
}
<|endoftext|>


<commit_before>package cache

import (
	"bytes"
	"crypto/sha1"
	"io"
	"log"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/gin-contrib/cache/persistence"
	"github.com/gin-gonic/gin"
)

const (
	CACHE_MIDDLEWARE_KEY = "gincontrib.cache"
)

var (
	PageCachePrefix = "gincontrib.page.cache"
)

type responseCache struct {
	Status int
	Header http.Header
	Data   []byte
}

type cachedWriter struct {
	gin.ResponseWriter
	status  int
	written bool
	store   persistence.CacheStore
	expire  time.Duration
	key     string
}

var _ gin.ResponseWriter = &cachedWriter{}

func urlEscape(prefix string, u string) string {
	key := url.QueryEscape(u)
	if len(key) > 200 {
		h := sha1.New()
		io.WriteString(h, u)
		key = string(h.Sum(nil))
	}
	var buffer bytes.Buffer
	buffer.WriteString(prefix)
	buffer.WriteString(":")
	buffer.WriteString(key)
	return buffer.String()
}

func newCachedWriter(store persistence.CacheStore, expire time.Duration, writer gin.ResponseWriter, key string) *cachedWriter {
	return &cachedWriter{writer, 0, false, store, expire, key}
}

func (w *cachedWriter) WriteHeader(code int) {
	w.status = code
	w.written = true
	w.ResponseWriter.WriteHeader(code)
}

func (w *cachedWriter) Status() int {
	return w.ResponseWriter.Status()
}

func (w *cachedWriter) Written() bool {
	return w.ResponseWriter.Written()
}

func (w *cachedWriter) Write(data []byte) (int, error) {
	ret, err := w.ResponseWriter.Write(data)
	if err == nil {
		store := w.store
		var cache responseCache
		if err := store.Get(w.key, &cache); err == nil {
			data = append(cache.Data, data...)
		}

		//cache response
		val := responseCache{
			w.status,
			w.Header(),
			data,
		}
		err = store.Set(w.key, val, w.expire)
		if err != nil {
			// need logger
		}
	}
	return ret, err
}

func (w *cachedWriter) WriteString(data string) (n int, err error) {
	ret, err := w.ResponseWriter.WriteString(data)
	if err == nil {
		//cache response
		store := w.store
		val := responseCache{
			w.status,
			w.Header(),
			[]byte(data),
		}
		store.Set(w.key, val, w.expire)
	}
	return ret, err
}

// Cache Middleware
func Cache(store *persistence.CacheStore) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Set(CACHE_MIDDLEWARE_KEY, store)
		c.Next()
	}
}

func SiteCache(store persistence.CacheStore, expire time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		var cache responseCache
		url := c.Request.URL
		key := urlEscape(PageCachePrefix, url.RequestURI())
		if err := store.Get(key, &cache); err != nil {
			c.Next()
		} else {
			c.Writer.WriteHeader(cache.Status)
			for k, vals := range cache.Header {
				for _, v := range vals {
					c.Writer.Header().Add(k, v)
				}
			}
			c.Writer.Write(cache.Data)
		}
	}
}

// CachePage Decorator
func CachePage(store persistence.CacheStore, expire time.Duration, handle gin.HandlerFunc) gin.HandlerFunc {

	return func(c *gin.Context) {
		var cache responseCache
		url := c.Request.URL
		key := urlEscape(PageCachePrefix, url.RequestURI())
		if err := store.Get(key, &cache); err != nil {
			if err != persistence.ErrCacheMiss {
				log.Println(err.Error())
			}
			// replace writer
			writer := newCachedWriter(store, expire, c.Writer, key)
			c.Writer = writer
			handle(c)
		} else {
			c.Writer.WriteHeader(cache.Status)
			for k, vals := range cache.Header {
				for _, v := range vals {
					c.Writer.Header().Add(k, v)
				}
			}
			c.Writer.Write(cache.Data)
		}
	}
}

// CachePageAtomic Decorator
func CachePageAtomic(store persistence.CacheStore, expire time.Duration, handle gin.HandlerFunc) gin.HandlerFunc {
	var m sync.Mutex
	p := CachePage(store, expire, handle)
	return func(c *gin.Context) {
		m.Lock()
		defer m.Unlock()
		p(c)
	}
}

func CachePageWithoutHeader(store persistence.CacheStore, expire time.Duration, handle gin.HandlerFunc) gin.HandlerFunc {
	return func(c *gin.Context) {
		var cache responseCache
		url := c.Request.URL
		key := urlEscape(PageCachePrefix, url.RequestURI())
		if err := store.Get(key, &cache); err != nil {
			if err != persistence.ErrCacheMiss {
				log.Println(err.Error())
			}
			// replace writer
			writer := newCachedWriter(store, expire, c.Writer, key)
			c.Writer = writer
			handle(c)
		} else {
			c.Writer.WriteHeader(cache.Status)
			c.Writer.Write(cache.Data)
		}
	}
}
<commit_msg>Export cache key creation method (#31)<commit_after>package cache

import (
	"bytes"
	"crypto/sha1"
	"io"
	"log"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/gin-contrib/cache/persistence"
	"github.com/gin-gonic/gin"
)

const (
	CACHE_MIDDLEWARE_KEY = "gincontrib.cache"
)

var (
	PageCachePrefix = "gincontrib.page.cache"
)

type responseCache struct {
	Status int
	Header http.Header
	Data   []byte
}

type cachedWriter struct {
	gin.ResponseWriter
	status  int
	written bool
	store   persistence.CacheStore
	expire  time.Duration
	key     string
}

var _ gin.ResponseWriter = &cachedWriter{}

// CreateKey creates a package specific key for a given string
func CreateKey(u string) string {
	return urlEscape(PageCachePrefix, u)
}

func urlEscape(prefix string, u string) string {
	key := url.QueryEscape(u)
	if len(key) > 200 {
		h := sha1.New()
		io.WriteString(h, u)
		key = string(h.Sum(nil))
	}
	var buffer bytes.Buffer
	buffer.WriteString(prefix)
	buffer.WriteString(":")
	buffer.WriteString(key)
	return buffer.String()
}

func newCachedWriter(store persistence.CacheStore, expire time.Duration, writer gin.ResponseWriter, key string) *cachedWriter {
	return &cachedWriter{writer, 0, false, store, expire, key}
}

func (w *cachedWriter) WriteHeader(code int) {
	w.status = code
	w.written = true
	w.ResponseWriter.WriteHeader(code)
}

func (w *cachedWriter) Status() int {
	return w.ResponseWriter.Status()
}

func (w *cachedWriter) Written() bool {
	return w.ResponseWriter.Written()
}

func (w *cachedWriter) Write(data []byte) (int, error) {
	ret, err := w.ResponseWriter.Write(data)
	if err == nil {
		store := w.store
		var cache responseCache
		if err := store.Get(w.key, &cache); err == nil {
			data = append(cache.Data, data...)
		}

		//cache response
		val := responseCache{
			w.status,
			w.Header(),
			data,
		}
		err = store.Set(w.key, val, w.expire)
		if err != nil {
			// need logger
		}
	}
	return ret, err
}

func (w *cachedWriter) WriteString(data string) (n int, err error) {
	ret, err := w.ResponseWriter.WriteString(data)
	if err == nil {
		//cache response
		store := w.store
		val := responseCache{
			w.status,
			w.Header(),
			[]byte(data),
		}
		store.Set(w.key, val, w.expire)
	}
	return ret, err
}

// Cache Middleware
func Cache(store *persistence.CacheStore) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Set(CACHE_MIDDLEWARE_KEY, store)
		c.Next()
	}
}

func SiteCache(store persistence.CacheStore, expire time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		var cache responseCache
		url := c.Request.URL
		key := CreateKey(url.RequestURI())
		if err := store.Get(key, &cache); err != nil {
			c.Next()
		} else {
			c.Writer.WriteHeader(cache.Status)
			for k, vals := range cache.Header {
				for _, v := range vals {
					c.Writer.Header().Add(k, v)
				}
			}
			c.Writer.Write(cache.Data)
		}
	}
}

// CachePage Decorator
func CachePage(store persistence.CacheStore, expire time.Duration, handle gin.HandlerFunc) gin.HandlerFunc {

	return func(c *gin.Context) {
		var cache responseCache
		url := c.Request.URL
		key := CreateKey(url.RequestURI())
		if err := store.Get(key, &cache); err != nil {
			if err != persistence.ErrCacheMiss {
				log.Println(err.Error())
			}
			// replace writer
			writer := newCachedWriter(store, expire, c.Writer, key)
			c.Writer = writer
			handle(c)
		} else {
			c.Writer.WriteHeader(cache.Status)
			for k, vals := range cache.Header {
				for _, v := range vals {
					c.Writer.Header().Add(k, v)
				}
			}
			c.Writer.Write(cache.Data)
		}
	}
}

// CachePageAtomic Decorator
func CachePageAtomic(store persistence.CacheStore, expire time.Duration, handle gin.HandlerFunc) gin.HandlerFunc {
	var m sync.Mutex
	p := CachePage(store, expire, handle)
	return func(c *gin.Context) {
		m.Lock()
		defer m.Unlock()
		p(c)
	}
}

func CachePageWithoutHeader(store persistence.CacheStore, expire time.Duration, handle gin.HandlerFunc) gin.HandlerFunc {
	return func(c *gin.Context) {
		var cache responseCache
		url := c.Request.URL
		key := CreateKey(url.RequestURI())
		if err := store.Get(key, &cache); err != nil {
			if err != persistence.ErrCacheMiss {
				log.Println(err.Error())
			}
			// replace writer
			writer := newCachedWriter(store, expire, c.Writer, key)
			c.Writer = writer
			handle(c)
		} else {
			c.Writer.WriteHeader(cache.Status)
			c.Writer.Write(cache.Data)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"fmt"
	"sync"
	"time"

	"github.com/miekg/dns"
)

type Mesg struct {
	Msg    *dns.Msg
	Expire time.Time
}

type Cache interface {
	Get(key string) (Msg *dns.Msg, err error)
	Set(key string, Msg *dns.Msg) error
	Exists(key string) bool
	Remove(key string)
	Length() int
}

type CacheMap struct {
	Backend  map[string]Mesg
	Expire   time.Duration
	Maxcount int
	mu       sync.RWMutex
}

type keyNotFound struct {
	key string
}

func (k keyNotFound) Error() string {
	return "Key " + k.key + " not found."
}

type keyExpired struct {
	key string
}

func (k keyExpired) Error() string {
	return "Key " + k.key + " expired."
}

type CacheIsFull struct {
}

func (c CacheIsFull) Error() string {
	return "Cache is Full"
}

func (m *CacheMap) Get(key string) (*dns.Msg, error) {
	m.mu.RLock()
	mesg, ok := m.Backend[key]
	m.mu.RUnlock()

	if !ok {
		return nil, "Key " + key + " not found."
	}

	if mesg.Expire.Before(time.Now()) {
		m.Remove(key)
		return nil, "Key " + key + " expired."
	}

	return mesg.Msg, nil
}

func (m *CacheMap) Set(key string, msg *dns.Msg) error {
	if m.Full() && !m.Exists(key) {
		return CacheIsFull{}
	}

	expire := time.Now().Add(m.Expire)
	mesg := Mesg{msg, expire}
	m.mu.Lock()
	m.Backend[key] = mesg
	m.mu.Unlock()

	return nil
}

func (m *CacheMap) Remove(key string) {
	m.mu.Lock()
	delete(m.Backend, key)
	m.mu.Unlock()
}

func (m *CacheMap) Exists(key string) bool {
	m.mu.RLock()
	_, ok := m.Backend[key]
	m.mu.RUnlock()
	return ok
}

func (m *CacheMap) Length() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.Backend)
}

func (m *CacheMap) Full() bool {
	if m.Maxcount == 0 {
		return false
	}
	return m.Length() >= m.Maxcount
}

func KeyGen(q Question) string {
	hash := md5.New()
	hash.Write([]byte(q.String()))
	hsum := hash.Sum(nil)
	key := fmt.Sprintf("%x", hsum)
	return key
}
<commit_msg>:gnocco: fixed cache to actually build.<commit_after>package main

import (
	"crypto/md5"
	"fmt"
	"sync"
	"time"

	"github.com/miekg/dns"
)

type Mesg struct {
	Msg    *dns.Msg
	Expire time.Time
}

type Cache interface {
	Get(key string) (Msg *dns.Msg, err error)
	Set(key string, Msg *dns.Msg) error
	Exists(key string) bool
	Remove(key string)
	Length() int
}

type CacheMap struct {
	Backend  map[string]Mesg
	Expire   time.Duration
	Maxcount int
	mu       sync.RWMutex
}

type keyNotFound struct {
	key string
}

func (k keyNotFound) Error() string {
	return "Key " + k.key + " not found."
}

type keyExpired struct {
	key string
}

func (k keyExpired) Error() string {
	return "Key " + k.key + " expired."
}

type CacheIsFull struct {
}

func (c CacheIsFull) Error() string {
	return "Cache is Full"
}

func (m *CacheMap) Get(key string) (*dns.Msg, error) {
	m.mu.RLock()
	mesg, ok := m.Backend[key]
	m.mu.RUnlock()

	if !ok {
		return nil, keyNotFound{key}
	}

	if mesg.Expire.Before(time.Now()) {
		m.Remove(key)
		return nil, keyExpired{key}
	}

	return mesg.Msg, nil
}

func (m *CacheMap) Set(key string, msg *dns.Msg) error {
	if m.Full() && !m.Exists(key) {
		return CacheIsFull{}
	}

	expire := time.Now().Add(m.Expire)
	mesg := Mesg{msg, expire}
	m.mu.Lock()
	m.Backend[key] = mesg
	m.mu.Unlock()

	return nil
}

func (m *CacheMap) Remove(key string) {
	m.mu.Lock()
	delete(m.Backend, key)
	m.mu.Unlock()
}

func (m *CacheMap) Exists(key string) bool {
	m.mu.RLock()
	_, ok := m.Backend[key]
	m.mu.RUnlock()
	return ok
}

func (m *CacheMap) Length() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.Backend)
}

func (m *CacheMap) Full() bool {
	if m.Maxcount == 0 {
		return false
	}
	return m.Length() >= m.Maxcount
}

func KeyGen(q Question) string {
	hash := md5.New()
	hash.Write([]byte(q.String()))
	hsum := hash.Sum(nil)
	key := fmt.Sprintf("%x", hsum)
	return key
}
<|endoftext|>