<commit_before>package main

import (
	"bitbucket.org/ant512/gobble/akismet"
	"errors"
	"fmt"
	"github.com/russross/blackfriday"
	"html"
	"io/ioutil"
	"log"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

type FileRepository struct {
	directory string
	posts     BlogPosts
	tags      map[string]int
	mutex     sync.RWMutex
}

func NewFileRepository(directory string) *FileRepository {

	f := new(FileRepository)
	f.directory = directory

	go f.update()

	return f
}

func (f *FileRepository) AllTags() map[string]int {
	return f.tags
}

func (f *FileRepository) AllPosts() BlogPosts {
	return f.posts
}

func (f *FileRepository) PostWithUrl(url string) (*BlogPost, error) {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	for i := range f.posts {
		if f.posts[i].Url() == url {
			return f.posts[i], nil
		}
	}

	err := errors.New("Could not find post")

	return nil, err
}

func (f *FileRepository) PostWithId(id int) (*BlogPost, error) {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	for i := range f.posts {
		if f.posts[i].Id == id {
			return f.posts[i], nil
		}
	}

	err := errors.New("Could not find post")

	return nil, err
}

func (f *FileRepository) PostsWithTag(tag string, start int, count int) (BlogPosts, int) {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	filteredPosts := BlogPosts{}

	for i := range f.posts {
		if f.posts[i].ContainsTag(tag) {
			filteredPosts = append(filteredPosts, f.posts[i])
		}
	}

	if start > len(filteredPosts) {
		return BlogPosts{}, 0
	}

	if start+count > len(filteredPosts) {
		count = len(filteredPosts) - start
	}

	return filteredPosts[start : start+count], len(filteredPosts)
}

func (f *FileRepository) SearchPosts(term string, start int, count int) (BlogPosts, int) {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	filteredPosts := BlogPosts{}

	if len(term) > 0 {
		for i := range f.posts {
			if f.posts[i].ContainsTerm(term) {
				filteredPosts = append(filteredPosts, f.posts[i])
			}
		}
	} else {
		filteredPosts = f.posts
	}

	if start > len(filteredPosts) {
		return BlogPosts{}, 0
	}

	if start+count > len(filteredPosts) {
		count = len(filteredPosts) - start
	}

	return filteredPosts[start : start+count], len(filteredPosts)
}

func (f *FileRepository) SaveComment(post *BlogPost, akismetAPIKey, serverAddress, remoteAddress, userAgent, referer, author, email, body string) {

	// TODO: Ensure file name is unique

	comment := new(Comment)

	comment.Author = author
	comment.Email = email
	comment.Date = time.Now()
	comment.Body = html.EscapeString(body)

	f.mutex.Lock()
	post.Comments = append(post.Comments, comment)
	f.mutex.Unlock()

	postPath := post.FilePath[:len(post.FilePath)-3]

	dirname := postPath + string(filepath.Separator) + "comments" + string(filepath.Separator)

	filename := timeToFilename(comment.Date)

	log.Println(dirname + filename)

	isSpam, _ := akismet.IsSpamComment(comment.Body, serverAddress, remoteAddress, userAgent, referer, author, email, akismetAPIKey)

	content := "Author: " + comment.Author + "\n"
	content += "Email: " + comment.Email + "\n"
	content += "Date: " + timeToString(comment.Date) + "\n"

	if isSpam {
		content += "Spam: true\n"
	}

	content += "\n"

	content += comment.Body

	err := ioutil.WriteFile(dirname+filename, []byte(content), 0644)

	if err != nil {
		log.Println(err)
	}
}

func (f *FileRepository) update() {
	for {
		start := time.Now()

		f.fetchAllPosts()
		f.fetchAllTags()

		end := time.Now()
		log.Printf("Cached %v posts in %v", len(f.posts), end.Sub(start))

		time.Sleep(10 * time.Minute)
	}
}

func (f *FileRepository) fetchAllPosts() error {

	dirname := f.directory + string(filepath.Separator)

	files, err := ioutil.ReadDir(dirname)

	posts := BlogPosts{}

	for i := range files {

		if files[i].IsDir() {
			continue
		}

		if filepath.Ext(files[i].Name()) != ".md" {
			continue
		}

		post, err := f.fetchPost(dirname + files[i].Name())

		if err != nil {
			return err
		}

		posts = append(posts, post)
	}

	sort.Sort(posts)

	f.mutex.Lock()
	f.posts = posts
	f.mutex.Unlock()

	return err
}

func (f *FileRepository) fetchAllTags() {

	// We're using a map to simulate a set
	tags := make(map[string]int)

	for i := range f.posts {
		for j := range f.posts[i].Tags {

			value := tags[strings.ToLower(f.posts[i].Tags[j])] + 1
			tags[strings.ToLower(f.posts[i].Tags[j])] = value
		}
	}

	f.mutex.Lock()
	f.tags = tags
	f.mutex.Unlock()
}

func (f *FileRepository) fetchPost(filename string) (*BlogPost, error) {

	post := new(BlogPost)
	post.FilePath = filename

	file, err := ioutil.ReadFile(filename)

	if err != nil {
		return post, err
	}

	file = []byte(extractPostHeader(string(file), post))

	htmlFlags := blackfriday.HTML_USE_SMARTYPANTS
	extensions := blackfriday.EXTENSION_AUTOLINK | blackfriday.EXTENSION_HARD_LINE_BREAK | blackfriday.EXTENSION_FENCED_CODE | blackfriday.EXTENSION_NO_INTRA_EMPHASIS

	renderer := blackfriday.HtmlRenderer(htmlFlags, post.Title, "")

	output := blackfriday.Markdown(file, renderer, extensions)

	post.Body = string(output)

	f.fetchCommentsForPost(post, filename)

	return post, nil
}

func (f *FileRepository) fetchCommentsForPost(post *BlogPost, filename string) {

	dirname := filename[:len(filename)-3] + string(filepath.Separator) + "comments" + string(filepath.Separator)

	files, err := ioutil.ReadDir(dirname)

	if err != nil {
		return
	}

	post.Comments = Comments{}

	for i := range files {

		if files[i].IsDir() {
			continue
		}

		if filepath.Ext(files[i].Name()) != ".md" {
			continue
		}

		comment, err := f.fetchComment(dirname + files[i].Name())

		if err != nil {
			log.Fatal(err)
			return
		}

		post.Comments = append(post.Comments, comment)
	}
}

func (f *FileRepository) fetchComment(filename string) (*Comment, error) {
	comment := new(Comment)

	file, err := ioutil.ReadFile(filename)

	if err != nil {
		return comment, err
	}

	file = []byte(extractCommentHeader(string(file), comment))

	htmlFlags := blackfriday.HTML_USE_SMARTYPANTS
	extensions := blackfriday.EXTENSION_AUTOLINK | blackfriday.EXTENSION_HARD_LINE_BREAK | blackfriday.EXTENSION_FENCED_CODE | blackfriday.EXTENSION_NO_INTRA_EMPHASIS

	renderer := blackfriday.HtmlRenderer(htmlFlags, "", "")

	output := blackfriday.Markdown(file, renderer, extensions)

	comment.Body = string(output)

	return comment, nil
}

func extractCommentHeader(text string, comment *Comment) string {

	lines := strings.Split(text, "\n")

	headerSize := 0

	for i := range lines {
		if strings.Contains(lines[i], ":") {
			components := strings.Split(lines[i], ":")

			header := strings.ToLower(strings.Trim(components[0], " "))
			separatorIndex := strings.Index(lines[i], ":") + 1
			data := strings.Trim(lines[i][separatorIndex:], " ")

			switch header {
			case "author":
				comment.Author = data
			case "email":
				comment.Email = data
			case "date":
				comment.Date = stringToTime(data)
			case "spam":
				comment.IsSpam = data == "true"
			default:
				continue
			}

			headerSize += len(lines[i]) + 1
		} else {
			break
		}
	}

	return text[headerSize:]
}

func extractPostHeader(text string, post *BlogPost) string {

	lines := strings.Split(text, "\n")

	headerSize := 0

	for i := range lines {
		if strings.Contains(lines[i], ":") {
			components := strings.Split(lines[i], ":")

			header := strings.ToLower(strings.Trim(components[0], " "))
			separatorIndex := strings.Index(lines[i], ":") + 1
			data := strings.Trim(lines[i][separatorIndex:], " ")

			switch header {
			case "title":
				post.Title = data
			case "id":
				post.Id, _ = strconv.Atoi(data)
			case "tags":

				tags := strings.Split(data, ",")

				formattedTags := []string{}

				for j := range tags {
					tags[j] = strings.Trim(tags[j], " ")
					tags[j] = strings.Replace(tags[j], " ", "-", -1)
					tags[j] = strings.Replace(tags[j], "/", "-", -1)
					tags[j] = strings.ToLower(tags[j])

					if tags[j] != "" {
						formattedTags = append(formattedTags, tags[j])
					}
				}

				post.Tags = formattedTags
			case "date":
				post.PublishDate = stringToTime(data)
			default:
				continue
			}

			headerSize += len(lines[i]) + 1
		} else {
			break
		}
	}

	return text[headerSize:]
}

func timeToFilename(t time.Time) string {
	return fmt.Sprintf("%04d-%02d-%02d_%02d-%02d-%02d.md", t.Year(), int(t.Month()), t.Day(), t.Hour(), t.Minute(), t.Second())
}

func timeToString(t time.Time) string {
	return fmt.Sprintf("%04d-%02d-%02d %02d:%02d:%02d", t.Year(), int(t.Month()), t.Day(), t.Hour(), t.Minute(), t.Second())
}

func stringToTime(s string) time.Time {

	year, err := strconv.Atoi(s[:4])

	if err != nil {
		return time.Unix(0, 0)
	}

	month, err := strconv.Atoi(s[5:7])

	if err != nil {
		return time.Unix(0, 0)
	}

	day, err := strconv.Atoi(s[8:10])

	if err != nil {
		return time.Unix(0, 0)
	}

	hour, err := strconv.Atoi(s[11:13])

	if err != nil {
		return time.Unix(0, 0)
	}

	minute, err := strconv.Atoi(s[14:16])

	if err != nil {
		return time.Unix(0, 0)
	}

	seconds, err := strconv.Atoi(s[17:19])

	if err != nil {
		return time.Unix(0, 0)
	}

	location, err := time.LoadLocation("UTC")

	return time.Date(year, time.Month(month), day, hour, minute, seconds, 0, location)
}
<commit_msg>Fixed possible concurrency issues.<commit_after>package main

import (
	"bitbucket.org/ant512/gobble/akismet"
	"errors"
	"fmt"
	"github.com/russross/blackfriday"
	"html"
	"io/ioutil"
	"log"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

type FileRepository struct {
	directory string
	posts     BlogPosts
	tags      map[string]int
	mutex     sync.RWMutex
}

func NewFileRepository(directory string) *FileRepository {

	f := new(FileRepository)
	f.directory = directory

	go f.update()

	return f
}

func (f *FileRepository) AllTags() map[string]int {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	tags := make(map[string]int)

	for i := range f.tags {
		tags[i] = f.tags[i]
	}

	return tags
}

func (f *FileRepository) AllPosts() BlogPosts {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	posts := BlogPosts{}

	for i := range f.posts {
		posts = append(posts, f.posts[i])
	}

	return posts
}

func (f *FileRepository) PostWithUrl(url string) (*BlogPost, error) {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	for i := range f.posts {
		if f.posts[i].Url() == url {
			return f.posts[i], nil
		}
	}

	err := errors.New("Could not find post")

	return nil, err
}

func (f *FileRepository) PostWithId(id int) (*BlogPost, error) {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	for i := range f.posts {
		if f.posts[i].Id == id {
			return f.posts[i], nil
		}
	}

	err := errors.New("Could not find post")

	return nil, err
}

func (f *FileRepository) PostsWithTag(tag string, start int, count int) (BlogPosts, int) {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	filteredPosts := BlogPosts{}

	for i := range f.posts {
		if f.posts[i].ContainsTag(tag) {
			filteredPosts = append(filteredPosts, f.posts[i])
		}
	}

	if start > len(filteredPosts) {
		return BlogPosts{}, 0
	}

	if start+count > len(filteredPosts) {
		count = len(filteredPosts) - start
	}

	return filteredPosts[start : start+count], len(filteredPosts)
}

func (f *FileRepository) SearchPosts(term string, start int, count int) (BlogPosts, int) {

	f.mutex.RLock()
	defer f.mutex.RUnlock()

	filteredPosts := BlogPosts{}

	if len(term) > 0 {
		for i := range f.posts {
			if f.posts[i].ContainsTerm(term) {
				filteredPosts = append(filteredPosts, f.posts[i])
			}
		}
	} else {
		filteredPosts = f.posts
	}

	if start > len(filteredPosts) {
		return BlogPosts{}, 0
	}

	if start+count > len(filteredPosts) {
		count = len(filteredPosts) - start
	}

	return filteredPosts[start : start+count], len(filteredPosts)
}

func (f *FileRepository) SaveComment(post *BlogPost, akismetAPIKey, serverAddress, remoteAddress, userAgent, referer, author, email, body string) {

	// TODO: Ensure file name is unique

	comment := new(Comment)

	comment.Author = author
	comment.Email = email
	comment.Date = time.Now()
	comment.Body = html.EscapeString(body)

	f.mutex.Lock()
	post.Comments = append(post.Comments, comment)
	f.mutex.Unlock()

	postPath := post.FilePath[:len(post.FilePath)-3]

	dirname := postPath + string(filepath.Separator) + "comments" + string(filepath.Separator)

	filename := timeToFilename(comment.Date)

	log.Println(dirname + filename)

	isSpam, _ := akismet.IsSpamComment(comment.Body, serverAddress, remoteAddress, userAgent, referer, author, email, akismetAPIKey)

	content := "Author: " + comment.Author + "\n"
	content += "Email: " + comment.Email + "\n"
	content += "Date: " + timeToString(comment.Date) + "\n"

	if isSpam {
		content += "Spam: true\n"
	}

	content += "\n"

	content += comment.Body

	err := ioutil.WriteFile(dirname+filename, []byte(content), 0644)

	if err != nil {
		log.Println(err)
	}
}

func (f *FileRepository) update() {
	for {
		start := time.Now()

		f.fetchAllPosts()
		f.fetchAllTags()

		end := time.Now()
		log.Printf("Cached %v posts in %v", len(f.posts), end.Sub(start))

		time.Sleep(10 * time.Minute)
	}
}

func (f *FileRepository) fetchAllPosts() error {

	dirname := f.directory + string(filepath.Separator)

	files, err := ioutil.ReadDir(dirname)

	posts := BlogPosts{}

	for i := range files {

		if files[i].IsDir() {
			continue
		}

		if filepath.Ext(files[i].Name()) != ".md" {
			continue
		}

		post, err := f.fetchPost(dirname + files[i].Name())

		if err != nil {
			return err
		}

		posts = append(posts, post)
	}

	sort.Sort(posts)

	f.mutex.Lock()
	f.posts = posts
	f.mutex.Unlock()

	return err
}

func (f *FileRepository) fetchAllTags() {

	// We're using a map to simulate a set
	tags := make(map[string]int)

	f.mutex.RLock()

	for i := range f.posts {
		for j := range f.posts[i].Tags {

			value := tags[strings.ToLower(f.posts[i].Tags[j])] + 1
			tags[strings.ToLower(f.posts[i].Tags[j])] = value
		}
	}

	f.mutex.RUnlock()

	f.mutex.Lock()
	f.tags = tags
	f.mutex.Unlock()
}

func (f *FileRepository) fetchPost(filename string) (*BlogPost, error) {

	post := new(BlogPost)
	post.FilePath = filename

	file, err := ioutil.ReadFile(filename)

	if err != nil {
		return post, err
	}

	file = []byte(extractPostHeader(string(file), post))

	htmlFlags := blackfriday.HTML_USE_SMARTYPANTS
	extensions := blackfriday.EXTENSION_AUTOLINK | blackfriday.EXTENSION_HARD_LINE_BREAK | blackfriday.EXTENSION_FENCED_CODE | blackfriday.EXTENSION_NO_INTRA_EMPHASIS

	renderer := blackfriday.HtmlRenderer(htmlFlags, post.Title, "")

	output := blackfriday.Markdown(file, renderer, extensions)

	post.Body = string(output)

	f.fetchCommentsForPost(post, filename)

	return post, nil
}

func (f *FileRepository) fetchCommentsForPost(post *BlogPost, filename string) {

	dirname := filename[:len(filename)-3] + string(filepath.Separator) + "comments" + string(filepath.Separator)

	files, err := ioutil.ReadDir(dirname)

	if err != nil {
		return
	}

	post.Comments = Comments{}

	for i := range files {

		if files[i].IsDir() {
			continue
		}

		if filepath.Ext(files[i].Name()) != ".md" {
			continue
		}

		comment, err := f.fetchComment(dirname + files[i].Name())

		if err != nil {
			log.Fatal(err)
			return
		}

		post.Comments = append(post.Comments, comment)
	}
}

func (f *FileRepository) fetchComment(filename string) (*Comment, error) {
	comment := new(Comment)

	file, err := ioutil.ReadFile(filename)

	if err != nil {
		return comment, err
	}

	file = []byte(extractCommentHeader(string(file), comment))

	htmlFlags := blackfriday.HTML_USE_SMARTYPANTS
	extensions := blackfriday.EXTENSION_AUTOLINK | blackfriday.EXTENSION_HARD_LINE_BREAK | blackfriday.EXTENSION_FENCED_CODE | blackfriday.EXTENSION_NO_INTRA_EMPHASIS

	renderer := blackfriday.HtmlRenderer(htmlFlags, "", "")

	output := blackfriday.Markdown(file, renderer, extensions)

	comment.Body = string(output)

	return comment, nil
}

func extractCommentHeader(text string, comment *Comment) string {

	lines := strings.Split(text, "\n")

	headerSize := 0

	for i := range lines {
		if strings.Contains(lines[i], ":") {
			components := strings.Split(lines[i], ":")

			header := strings.ToLower(strings.Trim(components[0], " "))
			separatorIndex := strings.Index(lines[i], ":") + 1
			data := strings.Trim(lines[i][separatorIndex:], " ")

			switch header {
			case "author":
				comment.Author = data
			case "email":
				comment.Email = data
			case "date":
				comment.Date = stringToTime(data)
			case "spam":
				comment.IsSpam = data == "true"
			default:
				continue
			}

			headerSize += len(lines[i]) + 1
		} else {
			break
		}
	}

	return text[headerSize:]
}

func extractPostHeader(text string, post *BlogPost) string {

	lines := strings.Split(text, "\n")

	headerSize := 0

	for i := range lines {
		if strings.Contains(lines[i], ":") {
			components := strings.Split(lines[i], ":")

			header := strings.ToLower(strings.Trim(components[0], " "))
			separatorIndex := strings.Index(lines[i], ":") + 1
			data := strings.Trim(lines[i][separatorIndex:], " ")

			switch header {
			case "title":
				post.Title = data
			case "id":
				post.Id, _ = strconv.Atoi(data)
			case "tags":

				tags := strings.Split(data, ",")

				formattedTags := []string{}

				for j := range tags {
					tags[j] = strings.Trim(tags[j], " ")
					tags[j] = strings.Replace(tags[j], " ", "-", -1)
					tags[j] = strings.Replace(tags[j], "/", "-", -1)
					tags[j] = strings.ToLower(tags[j])

					if tags[j] != "" {
						formattedTags = append(formattedTags, tags[j])
					}
				}

				post.Tags = formattedTags
			case "date":
				post.PublishDate = stringToTime(data)
			default:
				continue
			}

			headerSize += len(lines[i]) + 1
		} else {
			break
		}
	}

	return text[headerSize:]
}

func timeToFilename(t time.Time) string {
	return fmt.Sprintf("%04d-%02d-%02d_%02d-%02d-%02d.md", t.Year(), int(t.Month()), t.Day(), t.Hour(), t.Minute(), t.Second())
}

func timeToString(t time.Time) string {
	return fmt.Sprintf("%04d-%02d-%02d %02d:%02d:%02d", t.Year(), int(t.Month()), t.Day(), t.Hour(), t.Minute(), t.Second())
}

func stringToTime(s string) time.Time {

	year, err := strconv.Atoi(s[:4])

	if err != nil {
		return time.Unix(0, 0)
	}

	month, err := strconv.Atoi(s[5:7])

	if err != nil {
		return time.Unix(0, 0)
	}

	day, err := strconv.Atoi(s[8:10])

	if err != nil {
		return time.Unix(0, 0)
	}

	hour, err := strconv.Atoi(s[11:13])

	if err != nil {
		return time.Unix(0, 0)
	}

	minute, err := strconv.Atoi(s[14:16])

	if err != nil {
		return time.Unix(0, 0)
	}

	seconds, err := strconv.Atoi(s[17:19])

	if err != nil {
		return time.Unix(0, 0)
	}

	location, err := time.LoadLocation("UTC")

	return time.Date(year, time.Month(month), day, hour, minute, seconds, 0, location)
}
<|endoftext|>


<commit_before>package filestore

import (
	"fmt"
	"io"
	"os"

	ds "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/ipfs/go-datastore"
	"github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/ipfs/go-datastore/query"
	b58 "gx/ipfs/QmT8rehPR3F6bmwL6zjUN8XpiDBFFpMP2myPdC6ApsWfJf/go-base58"
)

const (
	StatusOk      = 1
	StatusMissing = 2
	StatusInvalid = 3
	StatusError   = 4
)

func statusStr(status int) string {
	switch status {
	case 0:
		return ""
	case 1:
		return "ok      "
	case 2:
		return "missing "
	case 3:
		return "invalid "
	case 4:
		return "error   "
	default:
		return "??      "
	}
}

type ListRes struct {
	Key []byte
	DataObj
	Status int
}

func (r *ListRes) Format() string {
	mhash := b58.Encode(r.Key)
	return fmt.Sprintf("%s%s %s\n", statusStr(r.Status), mhash, r.DataObj.Format())
}

func list(d *Datastore, out chan<- *ListRes, verify bool) error {
	qr, err := d.Query(query.Query{KeysOnly: true})
	if err != nil {
		return err
	}
	for r := range qr.Next() {
		if r.Error != nil {
			return r.Error
		}
		key := ds.NewKey(r.Key)
		val, _ := d.GetDirect(key)
		status := 0
		if verify {
			_, err := d.GetData(key, val, true)
			if err == nil {
				status = StatusOk
			} else if os.IsNotExist(err) {
				status = StatusMissing
			} else if _, ok := err.(InvalidBlock); ok || err == io.EOF || err == io.ErrUnexpectedEOF {
				status = StatusInvalid
			} else {
				status = StatusError
			}
		}
		out <- &ListRes{key.Bytes()[1:], val.StripData(), status}
	}
	return nil
}

func List(d *Datastore, out chan<- *ListRes) error { return list(d, out, false) }

func Verify(d *Datastore, out chan<- *ListRes) error { return list(d, out, true) }
<commit_msg>"filestore verify": change "invalid" status to "changed".<commit_after>package filestore

import (
	"fmt"
	"io"
	"os"

	ds "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/ipfs/go-datastore"
	"github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/ipfs/go-datastore/query"
	b58 "gx/ipfs/QmT8rehPR3F6bmwL6zjUN8XpiDBFFpMP2myPdC6ApsWfJf/go-base58"
)

const (
	StatusOk      = 1
	StatusChanged = 2
	StatusMissing = 3
	StatusError   = 4
)

func statusStr(status int) string {
	switch status {
	case 0:
		return ""
	case StatusOk:
		return "ok       "
	case StatusChanged:
		return "changed  "
	case StatusMissing:
		return "missing  "
	case StatusError:
		return "error    "
	default:
		return "??       "
	}
}

type ListRes struct {
	Key []byte
	DataObj
	Status int
}

func (r *ListRes) Format() string {
	mhash := b58.Encode(r.Key)
	return fmt.Sprintf("%s%s %s\n", statusStr(r.Status), mhash, r.DataObj.Format())
}

func list(d *Datastore, out chan<- *ListRes, verify bool) error {
	qr, err := d.Query(query.Query{KeysOnly: true})
	if err != nil {
		return err
	}
	for r := range qr.Next() {
		if r.Error != nil {
			return r.Error
		}
		key := ds.NewKey(r.Key)
		val, _ := d.GetDirect(key)
		status := 0
		if verify {
			_, err := d.GetData(key, val, true)
			if err == nil {
				status = StatusOk
			} else if os.IsNotExist(err) {
				status = StatusMissing
			} else if _, ok := err.(InvalidBlock); ok || err == io.EOF || err == io.ErrUnexpectedEOF {
				status = StatusChanged
			} else {
				status = StatusError
			}
		}
		out <- &ListRes{key.Bytes()[1:], val.StripData(), status}
	}
	return nil
}

func List(d *Datastore, out chan<- *ListRes) error { return list(d, out, false) }

func Verify(d *Datastore, out chan<- *ListRes) error { return list(d, out, true) }
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"os/exec"
	"strings"
)

func OpenDB() *sql.DB {
	db, err := sql.Open("sqlite3", "toril.db")
	if err != nil {
		log.Fatalln("Fatal Error: Cannot open DB: ", err)
	}
	_, err = db.Exec("PRAGMA foreign_keys = ON;")
	if err != nil {
		log.Fatalln("Fatal Error: Cannot enforce foreign keys: ", err)
	}
	return db
}

func ChkRows(rows *sql.Rows) {
	err := rows.Err()
	if err != nil {
		log.Fatalln("Fatal Error: Rows returned error: ", err)
	}
	rows.Close()
}

func BackupDB() {
	cmd := exec.Command("sh", "-c",
		"echo '.dump' | sqlite3 toril.db | "+
			"gzip -c >bak/toril.db.`date +\"%Y-%m-%d\"`.gz")
	err := cmd.Run()
	if err != nil {
		log.Fatalln("Fatal Error: Cannot backup DB: ", err)
	}
	// restore: cat dumpfile.sql | sqlite3 my_database.sqlite
}

func RestoreDB(file string) { // this doesn't work on Mac OS X
	cmd := exec.Command("sh", "-c", "zcat "+file+" | sqlite3 toril.db")
	err := cmd.Run()
	if err != nil {
		log.Fatalln("Fatal Error: Cannot restore DB: ", err)
	}
}
<commit_msg>removed unnused strings import<commit_after>package main

import (
	"database/sql"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"os/exec"
)

func OpenDB() *sql.DB {
	db, err := sql.Open("sqlite3", "toril.db")
	if err != nil {
		log.Fatalln("Fatal Error: Cannot open DB: ", err)
	}
	_, err = db.Exec("PRAGMA foreign_keys = ON;")
	if err != nil {
		log.Fatalln("Fatal Error: Cannot enforce foreign keys: ", err)
	}
	return db
}

func ChkRows(rows *sql.Rows) {
	err := rows.Err()
	if err != nil {
		log.Fatalln("Fatal Error: Rows returned error: ", err)
	}
	rows.Close()
}

func BackupDB() {
	cmd := exec.Command("sh", "-c",
		"echo '.dump' | sqlite3 toril.db | "+
			"gzip -c >bak/toril.db.`date +\"%Y-%m-%d\"`.gz")
	err := cmd.Run()
	if err != nil {
		log.Fatalln("Fatal Error: Cannot backup DB: ", err)
	}
	// restore: cat dumpfile.sql | sqlite3 my_database.sqlite
}

func RestoreDB(file string) { // this doesn't work on Mac OS X
	cmd := exec.Command("sh", "-c", "zcat "+file+" | sqlite3 toril.db")
	err := cmd.Run()
	if err != nil {
		log.Fatalln("Fatal Error: Cannot restore DB: ", err)
	}
}
<|endoftext|>


<commit_before>package fam100

import (
	"fmt"
	"hash/crc32"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/rcrowley/go-metrics"
)

type db interface {
	Reset() error
	Init() (err error)
	ChannelRanking(chanID string, limit int) (ranking Rank, err error)
	ChannelCount() (total int, err error)
	Channels() (channels map[string]string, err error)
	ChannelConfig(chanID, key, defaultValue string) (config string, err error)
	GlobalConfig(key, defaultValue string) (config string, err error)

	PlayerCount() (total int, err error)
	PlayerChannelScore(chanID string, playerID PlayerID) (PlayerScore, error)

	// stats command
	incStats(key string) error
	incChannelStats(chanID, key string) error
	incPlayerStats(playerID PlayerID, key string) error
	stats(key string) (interface{}, error)
	channelStats(chanID, key string) (interface{}, error)
	playerStats(playerID, key string) (interface{}, error)

	nextGame(chanID string) (seed int64, nextRound int, err error)
	incRoundPlayed(chanID string) error

	// scores
	saveScore(chanID, chanName string, scores Rank) error
	playerRanking(limit int) (Rank, error)
	playerScore(playerID PlayerID) (ps PlayerScore, err error)
}

var (
	redisPrefix = "fam100"

	gStatsKey, cStatsKey, pStatsKey, cRankKey, pNameKey, pRankKey string
	cNameKey, cConfigKey, gConfigKey                              string
)

var DefaultDB db

func init() {
	DefaultDB = new(RedisDB)
}

func SetRedisPrefix(prefix string) {
	redisPrefix = prefix
	// g: global, c: channel, p:player
	gStatsKey = fmt.Sprintf("%s_stats_", redisPrefix)
	cStatsKey = fmt.Sprintf("%s_chan_stats_", redisPrefix)
	pStatsKey = fmt.Sprintf("%s_player_stats_", redisPrefix)
	cRankKey = fmt.Sprintf("%s_chan_rank_", redisPrefix)

	cNameKey = fmt.Sprintf("%s_chan_name", redisPrefix)
	pNameKey = fmt.Sprintf("%s_player_name", redisPrefix)
	pRankKey = fmt.Sprintf("%s_player_rank", redisPrefix)

	cConfigKey = fmt.Sprintf("%s_chan_config_", redisPrefix)
	gConfigKey = fmt.Sprintf("%s_config", redisPrefix)
}

type RedisDB struct {
	pool *redis.Pool
}

func (r *RedisDB) Reset() error {
	_, err := r.pool.Get().Do("FLUSHALL")
	return err
}

func (r *RedisDB) Init() (err error) {
	r.pool = &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", ":6379")
			if err != nil {
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
	if _, err := r.pool.Get().Do("PING"); err != nil {
		return err
	}
	SetRedisPrefix(redisPrefix)

	go func() {
		redisConnCount := metrics.NewRegisteredGauge("redis.pool.count", metrics.DefaultRegistry)
		tick := time.Tick(5 * time.Second)
		for range tick {
			redisConnCount.Update(int64(r.pool.ActiveCount()))
		}
	}()

	return nil
}

func (r *RedisDB) ChannelCount() (total int, err error) {
	return redis.Int(r.pool.Get().Do("HLEN", cNameKey))
}

func (r *RedisDB) Channels() (channels map[string]string, err error) {
	return redis.StringMap(r.pool.Get().Do("HGETALL", cNameKey))
}

func (r *RedisDB) ChannelConfig(chanID, key, defaultValue string) (config string, err error) {
	rkey := fmt.Sprintf("%s%s", cConfigKey, chanID)
	config, err = redis.String(r.pool.Get().Do("HGET", rkey, key))

	if err != nil || config == "" {
		return defaultValue, err
	}

	return config, nil
}

func (r *RedisDB) GlobalConfig(key, defaultValue string) (config string, err error) {
	rkey := gConfigKey
	config, err = redis.String(r.pool.Get().Do("HGET", rkey, key))

	if err != nil || config == "" {
		return defaultValue, err
	}

	return config, nil
}

func (r *RedisDB) PlayerCount() (total int, err error) {
	return redis.Int(r.pool.Get().Do("HLEN", pNameKey))
}

func (r *RedisDB) nextGame(chanID string) (seed int64, nextRound int, err error) {
	seed = int64(crc32.ChecksumIEEE([]byte(chanID)))
	v, err := r.channelStats(chanID, "played")
	if err != nil {
		return 0, 0, err
	}
	if v == nil {
		return seed, 0, nil
	}
	nextRound, err = redis.Int(v, err)
	if err != nil {
		return 0, 0, err
	}

	return seed, nextRound + 1, nil
}

func (r RedisDB) incStats(key string) error {
	rkey := fmt.Sprintf("%s%s", gStatsKey, key)
	_, err := r.pool.Get().Do("INCR", rkey)

	return err
}

func (r RedisDB) incChannelStats(chanID, key string) error {
	rkey := fmt.Sprintf("%s%s_%s", cStatsKey, key, chanID)
	_, err := r.pool.Get().Do("INCR", rkey)

	return err
}

func (r RedisDB) incPlayerStats(playerID PlayerID, key string) error {
	rkey := fmt.Sprintf("%s%s_%s", pStatsKey, key, playerID)
	_, err := r.pool.Get().Do("INCR", rkey)

	return err
}

func (r RedisDB) stats(key string) (interface{}, error) {
	rkey := fmt.Sprintf("%s%s", gStatsKey, key)
	return r.pool.Get().Do("GET", rkey)
}

func (r RedisDB) channelStats(chanID, key string) (interface{}, error) {
	rkey := fmt.Sprintf("%s%s_%s", cStatsKey, key, chanID)
	return r.pool.Get().Do("GET", rkey)
}

func (r RedisDB) playerStats(playerID, key string) (interface{}, error) {
	rkey := fmt.Sprintf("%s%s_%s", pStatsKey, key, playerID)
	return r.pool.Get().Do("GET", rkey)
}

func (r *RedisDB) incRoundPlayed(chanID string) error {
	return r.incChannelStats(chanID, "played")
}

func (r RedisDB) saveScore(chanID, chanName string, scores Rank) error {
	conn := r.pool.Get()
	for _, score := range scores {
		conn.Send("HSET", cNameKey, chanID, chanName)
		conn.Send("HSET", pNameKey, score.PlayerID, score.Name)
		conn.Send("ZINCRBY", pRankKey, score.Score, score.PlayerID)
		conn.Send("ZINCRBY", cRankKey+chanID, score.Score, score.PlayerID)
	}
	return conn.Flush()
}

func (r RedisDB) ChannelRanking(chanID string, limit int) (ranking Rank, err error) {
	return r.getRanking(cRankKey+chanID, limit)
}

func (r RedisDB) playerRanking(limit int) (Rank, error) {
	return r.getRanking(pRankKey, limit)
}

func (r RedisDB) getRanking(key string, limit int) (ranking Rank, err error) {
	values, err := redis.Values(r.pool.Get().Do("ZREVRANGE", key, 0, limit, "WITHSCORES"))
	if err != nil {
		return nil, err
	}

	ids := make([]interface{}, 0, len(values))
	ids = append(ids, pNameKey)
	pos := 0
	for len(values) > 0 {
		var ps PlayerScore
		values, err = redis.Scan(values, &ps.PlayerID, &ps.Score)
		if err != nil {
			return nil, err
		}
		pos++
		ps.Position = pos
		ids = append(ids, ps.PlayerID)
		ranking = append(ranking, ps)
	}

	// get all name
	if len(ranking) > 0 {
		names, err := redis.Strings(r.pool.Get().Do("HMGET", ids...))
		if err != nil {
			return nil, err
		}
		for i := range ranking {
			ranking[i].Name = names[i]
		}
	}

	return ranking, nil
}

func (r RedisDB) playerScore(playerID PlayerID) (ps PlayerScore, err error) {
	return r.getScore(pRankKey, playerID)
}

func (r RedisDB) PlayerChannelScore(chanID string, playerID PlayerID) (PlayerScore, error) {
	return r.getScore(cRankKey+chanID, playerID)
}

func (r RedisDB) getScore(key string, playerID PlayerID) (ps PlayerScore, err error) {
	ps.PlayerID = playerID
	if ps.Name, err = redis.String(r.pool.Get().Do("HGET", pNameKey, playerID)); err != nil {
		return ps, err
	}
	if ps.Score, err = redis.Int(r.pool.Get().Do("ZSCORE", key, playerID)); err != nil {
		return ps, err
	}
	if ps.Position, err = redis.Int(r.pool.Get().Do("ZRANK", key, playerID)); err != nil {
		return ps, err
	}

	return ps, nil
}

// MemoryDb stores data in non persistence way
type MemoryDB struct {
	Seed   int64
	played int
}

func (m *MemoryDB) Reset() error                                                      { return nil }
func (m *MemoryDB) Init() (err error)                                                 { return nil }
func (m *MemoryDB) ChannelRanking(chanID string, limit int) (ranking Rank, err error) { return nil, nil }
func (m *MemoryDB) ChannelCount() (total int, err error)                              { return 0, nil }
func (m *MemoryDB) Channels() (channels map[string]string, err error)                 { return nil, nil }
func (m *MemoryDB) ChannelConfig(chanID, key, defaultValue string) (string, error)    { return "", nil }
func (m *MemoryDB) GlobalConfig(key, defaultValue string) (string, error)             { return "", nil }
func (m *MemoryDB) PlayerCount() (total int, err error)                               { return 0, nil }
func (m *MemoryDB) incStats(key string) error                                         { return nil }
func (m *MemoryDB) incChannelStats(chanID, key string) error                          { return nil }
func (m *MemoryDB) incPlayerStats(playerID PlayerID, key string) error                { return nil }
func (m *MemoryDB) stats(key string) (interface{}, error)                             { return nil, nil }
func (m *MemoryDB) channelStats(chanID, key string) (interface{}, error)              { return nil, nil }
func (m *MemoryDB) playerStats(playerID, key string) (interface{}, error)             { return nil, nil }
func (m *MemoryDB) saveScore(chanID, chanName string, scores Rank) error              { return nil }
func (m *MemoryDB) playerRanking(limit int) (Rank, error)                             { return nil, nil }
func (m *MemoryDB) playerScore(playerID PlayerID) (ps PlayerScore, err error) {
	return PlayerScore{}, nil
}
func (m *MemoryDB) PlayerChannelScore(chanID string, playerID PlayerID) (PlayerScore, error) {
	return PlayerScore{}, nil
}

func (m *MemoryDB) nextGame(chanID string) (seed int64, nextRound int, err error) {
	return m.Seed, m.played + 1, nil
}
func (m *MemoryDB) incRoundPlayed(chanID string) error {
	m.played++
	return nil
}
<commit_msg>Fix redis connection leak<commit_after>package fam100

import (
	"fmt"
	"hash/crc32"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/rcrowley/go-metrics"
)

type db interface {
	Reset() error
	Init() (err error)
	ChannelRanking(chanID string, limit int) (ranking Rank, err error)
	ChannelCount() (total int, err error)
	Channels() (channels map[string]string, err error)
	ChannelConfig(chanID, key, defaultValue string) (config string, err error)
	GlobalConfig(key, defaultValue string) (config string, err error)

	PlayerCount() (total int, err error)
	PlayerChannelScore(chanID string, playerID PlayerID) (PlayerScore, error)

	// stats command
	incStats(key string) error
	incChannelStats(chanID, key string) error
	incPlayerStats(playerID PlayerID, key string) error
	stats(key string) (interface{}, error)
	channelStats(chanID, key string) (interface{}, error)
	playerStats(playerID, key string) (interface{}, error)

	nextGame(chanID string) (seed int64, nextRound int, err error)
	incRoundPlayed(chanID string) error

	// scores
	saveScore(chanID, chanName string, scores Rank) error
	playerRanking(limit int) (Rank, error)
	playerScore(playerID PlayerID) (ps PlayerScore, err error)
}

var (
	redisPrefix = "fam100"

	gStatsKey, cStatsKey, pStatsKey, cRankKey, pNameKey, pRankKey string
	cNameKey, cConfigKey, gConfigKey                              string
)

var DefaultDB db

func init() {
	DefaultDB = new(RedisDB)
}

func SetRedisPrefix(prefix string) {
	redisPrefix = prefix
	// g: global, c: channel, p:player
	gStatsKey = fmt.Sprintf("%s_stats_", redisPrefix)
	cStatsKey = fmt.Sprintf("%s_chan_stats_", redisPrefix)
	pStatsKey = fmt.Sprintf("%s_player_stats_", redisPrefix)
	cRankKey = fmt.Sprintf("%s_chan_rank_", redisPrefix)

	cNameKey = fmt.Sprintf("%s_chan_name", redisPrefix)
	pNameKey = fmt.Sprintf("%s_player_name", redisPrefix)
	pRankKey = fmt.Sprintf("%s_player_rank", redisPrefix)

	cConfigKey = fmt.Sprintf("%s_chan_config_", redisPrefix)
	gConfigKey = fmt.Sprintf("%s_config", redisPrefix)
}

type RedisDB struct {
	pool *redis.Pool
}

func (r *RedisDB) Reset() error {
	conn := r.pool.Get()
	defer conn.Close()

	_, err := conn.Do("FLUSHALL")
	return err
}

func (r *RedisDB) Init() (err error) {
	r.pool = &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", ":6379")
			if err != nil {
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
	if _, err := r.pool.Get().Do("PING"); err != nil {
		return err
	}
	SetRedisPrefix(redisPrefix)

	go func() {
		redisConnCount := metrics.NewRegisteredGauge("redis.pool.count", metrics.DefaultRegistry)
		tick := time.Tick(5 * time.Second)
		for range tick {
			redisConnCount.Update(int64(r.pool.ActiveCount()))
		}
	}()

	return nil
}

func (r *RedisDB) ChannelCount() (total int, err error) {
	conn := r.pool.Get()
	defer conn.Close()
	return redis.Int(conn.Do("HLEN", cNameKey))
}

func (r *RedisDB) Channels() (channels map[string]string, err error) {
	conn := r.pool.Get()
	defer conn.Close()
	return redis.StringMap(conn.Do("HGETALL", cNameKey))
}

func (r *RedisDB) ChannelConfig(chanID, key, defaultValue string) (config string, err error) {
	conn := r.pool.Get()
	defer conn.Close()

	rkey := fmt.Sprintf("%s%s", cConfigKey, chanID)
	config, err = redis.String(conn.Do("HGET", rkey, key))

	if err != nil || config == "" {
		return defaultValue, err
	}

	return config, nil
}

func (r *RedisDB) GlobalConfig(key, defaultValue string) (config string, err error) {
	conn := r.pool.Get()
	defer conn.Close()

	rkey := gConfigKey
	config, err = redis.String(conn.Do("HGET", rkey, key))

	if err != nil || config == "" {
		return defaultValue, err
	}

	return config, nil
}

func (r *RedisDB) PlayerCount() (total int, err error) {
	conn := r.pool.Get()
	defer conn.Close()

	return redis.Int(conn.Do("HLEN", pNameKey))
}

func (r *RedisDB) nextGame(chanID string) (seed int64, nextRound int, err error) {
	seed = int64(crc32.ChecksumIEEE([]byte(chanID)))
	v, err := r.channelStats(chanID, "played")
	if err != nil {
		return 0, 0, err
	}
	if v == nil {
		return seed, 0, nil
	}
	nextRound, err = redis.Int(v, err)
	if err != nil {
		return 0, 0, err
	}

	return seed, nextRound + 1, nil
}

func (r RedisDB) incStats(key string) error {
	conn := r.pool.Get()
	defer conn.Close()

	rkey := fmt.Sprintf("%s%s", gStatsKey, key)
	_, err := conn.Do("INCR", rkey)

	return err
}

func (r RedisDB) incChannelStats(chanID, key string) error {
	conn := r.pool.Get()
	defer conn.Close()

	rkey := fmt.Sprintf("%s%s_%s", cStatsKey, key, chanID)
	_, err := conn.Do("INCR", rkey)

	return err
}

func (r RedisDB) incPlayerStats(playerID PlayerID, key string) error {
	conn := r.pool.Get()
	defer conn.Close()

	rkey := fmt.Sprintf("%s%s_%s", pStatsKey, key, playerID)
	_, err := conn.Do("INCR", rkey)

	return err
}

func (r RedisDB) stats(key string) (interface{}, error) {
	conn := r.pool.Get()
	defer conn.Close()

	rkey := fmt.Sprintf("%s%s", gStatsKey, key)
	return conn.Do("GET", rkey)
}

func (r RedisDB) channelStats(chanID, key string) (interface{}, error) {
	conn := r.pool.Get()
	defer conn.Close()

	rkey := fmt.Sprintf("%s%s_%s", cStatsKey, key, chanID)
	return conn.Do("GET", rkey)
}

func (r RedisDB) playerStats(playerID, key string) (interface{}, error) {
	conn := r.pool.Get()
	defer conn.Close()

	rkey := fmt.Sprintf("%s%s_%s", pStatsKey, key, playerID)
	return conn.Do("GET", rkey)
}

func (r *RedisDB) incRoundPlayed(chanID string) error {
	return r.incChannelStats(chanID, "played")
}

func (r RedisDB) saveScore(chanID, chanName string, scores Rank) error {
	conn := r.pool.Get()
	defer conn.Close()
	for _, score := range scores {
		conn.Send("HSET", cNameKey, chanID, chanName)
		conn.Send("HSET", pNameKey, score.PlayerID, score.Name)
		conn.Send("ZINCRBY", pRankKey, score.Score, score.PlayerID)
		conn.Send("ZINCRBY", cRankKey+chanID, score.Score, score.PlayerID)
	}
	return conn.Flush()
}

func (r RedisDB) ChannelRanking(chanID string, limit int) (ranking Rank, err error) {
	return r.getRanking(cRankKey+chanID, limit)
}

func (r RedisDB) playerRanking(limit int) (Rank, error) {
	return r.getRanking(pRankKey, limit)
}

func (r RedisDB) getRanking(key string, limit int) (ranking Rank, err error) {
	conn := r.pool.Get()
	defer conn.Close()

	values, err := redis.Values(conn.Do("ZREVRANGE", key, 0, limit, "WITHSCORES"))
	if err != nil {
		return nil, err
	}

	ids := make([]interface{}, 0, len(values))
	ids = append(ids, pNameKey)
	pos := 0
	for len(values) > 0 {
		var ps PlayerScore
		values, err = redis.Scan(values, &ps.PlayerID, &ps.Score)
		if err != nil {
			return nil, err
		}
		pos++
		ps.Position = pos
		ids = append(ids, ps.PlayerID)
		ranking = append(ranking, ps)
	}

	// get all name
	if len(ranking) > 0 {
		names, err := redis.Strings(conn.Do("HMGET", ids...))
		if err != nil {
			return nil, err
		}
		for i := range ranking {
			ranking[i].Name = names[i]
		}
	}

	return ranking, nil
}

func (r RedisDB) playerScore(playerID PlayerID) (ps PlayerScore, err error) {
	return r.getScore(pRankKey, playerID)
}

func (r RedisDB) PlayerChannelScore(chanID string, playerID PlayerID) (PlayerScore, error) {
	return r.getScore(cRankKey+chanID, playerID)
}

func (r RedisDB) getScore(key string, playerID PlayerID) (ps PlayerScore, err error) {
	conn := r.pool.Get()
	defer conn.Close()

	ps.PlayerID = playerID
	if ps.Name, err = redis.String(conn.Do("HGET", pNameKey, playerID)); err != nil {
		return ps, err
	}
	if ps.Score, err = redis.Int(conn.Do("ZSCORE", key, playerID)); err != nil {
		return ps, err
	}
	if ps.Position, err = redis.Int(conn.Do("ZRANK", key, playerID)); err != nil {
		return ps, err
	}

	return ps, nil
}

// MemoryDb stores data in non persistence way
type MemoryDB struct {
	Seed   int64
	played int
}

func (m *MemoryDB) Reset() error                                                      { return nil }
func (m *MemoryDB) Init() (err error)                                                 { return nil }
func (m *MemoryDB) ChannelRanking(chanID string, limit int) (ranking Rank, err error) { return nil, nil }
func (m *MemoryDB) ChannelCount() (total int, err error)                              { return 0, nil }
func (m *MemoryDB) Channels() (channels map[string]string, err error)                 { return nil, nil }
func (m *MemoryDB) ChannelConfig(chanID, key, defaultValue string) (string, error)    { return "", nil }
func (m *MemoryDB) GlobalConfig(key, defaultValue string) (string, error)             { return "", nil }
func (m *MemoryDB) PlayerCount() (total int, err error)                               { return 0, nil }
func (m *MemoryDB) incStats(key string) error                                         { return nil }
func (m *MemoryDB) incChannelStats(chanID, key string) error                          { return nil }
func (m *MemoryDB) incPlayerStats(playerID PlayerID, key string) error                { return nil }
func (m *MemoryDB) stats(key string) (interface{}, error)                             { return nil, nil }
func (m *MemoryDB) channelStats(chanID, key string) (interface{}, error)              { return nil, nil }
func (m *MemoryDB) playerStats(playerID, key string) (interface{}, error)             { return nil, nil }
func (m *MemoryDB) saveScore(chanID, chanName string, scores Rank) error              { return nil }
func (m *MemoryDB) playerRanking(limit int) (Rank, error)                             { return nil, nil }
func (m *MemoryDB) playerScore(playerID PlayerID) (ps PlayerScore, err error) {
	return PlayerScore{}, nil
}
func (m *MemoryDB) PlayerChannelScore(chanID string, playerID PlayerID) (PlayerScore, error) {
	return PlayerScore{}, nil
}

func (m *MemoryDB) nextGame(chanID string) (seed int64, nextRound int, err error) {
	return m.Seed, m.played + 1, nil
}
func (m *MemoryDB) incRoundPlayed(chanID string) error {
	m.played++
	return nil
}
<|endoftext|>


<commit_before>package linux

import (
	"io/ioutil"
	"reflect"
	"strconv"
	"strings"
)

type NetStat struct {
	// TcpExt
	SyncookiesSent            uint64 `json:"syncookie_sent"`
	SyncookiesRecv            uint64 `json:"syncookies_recv"`
	SyncookiesFailed          uint64 `json:"syncookies_failed"`
	EmbryonicRsts             uint64 `json:"embryonic_rsts"`
	PruneCalled               uint64 `json:"prune_called"`
	RcvPruned                 uint64 `json:"rcv_pruned"`
	OfoPruned                 uint64 `json:"ofo_pruned"`
	OutOfWindowIcmps          uint64 `json:"out_of_window_icmps"`
	LockDroppedIcmps          uint64 `json:"lock_dropped_icmps"`
	ArpFilter                 uint64 `json:"arp_filter"`
	TW                        uint64 `json:"tw"`
	TWRecycled                uint64 `json:"tw_recycled"`
	TWKilled                  uint64 `json:"tw_killed"`
	PAWSPassive               uint64 `json:"paws_passive"`
	PAWSActive                uint64 `json:"paws_active"`
	PAWSEstab                 uint64 `json:"paws_estab"`
	DelayedACKs               uint64 `json:"delayed_acks"`
	DelayedACKLocked          uint64 `json:"delayed_ack_locked"`
	DelayedACKLost            uint64 `json:"delayed_ack_lost"`
	ListenOverflows           uint64 `json:"listen_overflows"`
	ListenDrops               uint64 `json:"listen_drops"`
	TCPPrequeued              uint64 `json:"tcp_prequeued"`
	TCPDirectCopyFromBacklog  uint64 `json:"tcp_direct_copy_from_backlog"`
	TCPDirectCopyFromPrequeue uint64 `json:"tcp_direct_copy_from_prequeue"`
	TCPPrequeueDropped        uint64 `json:"tcp_prequeue_dropped"`
	TCPHPHits                 uint64 `json:"tcp_hp_hits"`
	TCPHPHitsToUser           uint64 `json:"tcp_hp_hits_to_user"`
	TCPPureAcks               uint64 `json:"tcp_pure_acks"`
	TCPHPAcks                 uint64 `json:"tcp_hp_acks"`
	TCPRenoRecovery           uint64 `json:"tcp_reno_recovery"`
	TCPSackRecovery           uint64 `json:"tcp_sack_recovery"`
	TCPSACKReneging           uint64 `json:"tcp_sack_reneging"`
	TCPFACKReorder            uint64 `json:"tcp_fack_reorder"`
	TCPSACKReorder            uint64 `json:"tcp_sack_reorder"`
	TCPRenoReorder            uint64 `json:"tcp_reno_reorder"`
	TCPTSReorder              uint64 `json:"tcp_ts_reorder"`
	TCPFullUndo               uint64 `json:"tcp_full_undo"`
	TCPPartialUndo            uint64 `json:"tcp_partial_undo"`
	TCPDSACKUndo              uint64 `json:"tcp_dsack_undo"`
	TCPLossUndo               uint64 `json:"tcp_loss_undo"`
	TCPLoss                   uint64 `json:"tcp_loss"`
	TCPLostRetransmit         uint64 `json:"tcp_lost_retransmit"`
	TCPRenoFailures           uint64 `json:"tcp_reno_failures"`
	TCPSackFailures           uint64 `json:"tcp_sack_failures"`
	TCPLossFailures           uint64 `json:"tcp_loss_failures"`
	TCPFastRetrans            uint64 `json:"tcp_fast_retrans"`
	TCPForwardRetrans         uint64 `json:"tcp_forward_retrans"`
	TCPSlowStartRetrans       uint64 `json:"tcp_slow_start_retrans"`
	TCPTimeouts               uint64 `json:"tcp_timeouts"`
	TCPLossProbes             uint64 `json:"tcp_loss_probes"`
	TCPLossProbeRecovery      uint64 `json:"tcp_loss_probe_recovery"`
	TCPRenoRecoveryFail       uint64 `json:"tcp_reno_recovery_fail"`
	TCPSackRecoveryFail       uint64 `json:"tcp_sack_recovery_fail"`
	TCPSchedulerFailed        uint64 `json:"tcp_scheduler_failed"`
	TCPRcvCollapsed           uint64 `json:"tcp_rcv_collapsed"`
	TCPDSACKOldSent           uint64 `json:"tcp_dsack_old_sent"`
	TCPDSACKOfoSent           uint64 `json:"tcp_dsack_ofo_sent"`
	TCPDSACKRecv              uint64 `json:"tcp_dsack_recv"`
	TCPDSACKOfoRecv           uint64 `json:"tcp_dsack_ofo_recv"`
	TCPAbortOnSyn             uint64 `json:"tcp_abort_on_syn"`
	TCPAbortOnData            uint64 `json:"tcp_abort_on_data"`
	TCPAbortOnClose           uint64 `json:"tcp_abort_on_close"`
	TCPAbortOnMemory          uint64 `json:"tcp_abort_on_memory"`
	TCPAbortOnTimeout         uint64 `json:"tcp_abort_on_timeout"`
	TCPAbortOnLinger          uint64 `json:"tcp_abort_on_linger"`
	TCPAbortFailed            uint64 `json:"tcp_abort_failed"`
	TCPMemoryPressures        uint64 `json:"tcp_memory_pressures"`
	TCPSACKDiscard            uint64 `json:"tcp_sack_discard"`
	TCPDSACKIgnoredOld        uint64 `json:"tcp_dsack_ignored_old"`
	TCPDSACKIgnoredNoUndo     uint64 `json:"tcp_dsack_ignored_no_undo"`
	TCPSpuriousRTOs           uint64 `json:"tcp_spurious_rtos"`
	TCPMD5NotFound            uint64 `json:"tcp_md5_not_found"`
	TCPMD5Unexpected          uint64 `json:"tcp_md5_unexpected"`
	TCPSackShifted            uint64 `json:"tcp_sack_shifted"`
	TCPSackMerged             uint64 `json:"tcp_sack_merged"`
	TCPSackShiftFallback      uint64 `json:"tcp_sack_shift_fallback"`
	TCPBacklogDrop            uint64 `json:"tcp_backlog_drop"`
	TCPMinTTLDrop             uint64 `json:"tcp_min_ttl_drop"`
	TCPDeferAcceptDrop        uint64 `json:"tcp_defer_accept_drop"`
	IPReversePathFilter       uint64 `json:"ip_reverse_path_filter"`
	TCPTimeWaitOverflow       uint64 `json:"tcp_time_wait_overflow"`
	TCPReqQFullDoCookies      uint64 `json:"tcp_req_q_full_do_cookies"`
	TCPReqQFullDrop           uint64 `json:"tcp_req_q_full_drop"`
	TCPRetransFail            uint64 `json:"tcp_retrans_fail"`
	TCPRcvCoalesce            uint64 `json:"tcp_rcv_coalesce"`
	TCPOFOQueue               uint64 `json:"tcp_ofo_drop"`
	TCPOFODrop                uint64 `json:"tcp_ofo_drop"`
	TCPOFOMerge               uint64 `json:"tcp_ofo_merge"`
	TCPChallengeACK           uint64 `json:"tcp_challenge_ack"`
	TCPSYNChallenge           uint64 `json:"tcp_syn_challenge"`
	TCPFastOpenActive         uint64 `json:"tcp_fast_open_active"`
	TCPFastOpenActiveFail     uint64 `json:"tcp_fast_open_active_fail"`
	TCPFastOpenPassive        uint64 `json:"tcp_fast_open_passive"`
	TCPFastOpenPassiveFail    uint64 `json:"tcp_fast_open_passive_fail"`
	TCPFastOpenListenOverflow uint64 `json:"tcp_fast_open_listen_overflow"`
	TCPFastOpenCookieReqd     uint64 `json:"tcp_fast_open_cookie_reqd"`
	TCPSpuriousRtxHostQueues  uint64 `json:"tcp_spurious_rtx_host_queues"`
	BusyPollRxPackets         uint64 `json:"busy_poll_rx_packets"`
	TCPAutoCorking            uint64 `json:"tcp_auto_corking"`
	TCPFromZeroWindowAdv      uint64 `json:"tcp_from_zero_window_adv"`
	TCPToZeroWindowAdv        uint64 `json:"tcp_to_zero_window_adv"`
	TCPWantZeroWindowAdv      uint64 `json:"tcp_want_zero_window_adv"`
	TCPSynRetrans             uint64 `json:"tcp_syn_retrans"`
	TCPOrigDataSent           uint64 `json:"tcp_orig_data_sent"`
	// IpExt
	InNoRoutes      uint64 `json:"in_no_routes"`
	InTruncatedPkts uint64 `json:"in_truncated_pkts"`
	InMcastPkts     uint64 `json:"in_mcast_pkts"`
	OutMcastPkts    uint64 `json:"out_mcast_pkts"`
	InBcastPkts     uint64 `json:"in_bcast_pkts"`
	OutBcastPkts    uint64 `json:"out_bcast_pkts"`
	InOctets        uint64 `json:"in_octets"`
	OutOctets       uint64 `json:"out_octets"`
	InMcastOctets   uint64 `json:"in_mcast_octets"`
	OutMcastOctets  uint64 `json:"out_mcast_octets"`
	InBcastOctets   uint64 `json:"in_bcast_octets"`
	OutBcastOctets  uint64 `json:"out_bcast_octets"`
	InCsumErrors    uint64 `json:"in_csum_errors"`
	InNoECTPkts     uint64 `json:"in_no_ect_pkts"`
	InECT1Pkts      uint64 `json:"in_ect1_pkts"`
	InECT0Pkts      uint64 `json:"in_ect0_pkts"`
	InCEPkts        uint64 `json:"in_ce_pkts"`
}

func ReadNetStat(path string) (*NetStat, error) {
	data, err := ioutil.ReadFile(path)

	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(data), "\n")

	// Maps a netstat metric to its value (i.e. SyncookiesSent --> 0)
	statMap := make(map[string]string)

	// patterns
	// TcpExt: SyncookiesSent SyncookiesRecv SyncookiesFailed... <-- header
	// TcpExt: 0 0 1764... <-- values

	for i := 1; i < len(lines); i = i + 2 {
		headers := strings.Fields(lines[i-1][strings.Index(lines[i-1], ":")+1:])
		values := strings.Fields(lines[i][strings.Index(lines[i], ":")+1:])

		for j, header := range headers {
			statMap[header] = values[j]
		}
	}

	var netstat NetStat = NetStat{}

	elem := reflect.ValueOf(&netstat).Elem()
	typeOfElem := elem.Type()

	for i := 0; i < elem.NumField(); i++ {
		if val, ok := statMap[typeOfElem.Field(i).Name]; ok {
			parsedVal, _ := strconv.ParseUint(val, 10, 64)
			elem.Field(i).SetUint(parsedVal)
		}
	}

	return &netstat, nil
}
<commit_msg>Adding ReadNetStatFromBytes and refactor<commit_after>package linux

import (
	"bufio"
	"bytes"
	"io/ioutil"
	"reflect"
	"strconv"
	"strings"
)

type (
	NetStat struct {
		// TcpExt
		SyncookiesSent            uint64 `json:"syncookie_sent"`
		SyncookiesRecv            uint64 `json:"syncookies_recv"`
		SyncookiesFailed          uint64 `json:"syncookies_failed"`
		EmbryonicRsts             uint64 `json:"embryonic_rsts"`
		PruneCalled               uint64 `json:"prune_called"`
		RcvPruned                 uint64 `json:"rcv_pruned"`
		OfoPruned                 uint64 `json:"ofo_pruned"`
		OutOfWindowIcmps          uint64 `json:"out_of_window_icmps"`
		LockDroppedIcmps          uint64 `json:"lock_dropped_icmps"`
		ArpFilter                 uint64 `json:"arp_filter"`
		TW                        uint64 `json:"tw"`
		TWRecycled                uint64 `json:"tw_recycled"`
		TWKilled                  uint64 `json:"tw_killed"`
		PAWSPassive               uint64 `json:"paws_passive"`
		PAWSActive                uint64 `json:"paws_active"`
		PAWSEstab                 uint64 `json:"paws_estab"`
		DelayedACKs               uint64 `json:"delayed_acks"`
		DelayedACKLocked          uint64 `json:"delayed_ack_locked"`
		DelayedACKLost            uint64 `json:"delayed_ack_lost"`
		ListenOverflows           uint64 `json:"listen_overflows"`
		ListenDrops               uint64 `json:"listen_drops"`
		TCPPrequeued              uint64 `json:"tcp_prequeued"`
		TCPDirectCopyFromBacklog  uint64 `json:"tcp_direct_copy_from_backlog"`
		TCPDirectCopyFromPrequeue uint64 `json:"tcp_direct_copy_from_prequeue"`
		TCPPrequeueDropped        uint64 `json:"tcp_prequeue_dropped"`
		TCPHPHits                 uint64 `json:"tcp_hp_hits"`
		TCPHPHitsToUser           uint64 `json:"tcp_hp_hits_to_user"`
		TCPPureAcks               uint64 `json:"tcp_pure_acks"`
		TCPHPAcks                 uint64 `json:"tcp_hp_acks"`
		TCPRenoRecovery           uint64 `json:"tcp_reno_recovery"`
		TCPSackRecovery           uint64 `json:"tcp_sack_recovery"`
		TCPSACKReneging           uint64 `json:"tcp_sack_reneging"`
		TCPFACKReorder            uint64 `json:"tcp_fack_reorder"`
		TCPSACKReorder            uint64 `json:"tcp_sack_reorder"`
		TCPRenoReorder            uint64 `json:"tcp_reno_reorder"`
		TCPTSReorder              uint64 `json:"tcp_ts_reorder"`
		TCPFullUndo               uint64 `json:"tcp_full_undo"`
		TCPPartialUndo            uint64 `json:"tcp_partial_undo"`
		TCPDSACKUndo              uint64 `json:"tcp_dsack_undo"`
		TCPLossUndo               uint64 `json:"tcp_loss_undo"`
		TCPLoss                   uint64 `json:"tcp_loss"`
		TCPLostRetransmit         uint64 `json:"tcp_lost_retransmit"`
		TCPRenoFailures           uint64 `json:"tcp_reno_failures"`
		TCPSackFailures           uint64 `json:"tcp_sack_failures"`
		TCPLossFailures           uint64 `json:"tcp_loss_failures"`
		TCPFastRetrans            uint64 `json:"tcp_fast_retrans"`
		TCPForwardRetrans         uint64 `json:"tcp_forward_retrans"`
		TCPSlowStartRetrans       uint64 `json:"tcp_slow_start_retrans"`
		TCPTimeouts               uint64 `json:"tcp_timeouts"`
		TCPLossProbes             uint64 `json:"tcp_loss_probes"`
		TCPLossProbeRecovery      uint64 `json:"tcp_loss_probe_recovery"`
		TCPRenoRecoveryFail       uint64 `json:"tcp_reno_recovery_fail"`
		TCPSackRecoveryFail       uint64 `json:"tcp_sack_recovery_fail"`
		TCPSchedulerFailed        uint64 `json:"tcp_scheduler_failed"`
		TCPRcvCollapsed           uint64 `json:"tcp_rcv_collapsed"`
		TCPDSACKOldSent           uint64 `json:"tcp_dsack_old_sent"`
		TCPDSACKOfoSent           uint64 `json:"tcp_dsack_ofo_sent"`
		TCPDSACKRecv              uint64 `json:"tcp_dsack_recv"`
		TCPDSACKOfoRecv           uint64 `json:"tcp_dsack_ofo_recv"`
		TCPAbortOnSyn             uint64 `json:"tcp_abort_on_syn"`
		TCPAbortOnData            uint64 `json:"tcp_abort_on_data"`
		TCPAbortOnClose           uint64 `json:"tcp_abort_on_close"`
		TCPAbortOnMemory          uint64 `json:"tcp_abort_on_memory"`
		TCPAbortOnTimeout         uint64 `json:"tcp_abort_on_timeout"`
		TCPAbortOnLinger          uint64 `json:"tcp_abort_on_linger"`
		TCPAbortFailed            uint64 `json:"tcp_abort_failed"`
		TCPMemoryPressures        uint64 `json:"tcp_memory_pressures"`
		TCPSACKDiscard            uint64 `json:"tcp_sack_discard"`
		TCPDSACKIgnoredOld        uint64 `json:"tcp_dsack_ignored_old"`
		TCPDSACKIgnoredNoUndo     uint64 `json:"tcp_dsack_ignored_no_undo"`
		TCPSpuriousRTOs           uint64 `json:"tcp_spurious_rtos"`
		TCPMD5NotFound            uint64 `json:"tcp_md5_not_found"`
		TCPMD5Unexpected          uint64 `json:"tcp_md5_unexpected"`
		TCPSackShifted            uint64 `json:"tcp_sack_shifted"`
		TCPSackMerged             uint64 `json:"tcp_sack_merged"`
		TCPSackShiftFallback      uint64 `json:"tcp_sack_shift_fallback"`
		TCPBacklogDrop            uint64 `json:"tcp_backlog_drop"`
		TCPMinTTLDrop             uint64 `json:"tcp_min_ttl_drop"`
		TCPDeferAcceptDrop        uint64 `json:"tcp_defer_accept_drop"`
		IPReversePathFilter       uint64 `json:"ip_reverse_path_filter"`
		TCPTimeWaitOverflow       uint64 `json:"tcp_time_wait_overflow"`
		TCPReqQFullDoCookies      uint64 `json:"tcp_req_q_full_do_cookies"`
		TCPReqQFullDrop           uint64 `json:"tcp_req_q_full_drop"`
		TCPRetransFail            uint64 `json:"tcp_retrans_fail"`
		TCPRcvCoalesce            uint64 `json:"tcp_rcv_coalesce"`
		TCPOFOQueue               uint64 `json:"tcp_ofo_drop"`
		TCPOFODrop                uint64 `json:"tcp_ofo_drop"`
		TCPOFOMerge               uint64 `json:"tcp_ofo_merge"`
		TCPChallengeACK           uint64 `json:"tcp_challenge_ack"`
		TCPSYNChallenge           uint64 `json:"tcp_syn_challenge"`
		TCPFastOpenActive         uint64 `json:"tcp_fast_open_active"`
		TCPFastOpenActiveFail     uint64 `json:"tcp_fast_open_active_fail"`
		TCPFastOpenPassive        uint64 `json:"tcp_fast_open_passive"`
		TCPFastOpenPassiveFail    uint64 `json:"tcp_fast_open_passive_fail"`
		TCPFastOpenListenOverflow uint64 `json:"tcp_fast_open_listen_overflow"`
		TCPFastOpenCookieReqd     uint64 `json:"tcp_fast_open_cookie_reqd"`
		TCPSpuriousRtxHostQueues  uint64 `json:"tcp_spurious_rtx_host_queues"`
		BusyPollRxPackets         uint64 `json:"busy_poll_rx_packets"`
		TCPAutoCorking            uint64 `json:"tcp_auto_corking"`
		TCPFromZeroWindowAdv      uint64 `json:"tcp_from_zero_window_adv"`
		TCPToZeroWindowAdv        uint64 `json:"tcp_to_zero_window_adv"`
		TCPWantZeroWindowAdv      uint64 `json:"tcp_want_zero_window_adv"`
		TCPSynRetrans             uint64 `json:"tcp_syn_retrans"`
		TCPOrigDataSent           uint64 `json:"tcp_orig_data_sent"`
		// IpExt
		InNoRoutes      uint64 `json:"in_no_routes"`
		InTruncatedPkts uint64 `json:"in_truncated_pkts"`
		InMcastPkts     uint64 `json:"in_mcast_pkts"`
		OutMcastPkts    uint64 `json:"out_mcast_pkts"`
		InBcastPkts     uint64 `json:"in_bcast_pkts"`
		OutBcastPkts    uint64 `json:"out_bcast_pkts"`
		InOctets        uint64 `json:"in_octets"`
		OutOctets       uint64 `json:"out_octets"`
		InMcastOctets   uint64 `json:"in_mcast_octets"`
		OutMcastOctets  uint64 `json:"out_mcast_octets"`
		InBcastOctets   uint64 `json:"in_bcast_octets"`
		OutBcastOctets  uint64 `json:"out_bcast_octets"`
		InCsumErrors    uint64 `json:"in_csum_errors"`
		InNoECTPkts     uint64 `json:"in_no_ect_pkts"`
		InECT1Pkts      uint64 `json:"in_ect1_pkts"`
		InECT0Pkts      uint64 `json:"in_ect0_pkts"`
		InCEPkts        uint64 `json:"in_ce_pkts"`
	}
)

func ReadNetStatFromBytes(data []byte) (*NetStat, error) {

	stat := &NetStat{}
	v := reflect.ValueOf(stat).Elem()

	for s := bufio.NewScanner(bytes.NewReader(data)); s.Scan(); {

		fields := strings.Fields(s.Text())
		if !s.Scan() || s.Err() != nil {
			return nil, s.Err()
		}
		values := strings.Fields(s.Text())
		if s.Err() != nil {
			return nil, s.Err()
		}

		// trim off row title
		if len(fields) > 1 && len(values) > 1 {
			fields, values = fields[1:], values[1:]
		}

		for i, name := range fields {
			if f := v.FieldByName(name); f.CanSet() {
				n, _ := strconv.ParseUint(values[i], 10, 64)
				f.SetUint(n)
			}
		}

	}

	return stat, nil

}

func ReadNetStat(path string) (*NetStat, error) {

	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	return ReadNetStatFromBytes(data)

}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"gopkg.in/logex.v1"
)

var DefaultClient = &http.Client{}

type TaskConfig struct {
	MaxSpeed   int64
	Clean      bool
	Progress   bool
	ShowRealSp bool
	Proxy      []string
}

func (t *TaskConfig) init() {
}

type DnTask struct {
	*TaskConfig
	source *url.URL
	Meta   *Meta

	file     *os.File
	writeOp  chan *writeOp
	stopChan chan struct{}

	wg    sync.WaitGroup
	start time.Time
	sync.Mutex
	rateLimit *RateLimit

	downloadPerSecond int64

	l *Liner
}

func NewDnTaskAuto(url_, pwd string, bit uint, cfg *TaskConfig) (*DnTask, error) {
	_, err := os.Stat(url_)
	if !cfg.Clean && err == nil {
		if meta, _ := NewMetaFormFile(url_); meta != nil {
			logex.Info("downloading form", meta.Source)
			return NewDnTask(meta.Source, pwd, meta.BlkBit, cfg)
		}
	}

	return NewDnTask(url_, pwd, bit, cfg)
}

func NewDnTask(url_, pwd string, bit uint, cfg *TaskConfig) (*DnTask, error) {
	if url_ == "" {
		return nil, logex.NewError("url is empty")
	}
	if cfg == nil {
		cfg = new(TaskConfig)
	}
	cfg.init()

	source, err := url.Parse(url_)
	if err != nil {
		return nil, logex.Trace(err)
	}
	meta, err := NewMeta(pwd, url_, bit, cfg.Clean)
	if err != nil {
		return nil, logex.Trace(err)
	}

	dn := &DnTask{
		TaskConfig: cfg,
		rateLimit:  NewRateLimit(cfg.MaxSpeed),
		source:     source,
		Meta:       meta,
		writeOp:    make(chan *writeOp, 1<<3),
		stopChan:   make(chan struct{}),
		start:      time.Now(),
		l:          NewLiner(os.Stderr),
	}
	if cfg.Clean {
		os.Remove(dn.Meta.targetPath())
	}

	if err = dn.Meta.retrieveFromDisk(cfg.Proxy); err != nil {
		dn.Meta.Remove()
		return nil, logex.Trace(err)
	}

	if err = dn.Meta.Sync(); err != nil {
		return nil, logex.Trace(err)
	}

	if err = dn.openFile(); err != nil {
		return nil, logex.Trace(err)
	}

	go dn.ioloop()
	go dn.progress()
	return dn, nil
}

func (d *DnTask) openFile() error {
	f, err := os.OpenFile(d.Meta.targetPath(), os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		return logex.Trace(err)
	}
	d.file = f
	return nil
}

type writeOpReply struct {
	N   int
	Err error
}

type writeOp struct {
	Offset int64
	Buf    []byte
	Reply  chan *writeOpReply
}

func (d *DnTask) ioloop() {
	var w *writeOp
	for {
		select {
		case w = <-d.writeOp:
		case <-d.stopChan:
			return
		}
		n, err := d.file.WriteAt(w.Buf, w.Offset)
		if err != nil {
			d.file.Close()
			d.file = nil
			if err := d.openFile(); err != nil {
				panic(err)
			}
		}
		if d.MaxSpeed > 0 {
			d.rateLimit.Process(n)
		}
		w.Reply <- &writeOpReply{n, logex.Trace(err)}
	}
}

// [start, end)
func (d *DnTask) allocDnBlk(off int) (idx int, start, end int64) {
	d.Lock()
	defer d.Unlock()

	for i := off; i < len(d.Meta.Blocks); i++ {
		blk := d.Meta.Blocks[i]
		if blk == nil {
			blk = NewBlock()
			d.Meta.Blocks[i] = blk
		}
		if blk.State != STATE_INIT {
			continue
		}
		blk.State = STATE_PROCESS
		offset := int64(i << d.Meta.BlkBit)
		start = offset + int64(blk.Written)
		end = int64((i + 1) << d.Meta.BlkBit)
		if end > d.Meta.FileSize {
			end = d.Meta.FileSize
		}
		return i, start, end
	}
	return -1, -1, -1
}

func setRange(h http.Header, start, end int64) {
	h.Set(H_RANGE, fmt.Sprintf("bytes=%d-%d", start, end-1))
}

func (d *DnTask) checkWritten(written, start, end int64) error {
	w := end - start - 1
	if w > 0 && written != w {
		return logex.NewError("written not expected:", written, w)
	}
	return nil
}

// call after written, offset changed
func (d *DnTask) onWriteFunc(offset int64, written int) error {
	if !d.Meta.IsAccpetRange() {
		d.Meta.MarkFinishStream(int64(written))
		return nil
	}

	err := d.Meta.MarkFinishByN(offset, written, true)
	return logex.Trace(err)
}

func (d *DnTask) httpDn(client *http.Client, req *http.Request, op *writeOp, start, end int64) (int64, error) {
	resp, err := client.Do(req)
	if err != nil {
		return 0, logex.Trace(err)
	}
	rc := NewReader(resp.Body)
	defer rc.Close()

	r := bufio.NewReader(rc)
	w := NewFileWriter(d, start, op, d.writeOp, d.onWriteFunc)
	written, err := io.CopyN(w, r, end-start)
	if err != nil {
		return written, logex.Trace(err)
	}
	if resp.ContentLength != written {
		logex.Error("ContentLength is not expected:",
			resp.ContentLength, written,
			req.Header, resp.Status,
		)
		panic(1)
	}
	io.Copy(ioutil.Discard, rc)
	return written, nil
}

func (d *DnTask) proxyGet(client *http.Client, host string, idx int, op *writeOp, start, end int64) (int64, error) {
	proxy := proxyUrl(host, d.Meta.Source, start, end)
	logex.Info(host, idx, start, end, proxy)
	req, err := http.NewRequest("GET", proxy, nil)
	if err != nil {
		return 0, logex.Trace(err)
	}
	return d.httpDn(client, req, op, start, end)
}

func (d *DnTask) httpGet(client *http.Client, idx int, op *writeOp, start, end int64) (int64, error) {
	req, err := http.NewRequest("GET", d.Meta.Source, nil)
	if err != nil {
		panic(err)
	}

	if idx >= 0 {
		setRange(req.Header, start, end)
	} else {
		start = 0
	}
	return d.httpDn(client, req, op, start, end)
}

func (d *DnTask) download(t *DnType) {
	var (
		idx        int
		start, end int64
		err        error
		retry      int

		maxRetry = 3
		op       = new(writeOp)
	)
	op.Reply = make(chan *writeOpReply)

	if !d.Meta.IsAccpetRange() {
		_, err = d.httpGet(DefaultClient, -1, op, -1, -1)
		if err != nil {
			logex.Error(err)
		}
		return
	}

	for {
		idx, start, end = d.allocDnBlk(idx)
		if idx < 0 {
			break
		}
		if t.Proxy == "" {
			_, err = d.httpGet(DefaultClient, idx, op, start, end)
		} else {
			_, err = d.proxyGet(DefaultClient, t.Proxy, idx, op, start, end)
		}
		if err != nil {
			if retry > maxRetry {
				logex.Error(err)
				return
			}
			retry++
			d.Meta.MarkInit(idx)
			continue
		}

		idx++
	}
}

type DnType struct {
	Proxy string
}

func NewDnType(host string) *DnType {
	return &DnType{host}
}

func (d *DnTask) Schedule(n int) {
	if !d.Meta.IsAccpetRange() {
		n = 1
		logex.Info("range is not acceptable, turn to single thread")
	}
	if n > len(d.Meta.Blocks) {
		logex.Info("remote file size is too small to use", n, "threads, decrease to", len(d.Meta.Blocks))
		n = len(d.Meta.Blocks)
	}

	types := make([]*DnType, len(d.Proxy)+1)
	for i := 0; i < len(types); i++ {
		if i == len(types)-1 {
			types[i] = NewDnType("")
		} else {
			types[i] = NewDnType(d.Proxy[i])
		}
	}

	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		i := i
		go func() {
			d.download(types[i%len(types)])
			wg.Done()
		}()
	}
	wg.Wait()
}

func calUnit(u int64) string {
	units := []string{
		"B", "KB", "MB", "GB", "TB", "PB", "ZB",
	}
	idx := 0
	data := float64(u)
	for data > 10240 {
		idx++
		data /= 1024
	}
	if idx < 2 {
		return fmt.Sprintf("%d%s", int(data), units[idx])
	} else {
		return fmt.Sprintf("%.2f%s", data, units[idx])
	}

}

func (t *DnTask) Close() {
	close(t.stopChan)
	t.wg.Wait()
	t.Meta.Close()
	t.l.Finish()
}

func (t *DnTask) progress() {
	t.wg.Add(1)
	defer t.wg.Done()

	lastWritten := int64(0)
	fileSize := t.Meta.FileSize
	size := calUnit(fileSize)
	ticker := time.NewTicker(time.Second)
	stop := false

	var totalSp, totalN int64
	for !stop {
		select {
		case <-ticker.C:
		case <-t.stopChan:
			stop = true
		}
		written := atomic.LoadInt64(&t.Meta.written)
		if lastWritten > 0 {
			totalSp += written - lastWritten
		}
		atomic.StoreInt64(&t.downloadPerSecond, written-lastWritten)
		totalN += 1
		if t.MaxSpeed > 0 {
			t.rateLimit.Reset()
		}
		realDn := atomic.SwapInt64(&report, 0)

		extend := "\b"
		if t.ShowRealSp {
			extend += fmt.Sprintf(" RL:%v", calUnit(realDn))
		}

		if t.Progress {
			t.l.Print(fmt.Sprintf("[%v/%v(%v%%) DL:%v TIME:%v ETA:%v %v]",
				calUnit(written),
				size,
				calProgress(written, fileSize),
				calUnit(written-lastWritten),
				calTime(time.Now().Sub(t.start)),
				calTime(calRemainTime(fileSize-written, totalSp/totalN)),
				extend,
			))
		}
		lastWritten = written
	}
}

func calRemainTime(remain int64, speed int64) time.Duration {
	if speed == 0 {
		return time.Duration(0)
	}
	return time.Duration(remain/speed) * time.Second
}

func calProgress(a, b int64) int64 {
	return int64(a * 100 / b)
}

func calTime(d time.Duration) string {
	return (time.Duration(d.Seconds()) * time.Second).String()
}

type Liner struct {
	buf *bytes.Buffer

	io.Writer
	last int
}

func NewLiner(w io.Writer) *Liner {
	l := &Liner{Writer: w}
	return l
}

func (l *Liner) Print(objs ...interface{}) {
	if l.buf == nil {
		l.buf = bytes.NewBuffer(nil)
	}
	last := l.last
	if last > 0 {
		l.buf.Write(bytes.Repeat([]byte("\b"), last+2))
		last = 0
	}
	for _, o := range objs {
		n, _ := l.buf.WriteString(fmt.Sprintf("%v", o))
		last += n
	}
	if last < l.last {
		l.buf.Write(bytes.Repeat([]byte(" "), l.last-last))
	} else {
		l.last = last
	}
	l.Writer.Write(l.buf.Bytes())
	l.buf.Reset()
}

func (l *Liner) Finish() {
	fmt.Fprintln(l.Writer)
}
<commit_msg>fix download<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"gopkg.in/logex.v1"
)

var DefaultClient = &http.Client{}

type TaskConfig struct {
	MaxSpeed   int64
	Clean      bool
	Progress   bool
	ShowRealSp bool
	Proxy      []string
}

func (t *TaskConfig) init() {
}

type DnTask struct {
	*TaskConfig
	source *url.URL
	Meta   *Meta

	file     *os.File
	writeOp  chan *writeOp
	stopChan chan struct{}

	wg    sync.WaitGroup
	start time.Time
	sync.Mutex
	rateLimit *RateLimit

	downloadPerSecond int64

	l *Liner
}

func NewDnTaskAuto(url_, pwd string, bit uint, cfg *TaskConfig) (*DnTask, error) {
	_, err := os.Stat(url_)
	if !cfg.Clean && err == nil {
		if meta, _ := NewMetaFormFile(url_); meta != nil {
			logex.Info("downloading form", meta.Source)
			return NewDnTask(meta.Source, pwd, meta.BlkBit, cfg)
		}
	}

	return NewDnTask(url_, pwd, bit, cfg)
}

func NewDnTask(url_, pwd string, bit uint, cfg *TaskConfig) (*DnTask, error) {
	if url_ == "" {
		return nil, logex.NewError("url is empty")
	}
	if cfg == nil {
		cfg = new(TaskConfig)
	}
	cfg.init()

	source, err := url.Parse(url_)
	if err != nil {
		return nil, logex.Trace(err)
	}
	meta, err := NewMeta(pwd, url_, bit, cfg.Clean)
	if err != nil {
		return nil, logex.Trace(err)
	}

	dn := &DnTask{
		TaskConfig: cfg,
		rateLimit:  NewRateLimit(cfg.MaxSpeed),
		source:     source,
		Meta:       meta,
		writeOp:    make(chan *writeOp, 1<<3),
		stopChan:   make(chan struct{}),
		start:      time.Now(),
		l:          NewLiner(os.Stderr),
	}
	if cfg.Clean {
		os.Remove(dn.Meta.targetPath())
	}

	if err = dn.Meta.retrieveFromDisk(cfg.Proxy); err != nil {
		dn.Meta.Remove()
		return nil, logex.Trace(err)
	}

	if err = dn.Meta.Sync(); err != nil {
		return nil, logex.Trace(err)
	}

	if err = dn.openFile(); err != nil {
		return nil, logex.Trace(err)
	}

	go dn.ioloop()
	go dn.progress()
	return dn, nil
}

func (d *DnTask) openFile() error {
	f, err := os.OpenFile(d.Meta.targetPath(), os.O_RDWR|os.O_CREATE, 0666)
	if err != nil {
		return logex.Trace(err)
	}
	d.file = f
	return nil
}

type writeOpReply struct {
	N   int
	Err error
}

type writeOp struct {
	Offset int64
	Buf    []byte
	Reply  chan *writeOpReply
}

func (d *DnTask) ioloop() {
	var w *writeOp
	for {
		select {
		case w = <-d.writeOp:
		case <-d.stopChan:
			return
		}
		n, err := d.file.WriteAt(w.Buf, w.Offset)
		if err != nil {
			d.file.Close()
			d.file = nil
			if err := d.openFile(); err != nil {
				panic(err)
			}
		}
		if d.MaxSpeed > 0 {
			d.rateLimit.Process(n)
		}
		w.Reply <- &writeOpReply{n, logex.Trace(err)}
	}
}

// [start, end)
func (d *DnTask) allocDnBlk(off int) (idx int, start, end int64) {
	d.Lock()
	defer d.Unlock()

	for i := off; i < len(d.Meta.Blocks); i++ {
		blk := d.Meta.Blocks[i]
		if blk == nil {
			blk = NewBlock()
			d.Meta.Blocks[i] = blk
		}
		if blk.State != STATE_INIT {
			continue
		}
		blk.State = STATE_PROCESS
		offset := int64(i << d.Meta.BlkBit)
		start = offset + int64(blk.Written)
		end = int64((i + 1) << d.Meta.BlkBit)
		if end > d.Meta.FileSize {
			end = d.Meta.FileSize
		}
		return i, start, end
	}
	return -1, -1, -1
}

func setRange(h http.Header, start, end int64) {
	h.Set(H_RANGE, fmt.Sprintf("bytes=%d-%d", start, end-1))
}

func (d *DnTask) checkWritten(written, start, end int64) error {
	w := end - start - 1
	if w > 0 && written != w {
		return logex.NewError("written not expected:", written, w)
	}
	return nil
}

// call after written, offset changed
func (d *DnTask) onWriteFunc(offset int64, written int) error {
	if !d.Meta.IsAccpetRange() {
		d.Meta.MarkFinishStream(int64(written))
		return nil
	}

	err := d.Meta.MarkFinishByN(offset, written, true)
	return logex.Trace(err)
}

func (d *DnTask) httpDn(client *http.Client, req *http.Request, op *writeOp, start, end int64) (int64, error) {
	resp, err := client.Do(req)
	if err != nil {
		return 0, logex.Trace(err)
	}
	rc := NewReader(resp.Body)
	defer rc.Close()

	r := bufio.NewReader(rc)
	w := NewFileWriter(d, start, op, d.writeOp, d.onWriteFunc)
	written, err := io.CopyN(w, r, end-start)
	if err != nil {
		return written, logex.Trace(err)
	}
	if resp.ContentLength != written {
		logex.Error("ContentLength is not expected:",
			resp.ContentLength, written,
			req.Header, resp.Status, req.URL,
		)
		panic(1)
	}
	io.Copy(ioutil.Discard, rc)
	return written, nil
}

func (d *DnTask) proxyGet(client *http.Client, host string, idx int, op *writeOp, start, end int64) (int64, error) {
	proxy := proxyUrl(host, d.Meta.Source, start, end)
	req, err := http.NewRequest("GET", proxy, nil)
	if err != nil {
		return 0, logex.Trace(err)
	}
	return d.httpDn(client, req, op, start, end)
}

func (d *DnTask) httpGet(client *http.Client, idx int, op *writeOp, start, end int64) (int64, error) {
	req, err := http.NewRequest("GET", d.Meta.Source, nil)
	if err != nil {
		panic(err)
	}

	if idx >= 0 {
		setRange(req.Header, start, end)
	} else {
		start = 0
	}
	return d.httpDn(client, req, op, start, end)
}

func (d *DnTask) download(t *DnType) {
	var (
		idx        int
		start, end int64
		err        error
		retry      int

		maxRetry = 3
		op       = new(writeOp)
	)
	op.Reply = make(chan *writeOpReply)

	if !d.Meta.IsAccpetRange() {
		_, err = d.httpGet(DefaultClient, -1, op, -1, -1)
		if err != nil {
			logex.Error(err)
		}
		return
	}

	for {
		idx, start, end = d.allocDnBlk(idx)
		if idx < 0 {
			break
		}
		if t.Proxy == "" {
			_, err = d.httpGet(DefaultClient, idx, op, start, end)
		} else {
			_, err = d.proxyGet(DefaultClient, t.Proxy, idx, op, start, end)
		}
		if err != nil {
			if retry > maxRetry {
				logex.Error(err)
				return
			}
			retry++
			d.Meta.MarkInit(idx)
			continue
		}

		idx++
	}
}

type DnType struct {
	Proxy string
}

func NewDnType(host string) *DnType {
	return &DnType{host}
}

func (d *DnTask) Schedule(n int) {
	if !d.Meta.IsAccpetRange() {
		n = 1
		logex.Info("range is not acceptable, turn to single thread")
	}
	if n > len(d.Meta.Blocks) {
		logex.Info("remote file size is too small to use", n, "threads, decrease to", len(d.Meta.Blocks))
		n = len(d.Meta.Blocks)
	}

	types := make([]*DnType, len(d.Proxy)+1)
	for i := 0; i < len(types); i++ {
		if i == len(types)-1 {
			types[i] = NewDnType("")
		} else {
			types[i] = NewDnType(d.Proxy[i])
		}
	}

	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		i := i
		go func() {
			d.download(types[i%len(types)])
			wg.Done()
		}()
	}
	wg.Wait()
}

func calUnit(u int64) string {
	units := []string{
		"B", "KB", "MB", "GB", "TB", "PB", "ZB",
	}
	idx := 0
	data := float64(u)
	for data > 10240 {
		idx++
		data /= 1024
	}
	if idx < 2 {
		return fmt.Sprintf("%d%s", int(data), units[idx])
	} else {
		return fmt.Sprintf("%.2f%s", data, units[idx])
	}

}

func (t *DnTask) Close() {
	close(t.stopChan)
	t.wg.Wait()
	t.Meta.Close()
	t.l.Finish()
}

func (t *DnTask) progress() {
	t.wg.Add(1)
	defer t.wg.Done()

	lastWritten := int64(0)
	fileSize := t.Meta.FileSize
	size := calUnit(fileSize)
	ticker := time.NewTicker(time.Second)
	stop := false

	var totalSp, totalN int64
	for !stop {
		select {
		case <-ticker.C:
		case <-t.stopChan:
			stop = true
		}
		written := atomic.LoadInt64(&t.Meta.written)
		if lastWritten > 0 {
			totalSp += written - lastWritten
		}
		atomic.StoreInt64(&t.downloadPerSecond, written-lastWritten)
		totalN += 1
		if t.MaxSpeed > 0 {
			t.rateLimit.Reset()
		}
		realDn := atomic.SwapInt64(&report, 0)

		extend := "\b"
		if t.ShowRealSp {
			extend += fmt.Sprintf(" RL:%v", calUnit(realDn))
		}

		if t.Progress {
			t.l.Print(fmt.Sprintf("[%v/%v(%v%%) DL:%v TIME:%v ETA:%v %v]",
				calUnit(written),
				size,
				calProgress(written, fileSize),
				calUnit(written-lastWritten),
				calTime(time.Now().Sub(t.start)),
				calTime(calRemainTime(fileSize-written, totalSp/totalN)),
				extend,
			))
		}
		lastWritten = written
	}
}

func calRemainTime(remain int64, speed int64) time.Duration {
	if speed == 0 {
		return time.Duration(0)
	}
	return time.Duration(remain/speed) * time.Second
}

func calProgress(a, b int64) int64 {
	return int64(a * 100 / b)
}

func calTime(d time.Duration) string {
	return (time.Duration(d.Seconds()) * time.Second).String()
}

type Liner struct {
	buf *bytes.Buffer

	io.Writer
	last int
}

func NewLiner(w io.Writer) *Liner {
	l := &Liner{Writer: w}
	return l
}

func (l *Liner) Print(objs ...interface{}) {
	if l.buf == nil {
		l.buf = bytes.NewBuffer(nil)
	}
	last := l.last
	if last > 0 {
		l.buf.Write(bytes.Repeat([]byte("\b"), last+2))
		last = 0
	}
	for _, o := range objs {
		n, _ := l.buf.WriteString(fmt.Sprintf("%v", o))
		last += n
	}
	if last < l.last {
		l.buf.Write(bytes.Repeat([]byte(" "), l.last-last))
	} else {
		l.last = last
	}
	l.Writer.Write(l.buf.Bytes())
	l.buf.Reset()
}

func (l *Liner) Finish() {
	fmt.Fprintln(l.Writer)
}
<|endoftext|>


<commit_before>package tree

import (
	"fmt"

	ct "github.com/phil-mansfield/consistent_trees"
)

// HaloHistories takes a slice of Rockstar halo tree file names, a slice of
// the root halo IDs, and the "snapshot offset." The snapshot offset is the
// difference between the number of snpashots which contain a nonzero number
// of halos and the total number of snapshots. This can be calculated by
// env.Halos.SnapOffset(). HaloHistories will return slices of IDs and
// snapshots which correspond to the history of each of the given root IDs.
func HaloHistories(
	files []string, roots []int, snapOffset int,
) (ids [][]int, snaps [][]int, err error) {
	if len(roots) == 0 {
		return [][]int{}, [][]int{}, nil
	}

	ids, snaps = make([][]int, len(roots)), make([][]int, len(roots))

	foundCount := 0
	for _, file := range files {
		ct.ReadTree(file)
		var ok bool
		for i, id := range roots {
			if ids[i] != nil {
				continue
			}
			if ids[i], snaps[i], ok = findHistory(id); ok {
				foundCount++
			}
		}
		ct.DeleteTree()
		if foundCount == len(roots) {
			break
		}
	}

	for i, idSnaps := range snaps {
		if idSnaps == nil {
			return nil, nil, fmt.Errorf(
				"Halo %d not found in given files.", roots[i],
			)
		}
	}

	for i := range snaps {
		for j := range snaps[i] {
			snaps[i][j] += snapOffset
		}
	}
	return ids, snaps, nil
}

// HaloSnaps takes a slice of Rockstar halo trees and a slice of IDs. It
// will return the snapshots which each of those IDs are from.
func HaloSnaps(files []string, ids []int) (snaps []int, err error) {
	if len(ids) == 0 {
		return []int{}, nil
	}

	snaps = make([]int, len(ids))
	for i := range snaps {
		snaps[i] = -1
	}

	foundCount := 0
	for _, file := range files {
		ct.ReadTree(file)
		for i, id := range ids {
			if snaps[i] != -1 {
				continue
			}
			if _, snap, ok := findHalo(id); ok {
				snaps[i] = snap
				foundCount++
			}
		}
		ct.DeleteTree()
		if foundCount == len(ids) {
			break
		}
	}

	for i, snap := range snaps {
		if snap == -1 {
			return nil, fmt.Errorf(
				"Halo %d not found in given files.", ids[i],
			)
		}
	}
	return snaps, nil
}

func findHalo(id int) (ct.Halo, int, bool) {
	tree := ct.GetHaloTree()
	for i := 0; i < tree.NumLists(); i++ {
		list := tree.HaloLists(i)
		h, ok := ct.LookupHaloInList(list, id)
		if ok {
			return h, tree.NumLists() - i, true
		}
	}
	return ct.Halo{}, 0, false
}

func findHistory(id int) (ids, snaps []int, ok bool) {
	h, snap, ok := findHalo(id)
	if !ok {
		return nil, nil, false
	}
	desc, descSnaps := descTree(h)
	prog, progSnaps := progTree(h)

	ids = combine(reverse(prog), []int{id}, desc)
	snaps = combine(reverse(progSnaps), []int{snap}, descSnaps)
	return ids, snaps, true
}

func descTree(h ct.Halo) (ids, snaps []int) {
	ids, snaps = []int{}, []int{}
	var ok bool
	numLists := ct.GetHaloTree().NumLists()
	for {
		h, ok = h.Desc()
		if !ok {
			break
		}
		ids = append(ids, h.ID())
		snaps = append(snaps, numLists-ct.LookupIndex(h.Scale()))
	}
	return ids, snaps
}

func progTree(h ct.Halo) (ids, snaps []int) {
	ids, snaps = []int{}, []int{}
	var ok bool
	numLists := ct.GetHaloTree().NumLists()
	for {
		h, ok = h.Prog()
		if !ok {
			break
		}
		ids = append(ids, h.ID())
		snaps = append(snaps, numLists-ct.LookupIndex(h.Scale()))
	}
	return ids, snaps
}

func reverse(xs []int) []int {
	out := make([]int, len(xs))
	for i := range xs {
		out[i] = xs[len(xs)-1-i]
	}
	return out
}

func combine(slices ...[]int) []int {
	out := []int{}
	for _, slice := range slices {
		out = append(out, slice...)
	}
	return out
}
<commit_msg>Removed HaloSnaps from los/tree.<commit_after>package tree

import (
	"fmt"

	ct "github.com/phil-mansfield/consistent_trees"
)

// HaloHistories takes a slice of Rockstar halo tree file names, a slice of
// the root halo IDs, and the "snapshot offset." The snapshot offset is the
// difference between the number of snpashots which contain a nonzero number
// of halos and the total number of snapshots. This can be calculated by
// env.Halos.SnapOffset(). HaloHistories will return slices of IDs and
// snapshots which correspond to the history of each of the given root IDs.
func HaloHistories(
	files []string, roots []int, snapOffset int,
) (ids [][]int, snaps [][]int, err error) {
	if len(roots) == 0 {
		return [][]int{}, [][]int{}, nil
	}

	ids, snaps = make([][]int, len(roots)), make([][]int, len(roots))

	foundCount := 0
	for _, file := range files {
		ct.ReadTree(file)
		var ok bool
		for i, id := range roots {
			if ids[i] != nil {
				continue
			}
			if ids[i], snaps[i], ok = findHistory(id); ok {
				foundCount++
			}
		}
		ct.DeleteTree()
		if foundCount == len(roots) {
			break
		}
	}

	for i, idSnaps := range snaps {
		if idSnaps == nil {
			return nil, nil, fmt.Errorf(
				"Halo %d not found in given files.", roots[i],
			)
		}
	}

	for i := range snaps {
		for j := range snaps[i] {
			snaps[i][j] += snapOffset
		}
	}
	return ids, snaps, nil
}

func findHalo(id int) (ct.Halo, int, bool) {
	tree := ct.GetHaloTree()
	for i := 0; i < tree.NumLists(); i++ {
		list := tree.HaloLists(i)
		h, ok := ct.LookupHaloInList(list, id)
		if ok {
			return h, tree.NumLists() - i, true
		}
	}
	return ct.Halo{}, 0, false
}

func findHistory(id int) (ids, snaps []int, ok bool) {
	h, snap, ok := findHalo(id)
	if !ok {
		return nil, nil, false
	}
	desc, descSnaps := descTree(h)
	prog, progSnaps := progTree(h)

	ids = combine(reverse(prog), []int{id}, desc)
	snaps = combine(reverse(progSnaps), []int{snap}, descSnaps)
	return ids, snaps, true
}

func descTree(h ct.Halo) (ids, snaps []int) {
	ids, snaps = []int{}, []int{}
	var ok bool
	numLists := ct.GetHaloTree().NumLists()
	for {
		h, ok = h.Desc()
		if !ok {
			break
		}
		ids = append(ids, h.ID())
		snaps = append(snaps, numLists-ct.LookupIndex(h.Scale()))
	}
	return ids, snaps
}

func progTree(h ct.Halo) (ids, snaps []int) {
	ids, snaps = []int{}, []int{}
	var ok bool
	numLists := ct.GetHaloTree().NumLists()
	for {
		h, ok = h.Prog()
		if !ok {
			break
		}
		ids = append(ids, h.ID())
		snaps = append(snaps, numLists-ct.LookupIndex(h.Scale()))
	}
	return ids, snaps
}

func reverse(xs []int) []int {
	out := make([]int, len(xs))
	for i := range xs {
		out[i] = xs[len(xs)-1-i]
	}
	return out
}

func combine(slices ...[]int) []int {
	out := []int{}
	for _, slice := range slices {
		out = append(out, slice...)
	}
	return out
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"runtime"
	"sync"
	"time"
)

type M map[string]interface{}

type Request struct {
	Method string
	Token  string
}

type Response struct {
	Token string      `json:"token"`
	Error string      `json:"error"`
	Data  interface{} `json:"data"`
}

type Broker struct {
	served uint
	start  time.Time
	rLck   sync.Mutex
	wLck   sync.Mutex
	r      io.Reader
	w      io.Writer
	in     *bufio.Reader
	out    *json.Encoder
}

func NewBroker(r io.Reader, w io.Writer) *Broker {
	return &Broker{
		r:   r,
		w:   w,
		in:  bufio.NewReader(r),
		out: json.NewEncoder(w),
	}
}

func (b *Broker) Send(resp Response) error {
	err := b.SendNoLog(resp)
	if err != nil {
		logger.Println("Cannot send result", err)
	}
	return err
}

func (b *Broker) SendNoLog(resp Response) error {
	b.wLck.Lock()
	defer b.wLck.Unlock()

	if resp.Data == nil {
		resp.Data = M{}
	}

	s, err := json.Marshal(resp)
	if err != nil {
		return err
	}

	// the only expected write failure are due to broken pipes
	// which usually means the client has gone away so just ignore the error
	b.w.Write(s)
	return nil
}

func (b *Broker) call(req *Request, cl Caller) {
	b.served++

	defer func() {
		err := recover()
		if err != nil {
			logger.Printf("%v#%v PANIC: %v\n", req.Method, req.Token, err)
			b.Send(Response{
				Token: req.Token,
				Error: "broker: " + req.Method + "#" + req.Token + " PANIC",
			})
		}
	}()

	res, err := cl.Call()
	b.Send(Response{
		Token: req.Token,
		Error: err,
		Data:  res,
	})
}

func (b *Broker) accept() (stopLooping bool) {
	line, err := b.in.ReadBytes('\n')

	if err == io.EOF {
		stopLooping = true
	} else if err != nil {
		logger.Println("Cannot read input", err)
		b.Send(Response{
			Error: err.Error(),
		})
		return
	}

	req := &Request{}
	dec := json.NewDecoder(bytes.NewReader(line))
	// if this fails, we are unable to return a useful error(no token to send it to)
	// so we'll simply/implicitly drop the request since it has no method
	// we can safely assume that all such cases will be empty lines and not an actual request
	dec.Decode(&req)

	if req.Method == "" {
		return
	}

	if req.Method == "bye-ni" {
		return true
	}

	m := registry.Lookup(req.Method)
	if m == nil {
		e := "Invalid method " + req.Method
		logger.Println(e)
		b.Send(Response{
			Token: req.Token,
			Error: e,
		})
		return
	}

	cl := m(b)
	err = dec.Decode(cl)
	if err != nil {
		logger.Println("Cannot decode arg", err)
		b.Send(Response{
			Token: req.Token,
			Error: err.Error(),
		})
		return
	}

	go b.call(req, cl)

	return
}

func (b *Broker) Loop(decorate bool) {
	b.start = time.Now()

	if decorate {
		go b.SendNoLog(Response{
			Token: "margo.hello",
			Data: M{
				"time": b.start.String(),
			},
		})
	}

	for {
		stopLooping := b.accept()
		if stopLooping {
			break
		}
		runtime.Gosched()
	}

	if decorate {
		b.SendNoLog(Response{
			Token: "margo.bye-ni",
			Data: M{
				"served": b.served,
				"uptime": time.Now().Sub(b.start).String(),
			},
		})
	}
}
<commit_msg>* write trailing LF - fixes python client not being able to decode the response :|<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"runtime"
	"sync"
	"time"
)

type M map[string]interface{}

type Request struct {
	Method string
	Token  string
}

type Response struct {
	Token string      `json:"token"`
	Error string      `json:"error"`
	Data  interface{} `json:"data"`
}

type Broker struct {
	served uint
	start  time.Time
	rLck   sync.Mutex
	wLck   sync.Mutex
	r      io.Reader
	w      io.Writer
	in     *bufio.Reader
	out    *json.Encoder
}

func NewBroker(r io.Reader, w io.Writer) *Broker {
	return &Broker{
		r:   r,
		w:   w,
		in:  bufio.NewReader(r),
		out: json.NewEncoder(w),
	}
}

func (b *Broker) Send(resp Response) error {
	err := b.SendNoLog(resp)
	if err != nil {
		logger.Println("Cannot send result", err)
	}
	return err
}

func (b *Broker) SendNoLog(resp Response) error {
	b.wLck.Lock()
	defer b.wLck.Unlock()

	if resp.Data == nil {
		resp.Data = M{}
	}

	s, err := json.Marshal(resp)
	if err != nil {
		return err
	}

	// the only expected write failure are due to broken pipes
	// which usually means the client has gone away so just ignore the error
	b.w.Write(s)
	b.w.Write([]byte{'\n'})
	return nil
}

func (b *Broker) call(req *Request, cl Caller) {
	b.served++

	defer func() {
		err := recover()
		if err != nil {
			logger.Printf("%v#%v PANIC: %v\n", req.Method, req.Token, err)
			b.Send(Response{
				Token: req.Token,
				Error: "broker: " + req.Method + "#" + req.Token + " PANIC",
			})
		}
	}()

	res, err := cl.Call()
	b.Send(Response{
		Token: req.Token,
		Error: err,
		Data:  res,
	})
}

func (b *Broker) accept() (stopLooping bool) {
	line, err := b.in.ReadBytes('\n')

	if err == io.EOF {
		stopLooping = true
	} else if err != nil {
		logger.Println("Cannot read input", err)
		b.Send(Response{
			Error: err.Error(),
		})
		return
	}

	req := &Request{}
	dec := json.NewDecoder(bytes.NewReader(line))
	// if this fails, we are unable to return a useful error(no token to send it to)
	// so we'll simply/implicitly drop the request since it has no method
	// we can safely assume that all such cases will be empty lines and not an actual request
	dec.Decode(&req)

	if req.Method == "" {
		return
	}

	if req.Method == "bye-ni" {
		return true
	}

	m := registry.Lookup(req.Method)
	if m == nil {
		e := "Invalid method " + req.Method
		logger.Println(e)
		b.Send(Response{
			Token: req.Token,
			Error: e,
		})
		return
	}

	cl := m(b)
	err = dec.Decode(cl)
	if err != nil {
		logger.Println("Cannot decode arg", err)
		b.Send(Response{
			Token: req.Token,
			Error: err.Error(),
		})
		return
	}

	go b.call(req, cl)

	return
}

func (b *Broker) Loop(decorate bool) {
	b.start = time.Now()

	if decorate {
		go b.SendNoLog(Response{
			Token: "margo.hello",
			Data: M{
				"time": b.start.String(),
			},
		})
	}

	for {
		stopLooping := b.accept()
		if stopLooping {
			break
		}
		runtime.Gosched()
	}

	if decorate {
		b.SendNoLog(Response{
			Token: "margo.bye-ni",
			Data: M{
				"served": b.served,
				"uptime": time.Now().Sub(b.start).String(),
			},
		})
	}
}
<|endoftext|>


<commit_before>package minion

import (
	"time"

	"code.google.com/p/go-uuid/uuid"
)

type Minion interface {
	// Get minion identifier
	GetUUID() uuid.UUID

	// Set name of minion
	SetName(name string) error

	// Get name of minion
	GetName() (string, error)

	// Set the time the minion was last seen in seconds since the Epoch
	SetLastseen(s int64) error

	// Get a classifier for a minion
	GetClassifier(key string) (MinionClassifier, error)

	// Classify minion a with given a key and value
	SetClassifier(c MinionClassifier) error

	// Runs periodic functions, e.g. refreshes classifies and lastseen
	Refresh(t *time.Ticker) error

	// Listens for new tasks and processes them
	TaskListener(c chan<- MinionTask) error

	// Start serving
	Serve() error
}

type MinionTask interface {
	// Gets the command to be executed
	GetCommand() (string, error)

	// Gets the time the task was sent for processing
	GetTimestamp() (int64, error)
}
<commit_msg>MinionTask interfaces requires implementing the Process() and Submit() methods<commit_after>package minion

import (
	"time"

	"code.google.com/p/go-uuid/uuid"
)

type Minion interface {
	// Get minion identifier
	GetUUID() uuid.UUID

	// Set name of minion
	SetName(name string) error

	// Get name of minion
	GetName() (string, error)

	// Set the time the minion was last seen in seconds since the Epoch
	SetLastseen(s int64) error

	// Get a classifier for a minion
	GetClassifier(key string) (MinionClassifier, error)

	// Classify minion a with given a key and value
	SetClassifier(c MinionClassifier) error

	// Runs periodic functions, e.g. refreshes classifies and lastseen
	Refresh(t *time.Ticker) error

	// Listens for new tasks and processes them
	TaskListener(c chan<- MinionTask) error

	// Start serving
	Serve() error
}

type MinionTask interface {
	// Gets the command to be executed
	GetCommand() (string, error)

	// Gets the time the task was sent for processing
	GetTimestamp() (int64, error)

	// Processes the task
	Process() error

	// Submits a task to a minion
	Submit(m Minion) error
}
<|endoftext|>